Fix shards in IterableDataset.from_generator

src/datasets/iterable_dataset.py

@@ -764,6 +764,8 @@ def from_generator(
             generator (:obj:`Callable`): A generator function that `yields` examples.
             features (:class:`Features`, optional): Dataset features.
             gen_kwargs(:obj:`dict`, optional): Keyword arguments to be passed to the `generator` callable.
+                You can define a sharded iterable dataset by passing the list of shards in `gen_kwargs`.
+                This can be used to improve shuffling and when iterating over the dataset with multiple workers.
 
         Returns:
             :class:`IterableDataset`
@@ -777,6 +779,20 @@ def from_generator(
         ...
         >>> ds = IterableDataset.from_generator(gen)
         ```
+
+        ```py
+        >>> def gen(shards):
+        ...     for shard in shards:
+        ...         with open(shard) as f:
+        ...             for line in f:
+        ...                 yield {"line": line}
+        ...
+        >>> shards = [f"data{i}.txt" for i in range(32)]
+        >>> ds = IterableDataset.from_generator(gen, gen_kwargs={"shards": shards})
+        >>> ds = ds.shuffle(seed=42, buffer_size=10_000)  # shuffles the shards order + uses a shuffle buffer
+        >>> from torch.utils.data import DataLoader
+        >>> dataloader = .DataLoader(ds.with_format("torch"), num_workers=4)  # give each worker a subset of 32/4=8 shards
+        ```
         """
         from .io.generator import GeneratorDatasetInputStream
 

src/datasets/packaged_modules/generator/generator.py

@@ -24,8 +24,8 @@ def _info(self):
         return datasets.DatasetInfo(features=self.config.features)
 
     def _split_generators(self, dl_manager):
-        return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={})]
+        return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs=self.config.gen_kwargs)]
 
-    def _generate_examples(self):
-        for idx, ex in enumerate(self.config.generator(**self.config.gen_kwargs)):
+    def _generate_examples(self, **gen_kwargs):
+        for idx, ex in enumerate(self.config.generator(**gen_kwargs)):
             yield idx, ex

tests/test_iterable_dataset.py

@@ -592,6 +592,18 @@ def gen():
     assert list(dataset) == data
 
 
+def test_iterable_dataset_from_generator_with_shards():
+    def gen(shard_names):
+        for shard_name in shard_names:
+            for i in range(10):
+                yield {"shard_name": shard_name, "i": i}
+
+    shard_names = [f"data{shard_idx}.txt" for shard_idx in range(4)]
+    dataset = IterableDataset.from_generator(gen, gen_kwargs={"shard_names": shard_names})
+    assert isinstance(dataset, IterableDataset)
+    assert dataset.n_shards == len(shard_names)
+
+
 @require_torch
 def test_iterable_dataset_factory_torch_integration():
     import torch