Bug fix in removal write kwargs + add input_task_limit in removal (#995)

Signed-off-by: Praateek Mahajan <[email protected]>

Author: praateekmahajan

nemo_curator/stages/text/deduplication/removal_workflow.py

@@ -42,6 +42,7 @@ class TextDuplicatesRemovalWorkflow:
     input_files_per_partition: int | None = None
     input_blocksize: str | None = None
     input_file_extensions: list[str] | None = None
+    input_task_limit: int | None = None
     input_kwargs: dict[str, Any] | None = None
 
     # ids_to_remove args
@@ -83,6 +84,7 @@ def _generate_stages(self, initial_tasks: list[FileGroupTask] | None = None) ->
                     blocksize=self.input_blocksize,
                     file_extensions=self.input_file_extensions,
                     storage_options=(self.input_kwargs or {}).get("storage_options"),
+                    limit=self.input_task_limit,
                 )
             )
         else:
@@ -135,7 +137,7 @@ def _generate_stages(self, initial_tasks: list[FileGroupTask] | None = None) ->
             write_stage(
                 path=self.output_path,
                 **({"file_extension": self.output_file_extension} if self.output_file_extension else {}),
-                write_kwargs=self.output_kwargs,
+                write_kwargs=self.output_kwargs or {},
                 fields=self.output_fields,
                 **({"mode": self.output_mode} if self.output_mode else {}),
             )
@@ -151,6 +153,11 @@ def run(
             description="Text duplicates removal workflow",
             stages=self._generate_stages(initial_tasks),
         )
+        if self.input_task_limit is not None and len(initial_tasks) > self.input_task_limit:
+            logger.warning(
+                f"Initial tasks provided ({len(initial_tasks)}) is greater than input_task_limit ({self.input_task_limit}), truncating to {self.input_task_limit}"
+            )
+            initial_tasks = initial_tasks[: self.input_task_limit]
 
         if executor is None:
             from nemo_curator.backends.xenna import XennaExecutor

tests/stages/text/deduplication/test_removal_workflow.py

@@ -230,6 +230,7 @@ def test_initial_tasks_partitioning(self, test_config: "TestTextDuplicateRemoval
             output_filetype="parquet",
             input_id_field=CURATOR_DEDUP_ID_STR,
             ids_to_remove_duplicate_id_field="id",
+            input_task_limit=10,  # truncate to 10 tasks only
             input_kwargs={},
             ids_to_remove_read_kwargs={},
             output_kwargs={},
@@ -239,14 +240,14 @@ def test_initial_tasks_partitioning(self, test_config: "TestTextDuplicateRemoval
         output_tasks = workflow.run(executor, initial_tasks=initial_tasks)
 
         # Verify we get 20 output tasks (one per input task)
-        assert len(output_tasks) == 20, (
-            f"Expected 20 output tasks, got {len(output_tasks)} for {test_config.executor_cls.__name__}"
+        assert len(output_tasks) == 10, (
+            f"Expected 10 output tasks, got {len(output_tasks)} for {test_config.executor_cls.__name__}"
         )
 
         # Verify correctness remains the same as other tests
         combined_output_df = pd.concat([pd.read_parquet(task.data) for task in output_tasks], ignore_index=True)
-        assert len(combined_output_df) == 800, (
-            f"Expected 800 records, got {len(combined_output_df)} for {test_config.executor_cls.__name__}"
+        assert len(combined_output_df) == 400, (
+            f"Expected 400 records, got {len(combined_output_df)} for {test_config.executor_cls.__name__}"
         )
 
         # Verify no IDs divisible by 5 remain