add split argument to Generator

[piercus]

Actual

When creating a multi-split dataset using generators like

datasets.DatasetDict({
  "val": datasets.Dataset.from_generator(
      generator=generator_val,
      features=features
  ),
  "test": datasets.Dataset.from_generator(
      generator=generator_test,
      features=features,
  )
})

It displays (for both test and val)

Generating train split

Expected

I would like to be able to improve this behavior by doing

datasets.DatasetDict({
  "val": datasets.Dataset.from_generator(
      generator=generator_val,
      features=features,
      split="val"
  ),
  "test": datasets.Dataset.from_generator(
      generator=generator_test,
      features=features,
      split="test"
  )
})

It would display

Generating val split

and

Generating test split

Proposal

Current PR is adding an explicit split argument and replace the implicit "train" split in the following classes/function :

• Generator
• from_generator
• AbstractDatasetInputStream
• GeneratorDatasetInputStream

Please share your feedbacks

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[albertvillanova]

Thanks for your proposed contribution @piercus.

This is a nice one!

Some comments below. Basically, I would propose to define the split parameter just as an attribute of GeneratorConfig instead of Generator, GeneratorDatasetInputStream, AbstractDatasetInputStream and SqlDatasetReader.

[albertvillanova]

This docstring should go below <Added version="2.7.0"/>, because the version added tag corresponds to the num_proc parameter above split.

I would suggest to align its type with the rest of the code as: ([`NamedSplit`], defaults to `Split.TRAIN`).

I would also add a specific version added tag for the split parameter: . We may eventually change this depending on the next release.

[piercus]

just used <Added version="2.21.0"/>, please cross-check

[albertvillanova]

Same comments as before.

[piercus]

I also added <Added version="2.21.0"/> please cross-check

[albertvillanova]

I would add a specific test_dataset_from_generator_split with a parametrized split values, such as not passing any value, passing NamedSplit("train"), passing literal "train", passing other NamedSplit, etc.

[piercus]

test_dataset_from_generator_split added, still i have impacted _check_generator_dataset to share the same generic check everywhere

[piercus]

@albertvillanova thanks for the review, please take a look

[albertvillanova]

Thanks! Good work.

Just a fix of the non-passing test and a nit.

[piercus]

@albertvillanova please take a look

[albertvillanova]

Thanks for your contribution!

Note the CI action to generate the docs is failing due to an unrelated issue: https://github.com/huggingface/datasets/actions/runs/9887484572/job/27309892176?pr=7015

Therefore, if we do not want to break the generation of docs, this other PR should be merged before yours:

• Fix doc generation when NamedSplit is used as parameter default value #7036