Workaround for training models with really big text files
Author: finiteautomata
Pull request: no


# 🚀 Feature request

Provide a workaround for `run_mlm.py` when working with big text files. 

## Motivation

I'm trying to train a `RoBERTa` model with a lot of big text files (~ 50GB of text). When doing so, I'm facing two obstacles:

1. Tokenization beforehand creates a lot of cache storage (See #10204), so one has to resort to on-the-fly tokenization using `set_transform`
2.  `datasets` is quite slow when working with really big text files (see https://github.com/huggingface/datasets/issues/2210, https://github.com/huggingface/datasets/issues/2252). Some fixes have been proposed but to my knowledge the issue persists

Can you provide an example on how to workaround these two issues? I suppose that using a custom torch Dataset could (temporarily) fix this.

## Your contribution

Inspired by https://github.com/huggingface/transformers/issues/10278#issuecomment-805245903, I replaced the `datasets` class by this one. As far as I could see, time improves vs using `datasets` with `set_transform`, but I'm not really sure if it is optimal, particularly regarding parallelism (I'm running this script with `python xla_spawn.py`)

```python
from torch.utils.data import IterableDataset

class BatchProcessedDataset(IterableDataset):
    def __init__(self, files, tokenizer, batch_size=4096, limit=-1):
        self.files = files
        self.batch_size = batch_size
        self.tokenizer = tokenizer
        self.limit = limit

    def __iter__(self):
        num_iter = 0
        for file_path in self.files:
            with open(file_path) as f:

                next_batch = [x.strip("\n") for _, x in zip(range(self.batch_size), f)]

                while next_batch:
                    tokenized_batch = self.tokenizer(next_batch, padding='max_length', truncation=True, return_special_tokens_mask=True)
                    for encoding in tokenized_batch.encodings:
                        if num_iter == self.limit:
                            return
                        yield {
                            "input_ids": encoding.ids,
                            "token_type_ids": encoding.type_ids,
                            "attention_mask": encoding.attention_mask,
                            "special_tokens_mask": encoding.special_tokens_mask
                        }
                        num_iter += 1
                    next_batch = [x.strip("\n") for _, x in zip(range(self.batch_size), f)]
```

Comments:
cc @sgugger @lhoestq

The cache issue should be mostly fixed, now that datasets stores the tokenized inputs with the right precision. If it's not, it should be discussed on the Datasets repo.
The second issue should also be discussed on the Datasets repo.

As mentioned on the main README, the examples provided here are just this: examples. You can adapt them to your use case (as you did) but we leave them as generic as possible on purpose.

Thanks @sgugger for your answer; it's true you can't just add every possible example there, and of course I don't intend to discuss `datasets` issues here. 

I share a gist with a modified version of `run_mlm.py` in case anyone is facing the same problem. 

https://gist.github.com/finiteautomata/bef480d508d12e2028fdeae19a92b350