Merge remote-tracking branch 'upstream/main' into praateek/ray_data_max_workers_per_node

Signed-off-by: Praateek <praateekm@gmail.com>

Author: praateekmahajan

.github/workflows/config/.secrets.baseline

@@ -518,7 +518,7 @@
         "filename": "tutorials/text/deduplication/semantic/semantic_e2e.ipynb",
         "hashed_secret": "00ad2e1afc84d1b7c3c139b68f00edd67998232c",
         "is_verified": false,
-        "line_number": 652
+        "line_number": 653
       }
     ],
     "tutorials/text/deduplication/semantic/semantic_step_by_step.ipynb": [
@@ -527,7 +527,7 @@
         "filename": "tutorials/text/deduplication/semantic/semantic_step_by_step.ipynb",
         "hashed_secret": "00ad2e1afc84d1b7c3c139b68f00edd67998232c",
         "is_verified": false,
-        "line_number": 670
+        "line_number": 671
       }
     ]
   },

tutorials/text/deduplication/fuzzy/fuzzy_e2e.ipynb

@@ -1123,7 +1123,7 @@
     "    def process(self, batch: DocumentBatch) -> DocumentBatch:\n",
     "        df = batch.to_pandas().merge(cc_df, how=\"inner\", on=[CURATOR_DEDUP_ID_STR])\n",
     "        return DocumentBatch(\n",
-    "            task_id=batch.task_id, dataset_name=batch.dataset_name, data=df, _stage_perf=batch._stage_perf\n",
+    "            dataset_name=batch.dataset_name, data=df, _stage_perf=batch._stage_perf\n",
     "        )\n",
     "\n",
     "\n",
@@ -1509,7 +1509,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.11"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,

tutorials/text/deduplication/semantic/semantic_e2e.ipynb

@@ -110,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "11448dfb",
    "metadata": {},
    "outputs": [],
@@ -128,6 +128,7 @@
     "    model_identifier=\"google/embeddinggemma-300m\",\n",
     "    embedding_max_chars=None,\n",
     "    embedding_pretokenize=False,\n",
+    "    embedding_vllm_init_kwargs={\"gpu_memory_utilization\": 0.9},  # lower if OOM occurs\n",
     "    # Semantic deduplication parameters\n",
     "    n_clusters=100,  # this number can be much higher when the data is large\n",
     "    # For large scale data removal we should use CURATOR_DEDUP_ID_STR\n",

tutorials/text/deduplication/semantic/semantic_step_by_step.ipynb

@@ -218,6 +218,7 @@
     "        ParquetReader(file_paths=input_path, files_per_partition=1, fields=[\"text\"], _generate_ids=True),\n",
     "        VLLMEmbeddingModelStage(\n",
     "            model_identifier=\"google/embeddinggemma-300m\",\n",
+    "            vllm_init_kwargs={\"gpu_memory_utilization\": 0.9},  # lower if OOM occurs\n",
     "            text_field=\"text\",\n",
     "            max_chars=None,\n",
     "            pretokenize=False,\n",