When resuming from checkpoint with Trainer using a streamed dataset, use the Datasets API to skip
Author: sinking-point
Pull request: no

### Feature request

Huggingface Datasets has a feature where you can instantiate datasets in streaming mode, so you don't have to download the whole thing onto your machine. The API has a skip function. The Transformers Trainer doesn't use this, it just iterates through all the batches to be skipped.

I propose Trainer checks whether the given dataset is a Datasets one in streaming mode, and if so, it uses the skip function.

### Motivation

I've been using the C4 dataset in streaming mode because of its size. Whenever I resume from a checkpoint, it takes a long time to skip. Around an hour for 200k batches. With this change, it should be effectively instant, which would save me a lot of time.

### Your contribution

I can make the change if no one else wants to. In which case, I'd like to be assured that the change will be reviewed and merged in a reasonable timeframe rather than being lost in the sea of pull requests.

Comments:
WDYT @lhoestq?

Unfortunately the `skip` function does the same. Though if you call `skip` before `map` it won't apply the preprocessing on the examples to skip and save time.

Therefore I don't think it using `skip` would have a big impact here

Thanks for your insight @lhoestq . Is there a technical reason the function is implemented that way? I assume random access is supported, given that shuffling of enormous datasets is allowed and fast.

Random access is not supported for streaming datasets. Shuffling is approximate by shuffling the dataset shards order, and using a shuffle buffer, see the documentation here: https://huggingface.co/docs/datasets/v2.4.0/en/stream#shuffle

Would it be possible to skip whole shards when a large number of samples need to be skipped?

In the general case we don't know in advance how many examples there are per shard (it depends on the data file format). In many cases we would need some extra metadata somewhere that says how many examples each shard contain.

For example the C4 dataset is made out gzipped JSON Lines files - you don't know in advance how many examples each shard contain, because you need to uncompress the data and count the EOL.

For certain file formats like Parquet or Arrow however, the number of examples is known for free, as metadata included in the file itself. So maybe for those specific formats we could do something

Ah, this needs deeper changes than I thought then. Still, the metadata would be nice to have. It could even be generated lazily when people stream the dataset, to avoid a bit upfront cost. Alternatively, the streamed dataset client could locally keep a cache of how many samples each shard it's iterated through contains. I'd prefer the former, so it doesn't matter if the cache is cleared or you switch machines.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.