Dataset.from_generator raises with sharded gen_args
Opened by hartmans (closed)

### Describe the bug

According to the docs of Datasets.from_generator:
```
        gen_kwargs(`dict`, *optional*):
            Keyword arguments to be passed to the `generator` callable.
            You can define a sharded dataset by passing the list of shards in `gen_kwargs`.
```
So I'd expect that if gen_kwargs was a list, then my generator would be called once for each element in the list with the dict in the list for that element.
It doesn't work that way though.

### Steps to reproduce the bug

```python
#!/usr/bin/python

from pathlib import Path
import datasets

def process_yaml(file):
    yield dict(example=42)


if __name__ == '__main__':
    import sys
    dir = Path(sys.argv[0]).parent
    ds = datasets.Dataset.from_generator(process_yaml, gen_kwargs=[{'file':f} for f in dir.glob('*.yml')],
        )
    ds.to_json('training.jsonl')
    

```
```
Generating train split: 0 examples [00:00, ? examples/s]
Traceback (most recent call last):
  File "/tmp/dataset_bug.py", line 13, in <module>
    ds = datasets.Dataset.from_generator(process_yaml, gen_kwargs=[{'file':f} for f in dir.glob('*.yml')],
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/arrow_dataset.py", line 1072, in from_generator
    ).read()
      ^^^^^^
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/io/generator.py", line 47, in read
    self.builder.download_and_prepare(
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/builder.py", line 954, in download_and_prepare
    self._download_and_prepare(
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/builder.py", line 1717, in _download_and_prepare
    super()._download_and_prepare(
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/builder.py", line 1049, in _download_and_prepare
    self._prepare_split(split_generator, **prepare_split_kwargs)
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/builder.py", line 1555, in _prepare_split
    for job_id, done, content in self._prepare_split_single(
  File "/home/hartmans/ai/venv/lib/python3.11/site-packages/datasets/builder.py", line 1656, in _prepare_split_single
    generator = self._generate_examples(**gen_kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: datasets.packaged_modules.generator.generator.Generator._generate_examples() argument after ** must be a ```
mapping, not list


### Expected behavior

I would expect that process_yaml would be called once for each yaml file in the directory where the script is run.
I also tried with the list being in gen_kwargs, but in that case process_yaml gets called with a list.


### Environment info

- `datasets` version: 2.14.6.dev0 (git commit 0cc77d7f45c7369; also tested with 2.14.0)
- Platform: Linux-6.1.0-10-amd64-x86_64-with-glibc2.36
- Python version: 3.11.2
- Huggingface_hub version: 0.16.4
- PyArrow version: 12.0.1
- Pandas version: 2.0.3


Comments:
`gen_kwargs` should be a `dict`, as stated in the docstring, but you are passing a `list`.

So, to fix the error, replace the list of dicts with a dict of lists (and slightly modify the generator function):
```python
from pathlib import Path
import datasets

def process_yaml(files):
    for f in files:
        # process
        yield dict(...)


if __name__ == '__main__':
    import sys
    dir = Path(sys.argv[0]).parent
    ds = datasets.Dataset.from_generator(process_yaml, gen_kwargs={'files': [f for f in dir.glob('*.yml')]})
    ds.to_json('training.jsonl')
```

That runs, and because my dataset is small, it's what I did to get past the problem.
However, it does not produce a sharded dataset.  From the doc string I expect there ought to be a way to call from_generator such that num_shards in the resulting data set is equal to the number of items in the list.
The part of the doc string that your suggestion is not responsive to is:
`            You can define a sharded dataset by passing the list of shards in *g
en_kwargs*.
`

What your suggestion does is calls the generator once, with the list argument, and produces a single shard dataset.


The sharding mentioned here refers to using this function with `num_proc` (multiprocessing splits the `kwargs` into shards and passes them to the generator function)

> That runs, and because my dataset is small, it's what I did to get past the problem.

`from_generator` generates a memory-mapped dataset (can be larger than RAM), so the dataset size should not be an issue unless the generator function's implementation does not properly free the memory.


It sounds like you are saying that num_proc affects the form of gen_kwargs.
Are you saying that for non-zero num_proc gen_kwargs should be a list whose length is the same as num_proc?
Or are you saying that for non-zero num_proc, gen_kwargs should be a dict whose elements are lists the length of num_proc?


I ran some tests.  So, it looks like with num_proc greater than 1, gen_kwargs is expected to be a dict of lists.  It calls the generator also with a dict of lists, but the lists are split.
I.E. if my original has `gen_kwargs=dict(a=[0,1,2])`, then my generator might get called with `gen_kwalrgs=dict([0])`.
That all makes sense, but I definitely think there is room for improvement in the doc string here.
In order to suggest improvements to the doc string, I need to look at how the gen_kwargs are split, and figure out if:
* num_proc needs to exactly equal the length of the lists
* num_proc needs to evenly divide the length of the lists
* Or there's no required relationship.
I'll look into that and then propose an improved doc string if no one else gets to it first.

Okay, that was fun; I took a dive through the dataset code and feel like I have a much better understanding.
Here is my understanding of the behavior:
* max_proc is an upper limit on the number of shards that `from_generator` produces
* If `max_proc` is greater than 1, then all lists in *gen_kwargs* must be the same length
* If the lists in *gen_kwargs* are shorter than *num_proc* elements, *num_proc* will be reduced and a warning produced.  Put another way, `min(list_length, num_shards)` shards will be produced
* The members of the lists in *gen_kwargs* will be  partitioned among the created jobs.
To validate the above, take a look at
`_number_of_shards_in_gen_kwargs` and `_distribute_shards`  and `_split_gen_kwargs` in utils/sharding.py.
I've also chased down starting at *from_generator* all the way through to GeneratorBuilder and the calls to the functions in sharding.py.
Tomorrow I'll take a look at the contributing guidelines and see what's involved in putting together a PR to improve the doc string.