sc

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

Author: praateekmahajan

nemo_curator/stages/deduplication/semantic/identify_duplicates.py

@@ -89,7 +89,8 @@ def process_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]:
 
         all_files = [file for task in tasks for file in task.data]
         # Read using filters
-
+        # We read file by file since list[files] when files are remote urls can fail
+        # See https://github.com/pandas-dev/pandas/issues/62922
         df: pd.DataFrame = pd.concat(
             [
                 pd.read_parquet(

nemo_curator/stages/text/deduplication/removal.py

@@ -76,9 +76,7 @@ def process(self, task: DocumentBatch) -> DocumentBatch:
         # Filter the parquet files for IDs to remove within this range
         read_dupes_t0 = time.perf_counter()
 
-        # we use pq.read_table instead of pd.read_parquet since ids_to_remove_path is a directory
-        # and it might error out when the directory is a cloud path
-        removal_table = pd.read_parquet(
+        removal_df = pd.read_parquet(
             self.ids_to_remove_path,
             filters=[(self.duplicate_id_field, ">=", min_id), (self.duplicate_id_field, "<=", max_id)],
             columns=[self.duplicate_id_field],
@@ -89,7 +87,7 @@ def process(self, task: DocumentBatch) -> DocumentBatch:
 
         # Filter out documents with IDs in the removal set using pandas
         time_to_remove_t0 = time.perf_counter()
-        removal_ids = set(removal_table[self.duplicate_id_field].tolist())
+        removal_ids = set(removal_df[self.duplicate_id_field].tolist())
         df = df[~df[self.id_field].isin(removal_ids)]
         removal_ids_time = time.perf_counter() - time_to_remove_t0
         self._log_metrics(