Dataset.from_generator raises with sharded gen_args

[hartmans]

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

#!/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

[mariosasko]

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):

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')

[hartmans]

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.

[mariosasko]

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.

[hartmans]

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?

[hartmans]

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.

[hartmans]

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.