Add shuffle option on huggingface reader

src/datatrove/pipeline/readers/huggingface.py

@@ -26,6 +26,8 @@ class HuggingFaceDatasetReader(BaseReader):
         text_key: key to use for the text in the default adapter (default: "text"). Ignored if you provide your own `adapter`
         id_key: key to use for the id in the default adapter (default: "id"). Ignored if you provide your own `adapter`
         default_metadata: default metadata to add to all documents
+        shuffle_files: shuffle the files within the returned shard. Mostly used for data viz. purposes, do not use
+            with dedup blocks
     """
 
     name = "🤗 HuggingFace"
@@ -44,13 +46,15 @@ def __init__(
         text_key: str = "text",
         id_key: str = "id",
         default_metadata: dict = None,
+        shuffle_files: bool = False,
     ):
         super().__init__(limit, skip, adapter, text_key, id_key, default_metadata)
         self.dataset = dataset
         self.dataset_options = dataset_options or {}
         self.batch_size = batch_size
         self.doc_progress = doc_progress
         self.streaming = streaming
+        self.shuffle_files = shuffle_files
 
     def get_document_from_dict(self, data: dict, source: str, id_in_file: int | str):
         document = super().get_document_from_dict(data, source, id_in_file)
@@ -93,6 +97,12 @@ def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1
             yield from data
         ds = load_dataset(self.dataset, **self.dataset_options, streaming=self.streaming)
 
+        if self.shuffle_files:
+            if not self.streaming:
+                ds = ds.shuffle(seed=42)
+            else:
+                ds = ds.shuffle(seed=42, buffer_size=1000)
+
         # In case the dataset is (Iterable)?DatasetDict, raise informative error
         if isinstance(ds, dict):
             raise ValueError(

tests/pipeline/test_hf_reader.py

@@ -14,6 +14,17 @@ def test_read_dataset(self):
         data = list(reader())
         self.assertEqual(len(data), 817)
 
+    def test_read_dataset_shuffle(self):
+        reader = HuggingFaceDatasetReader(
+            "truthful_qa",
+            dataset_options={"name": "generation", "split": "validation"},
+            text_key="question",
+            shuffle_files=True,
+        )
+        data = list(reader())
+        self.assertEqual(len(data[0].text), 69)
+        self.assertEqual(len(data[1].text), 46)
+
     def test_read_streaming_dataset(self):
         reader = HuggingFaceDatasetReader(
             "truthful_qa",
@@ -24,6 +35,18 @@ def test_read_streaming_dataset(self):
         data = list(reader())
         self.assertEqual(len(data), 817)
 
+    def test_read_streaming_dataset_shuffle(self):
+        reader = HuggingFaceDatasetReader(
+            "truthful_qa",
+            dataset_options={"name": "generation", "split": "validation"},
+            text_key="question",
+            streaming=True,
+            shuffle_files=True,
+        )
+        data = list(reader())
+        self.assertEqual(len(data[0].text), 69)
+        self.assertEqual(len(data[1].text), 46)
+
     def test_sharding(self):
         for shards in [1, 3]:
             for streaming in [True, False]: