Using DocumentDataset after removal in tutorials / examples (#724)

Author: praateekmahajan

examples/exact_deduplication.py

@@ -17,7 +17,7 @@
 
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.modules import ExactDuplicates
-from nemo_curator.utils.distributed_utils import get_client, write_to_disk
+from nemo_curator.utils.distributed_utils import get_client
 from nemo_curator.utils.script_utils import ArgumentHelper
 
 
@@ -60,7 +60,7 @@ def main(args: argparse.Namespace) -> None:
         duplicates = DocumentDataset.read_parquet(duplicates, backend=backend)
 
     result = exact_dup.remove(input_dataset, duplicates)
-    write_to_disk(result, output_dir, output_type="parquet")
+    result.to_parquet(output_dir)
     print(time.time() - t0)
 
 

examples/fuzzy_deduplication.py

@@ -19,7 +19,7 @@
 
 from nemo_curator import FuzzyDuplicates, FuzzyDuplicatesConfig
 from nemo_curator.datasets import DocumentDataset
-from nemo_curator.utils.distributed_utils import get_client, write_to_disk
+from nemo_curator.utils.distributed_utils import get_client
 from nemo_curator.utils.script_utils import ArgumentHelper
 
 
@@ -92,7 +92,7 @@ def main(args: argparse.Namespace) -> None:
             return
 
         result = fuzzy_dup.remove(input_dataset, duplicates)
-        write_to_disk(result, output_dir, output_type=filetype)
+        result.to_parquet(output_dir)
         print(f"Time taken:{time.time() - t0}s")
 
 

tutorials/dapt-curation/code/utils.py

@@ -274,8 +274,7 @@ def exact_dedupe(dataset: DocumentDataset) -> DocumentDataset:
     deduplicator = ExactDuplicates(id_field="id", text_field="text", hash_method="md5")
     # Find the duplicates
     duplicates = deduplicator(dataset)
-    deduped = deduplicator.remove(dataset, duplicates)
-    return DocumentDataset(deduped)
+    return deduplicator.remove(dataset, duplicates)
 
 
 def fuzzy_dedupe(dataset: DocumentDataset, cache_dir: str) -> DocumentDataset:

tutorials/multimodal_dapt_curation/curator/utils.py

@@ -162,8 +162,7 @@ def exact_dedupe(dataset: DocumentDataset) -> DocumentDataset:
     deduplicator = ExactDuplicates(id_field="id", text_field="text", hash_method="md5")
     # Find the duplicates
     duplicates = deduplicator(dataset)
-    deduped = deduplicator.remove(dataset, duplicates)
-    return DocumentDataset(deduped)
+    return deduplicator.remove(dataset, duplicates)
 
 
 def fuzzy_dedupe(dataset: DocumentDataset, cache: str) -> DocumentDataset:

tutorials/tinystories/main.py

@@ -152,8 +152,7 @@ def dedupe(dataset: DocumentDataset) -> DocumentDataset:
     deduplicator = ExactDuplicates(id_field="id", text_field="text", hash_method="md5")
     # Find the duplicates
     duplicates = deduplicator(dataset)
-    deduped = deduplicator.remove(dataset, duplicates)
-    return DocumentDataset(deduped)
+    return deduplicator.remove(dataset, duplicates)
 
 
 def run_curation_pipeline(args: argparse.Namespace, jsonl_dir: str) -> None:
@@ -173,7 +172,8 @@ def run_curation_pipeline(args: argparse.Namespace, jsonl_dir: str) -> None:
         keep_extensions="jsonl",
     )
     print("Reading the data...")
-    orig_dataset = DocumentDataset.read_json(files, add_filename=True)
+    # We don't read with add_filename because it already exists in the jsonl files.
+    orig_dataset = DocumentDataset.read_json(files)
     dataset = orig_dataset
 
     curation_steps = Sequential(