Datasets: [Feature request] Add dataset.ragged_map() function for many-to-many transformations

[jarednielsen]

dataset.map() enables one-to-one transformations. Input one example and output one example. This is helpful for tokenizing and cleaning individual lines.
dataset.filter() enables one-to-(one-or-none) transformations. Input one example and output either zero/one example. This is helpful for removing portions from the dataset.
However, some dataset transformations are many-to-many. Consider constructing BERT training examples from a dataset of sentences, where you map ["a", "b", "c"] -> ["a[SEP]b", "a[SEP]c", "b[SEP]c", "c[SEP]b", ...]

I propose a more general ragged_map() method that takes in a batch of examples of length N and return a batch of examples M. This is different from the map(batched=True) method, which takes examples of length N and returns a batch of length N, processing individual examples in parallel. I don't have a clear vision of how this would be implemented efficiently and lazily, but would love to hear the community's feedback on this.

My specific use case is creating an end-to-end ELECTRA data pipeline. I would like to take the raw WikiText data and generate training examples from this using the ragged_map() method, then export to TFRecords and train quickly. This would be a reproducible pipeline with no bash scripts. Currently I'm relying on scripts like https://github.com/google-research/electra/blob/master/build_pretraining_dataset.py, which are less general.

[thomwolf]

Actually map(batched=True) can already change the size of the dataset.
It can accept examples of length N and returns a batch of length M (can be null or greater than N).

I'll make that explicit in the doc that I'm currently writing.

[jarednielsen]

You're two steps ahead of me :) In my testing, it also works if M < N.

A batched map of different length seems to work if you directly overwrite all of the original keys, but fails if any of the original keys are preserved.

For example,

# Create a dummy dataset
dset = load_dataset("wikitext", "wikitext-2-raw-v1")["test"]
dset = dset.map(lambda ex: {"length": len(ex["text"]), "foo": 1})

# Do an allreduce on each batch, overwriting both keys
dset.map(lambda batch: {"length": [sum(batch["length"])], "foo": [1]})
# Dataset(schema: {'length': 'int64', 'foo': 'int64'}, num_rows: 5)

# Now attempt an allreduce without touching the `foo` key
dset.map(lambda batch: {"length": [sum(batch["length"])]})
# This fails with the error message below

  File "/path/to/nlp/src/nlp/arrow_dataset.py", line 728, in map
    arrow_schema = pa.Table.from_pydict(test_output).schema
  File "pyarrow/io.pxi", line 1532, in pyarrow.lib.Codec.detect
  File "pyarrow/table.pxi", line 1503, in pyarrow.lib.Table.from_arrays
  File "pyarrow/public-api.pxi", line 390, in pyarrow.lib.pyarrow_wrap_table
  File "pyarrow/error.pxi", line 85, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Column 1 named foo expected length 1 but got length 2

Adding the remove_columns=["length", "foo"] argument to map() solves the issue. Leaving the above error for future visitors. Perfect, thank you!