Datasets: How to augment data ?

[astariul]

Is there any clean way to augment data ?

For now my work-around is to use batched map, like this :

def aug(samples):
    # Simply copy the existing data to have x2 amount of data
    for k, v in samples.items():
        samples[k].extend(v)
    return samples

dataset = dataset.map(aug, batched=True)

[lhoestq]

Using batched map is probably the easiest way at the moment.
What kind of augmentation would you like to do ?

[astariul]

Some samples in the dataset are too long, I want to divide them in several samples.

[lhoestq]

Using batched map is the way to go then.
We'll make it clearer in the docs that map could be used for augmentation.

Let me know if you think there should be another way to do it. Or feel free to close the issue otherwise.

[astariul]

It just feels awkward to use map to augment data. Also it means it's not possible to augment data in a non-batched way.

But to be honest I have no idea of a good API...

[astariul]

Or for non-batched samples, how about returning a tuple ?

def aug(sample):
    # Simply copy the existing data to have x2 amount of data
    return sample, sample

dataset = dataset.map(aug)

It feels really natural and easy, but :

• it means the behavior with batched data is different
• I don't know how doable it is backend-wise

@lhoestq

[lhoestq]

As we're working with arrow's columnar format we prefer to play with batches that are dictionaries instead of tuples.
If we have tuple it implies to re-format the data each time we want to write to arrow, which can lower the speed of map for example.

It's also a matter of coherence, as we don't want users to be confused whether they have to return dictionaries for some functions and tuples for others when they're doing batches.