Adds option to shuffle input files in readers (#128)

* add todo

* rewrote c4 quality filter to match the official implementation

* add tokenizer_language

* added docstring

* fix actually keeping the document

* simplified a bit sent dedup

* small adhoc bugfixes

* made sentence dedup s2 parallelizable; bugfixes and other improvements

* reduce mem footprint

* add debug msg

* small perf improvs

* change folder structure

* added line_buffering option for dedup. Changed/optimized sentence dedup

* fix test

* updated example

* perf improvements for sent dedup filter

* remove print

* nit

* bugfixes and perf improvements

* added check

* use compiled struct

* added shuffle_files to readers

src/datatrove/pipeline/readers/base.py

@@ -1,3 +1,4 @@
+import random
 from abc import abstractmethod
 from contextlib import nullcontext
 from typing import Callable
@@ -134,6 +135,7 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         """
 
@@ -147,11 +149,14 @@ def __init__(
             default_metadata: a dictionary with any data that should be added to all sample's metadata
             recursive: whether to search recursively for files
             glob_pattern: pattern that all files must match exactly to be included (relative to data_folder)
+            shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
         """
         super().__init__(limit, progress, adapter, text_key, id_key, default_metadata)
         self.data_folder = get_datafolder(data_folder)
         self.recursive = recursive
         self.glob_pattern = glob_pattern
+        self.shuffle_files = shuffle_files
 
     def get_document_from_dict(self, data: dict, source_file: str, id_in_file: int):
         document = super().get_document_from_dict(data, source_file, id_in_file)
@@ -219,6 +224,8 @@ def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1
                 raise RuntimeError(f"No files found on {self.data_folder.path}!")
             # otherwise just a warning
             logger.warning(f"No files found on {self.data_folder.path} for {rank=}")
+        if self.shuffle_files:
+            random.shuffle(files_shard)
         for doc in self.read_files_shard(files_shard):
             self.update_doc_stats(doc)
             yield doc

src/datatrove/pipeline/readers/csv.py

@@ -22,6 +22,8 @@ class CsvReader(BaseDiskReader):
         default_metadata: default metadata to add to all documents
         recursive: if True, will read files recursively in subfolders (default: True)
         glob_pattern: a glob pattern to filter files to read (default: None)
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "🔢 Csv"
@@ -38,9 +40,19 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         super().__init__(
-            data_folder, limit, progress, adapter, text_key, id_key, default_metadata, recursive, glob_pattern
+            data_folder,
+            limit,
+            progress,
+            adapter,
+            text_key,
+            id_key,
+            default_metadata,
+            recursive,
+            glob_pattern,
+            shuffle_files,
         )
         self.compression = compression
         self.empty_warning = False

src/datatrove/pipeline/readers/ipc.py

@@ -20,6 +20,8 @@ class IpcReader(BaseDiskReader):
         default_metadata: default metadata to add to all documents
         recursive: if True, will read files recursively in subfolders (default: True)
         glob_pattern: a glob pattern to filter files to read (default: None)
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "🪶 Ipc"
@@ -37,9 +39,19 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         super().__init__(
-            data_folder, limit, progress, adapter, text_key, id_key, default_metadata, recursive, glob_pattern
+            data_folder,
+            limit,
+            progress,
+            adapter,
+            text_key,
+            id_key,
+            default_metadata,
+            recursive,
+            glob_pattern,
+            shuffle_files,
         )
         self.stream = stream
         # TODO: add option to disable reading metadata (https://github.com/apache/arrow/issues/13827 needs to be addressed first)

src/datatrove/pipeline/readers/jsonl.py

@@ -25,6 +25,8 @@ class JsonlReader(BaseDiskReader):
         default_metadata: default metadata to add to all documents
         recursive: if True, will read files recursively in subfolders (default: True)
         glob_pattern: a glob pattern to filter files to read (default: None)
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "🐿 Jsonl"
@@ -41,9 +43,19 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         super().__init__(
-            data_folder, limit, progress, adapter, text_key, id_key, default_metadata, recursive, glob_pattern
+            data_folder,
+            limit,
+            progress,
+            adapter,
+            text_key,
+            id_key,
+            default_metadata,
+            recursive,
+            glob_pattern,
+            shuffle_files,
         )
         self.compression = compression
 

src/datatrove/pipeline/readers/parquet.py

@@ -22,6 +22,8 @@ class ParquetReader(BaseDiskReader):
         default_metadata: default metadata to add to all documents
         recursive: if True, will read files recursively in subfolders (default: True)
         glob_pattern: a glob pattern to filter files to read (default: None)
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "📒 Parquet"
@@ -40,9 +42,19 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         super().__init__(
-            data_folder, limit, progress, adapter, text_key, id_key, default_metadata, recursive, glob_pattern
+            data_folder,
+            limit,
+            progress,
+            adapter,
+            text_key,
+            id_key,
+            default_metadata,
+            recursive,
+            glob_pattern,
+            shuffle_files,
         )
         self.batch_size = batch_size
         self.read_metadata = read_metadata

src/datatrove/pipeline/readers/warc.py

@@ -25,6 +25,8 @@ class WarcReader(BaseDiskReader):
         default_metadata: default metadata to add to all documents
         recursive: if True, will read files recursively in subfolders (default: True)
         glob_pattern: a glob pattern to filter files to read (default: None)
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "🕷 Warc"
@@ -42,10 +44,20 @@ def __init__(
         default_metadata: dict = None,
         recursive: bool = True,
         glob_pattern: str | None = None,
+        shuffle_files: bool = False,
     ):
         self.compression = compression
         super().__init__(
-            data_folder, limit, progress, adapter, text_key, id_key, default_metadata, recursive, glob_pattern
+            data_folder,
+            limit,
+            progress,
+            adapter,
+            text_key,
+            id_key,
+            default_metadata,
+            recursive,
+            glob_pattern,
+            shuffle_files,
         )
 
     def read_file(self, filepath: str):