add skip in decont index builder

swiss-ai · May 4, 2024 · d56d3c5 · d56d3c5
1 parent 15c8425
commit d56d3c5
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/datatrove/pipeline/decont/n_grams.py b/src/datatrove/pipeline/decont/n_grams.py
@@ -147,8 +147,14 @@ def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1
 
         for task_name, task in task_dict.items():
             for eval_doc in task.eval_docs():
-                for gold in eval_doc.get_golds():
-                    hashes[task_name].update(self.compute_hashes(gold, eval_doc.query))
+                try:
+                    golds = eval_doc.get_golds()
+                    query = eval_doc.query
+                except Exception as e:
+                    logger.warning(f"Error while fetching doc data: {e}")
+                    continue
+                for gold in golds:
+                    hashes[task_name].update(self.compute_hashes(gold, query))
 
         for task_name, task_hashes in hashes.items():
             hashes_array = np.array(list(task_hashes), dtype="<u8")
@@ -216,5 +222,7 @@ def filter(self, doc: Document) -> bool | Tuple[bool, str]:
                 doc.metadata["contaminated_ngram"] = n_gram
                 doc.metadata["contaminated_task"] = task
                 self.stat_update(f"contaminated_{task}")
+                if ":" in task:
+                    self.stat_update(f"contaminated_tg_{task[:task.index(':')]}")
                 return False, "contaminated"
         return True