[Feature request] Add dataset.ragged_map() function for many-to-many transformations
{'login': 'jarednielsen', 'id': 4564897, 'node_id': 'MDQ6VXNlcjQ1NjQ4OTc=', 'avatar_url': 'https://avatars.githubusercontent.com/u/4564897?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/jarednielsen', 'html_url': 'https://github.com/jarednielsen', 'followers_url': 'https://api.github.com/users/jarednielsen/followers', 'following_url': 'https://api.github.com/users/jarednielsen/following{/other_user}', 'gists_url': 'https://api.github.com/users/jarednielsen/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/jarednielsen/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/jarednielsen/subscriptions', 'organizations_url': 'https://api.github.com/users/jarednielsen/orgs', 'repos_url': 'https://api.github.com/users/jarednielsen/repos', 'events_url': 'https://api.github.com/users/jarednielsen/events{/privacy}', 'received_events_url': 'https://api.github.com/users/jarednielsen/received_events', 'type': 'User', 'site_admin': False}
[]
[]
["Actually `map(batched=True)` can already change the size of the dataset.\r\nIt can accept examples of length `N` and returns a batch of length `M` (can be null or greater than `N`).\r\n\r\nI'll make that explicit in the doc that I'm currently writing."
 'You\'re two steps ahead of me :) In my testing, it also works if `M` < `N`.\r\n\r\nA 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.\r\n\r\nFor example,\r\n```python\r\n# Create a dummy dataset\r\ndset = load_dataset("wikitext", "wikitext-2-raw-v1")["test"]\r\ndset = dset.map(lambda ex: {"length": len(ex["text"]), "foo": 1})\r\n\r\n# Do an allreduce on each batch, overwriting both keys\r\ndset.map(lambda batch: {"length": [sum(batch["length"])], "foo": [1]})\r\n# Dataset(schema: {\'length\': \'int64\', \'foo\': \'int64\'}, num_rows: 5)\r\n\r\n# Now attempt an allreduce without touching the `foo` key\r\ndset.map(lambda batch: {"length": [sum(batch["length"])]})\r\n# This fails with the error message below\r\n```\r\n\r\n```bash\r\n  File "/path/to/nlp/src/nlp/arrow_dataset.py", line 728, in map\r\n    arrow_schema = pa.Table.from_pydict(test_output).schema\r\n  File "pyarrow/io.pxi", line 1532, in pyarrow.lib.Codec.detect\r\n  File "pyarrow/table.pxi", line 1503, in pyarrow.lib.Table.from_arrays\r\n  File "pyarrow/public-api.pxi", line 390, in pyarrow.lib.pyarrow_wrap_table\r\n  File "pyarrow/error.pxi", line 85, in pyarrow.lib.check_status\r\npyarrow.lib.ArrowInvalid: Column 1 named foo expected length 1 but got length 2\r\n```\r\n\r\nAdding the `remove_columns=["length", "foo"]` argument to `map()` solves the issue. Leaving the above error for future visitors. Perfect, thank you!']
`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.


{'url': 'https://api.github.com/repos/huggingface/datasets/issues/360/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}