Datasets library not suitable for huge text datasets
Author: avacaondata
Pull request: no

## Environment info

- `transformers` version: 4.3.0.dev0
- Platform: Ubuntu 18
- Python version: 3.7
- PyTorch version (GPU?): 1.7.1
- Tensorflow version (GPU?):
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help
@n1t0, @LysandreJik @patrickvonplaten @sgugger


## Information

Model I am using (Bert, XLNet ...):

BERT; but the problem arises before using it.

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Get 187GB (or other big-sized) dataset
2. Try to tokenize it.
3. Wait until your whole server crashes due to memory or disk (probably this last one).

The code used to train the tokenizer is:

```{python}
from argparse import ArgumentParser

from datasets import load_dataset

from transformers import AutoTokenizer

# luego se puede cambiar esto con load_from_disk

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--train_file", type=str, required=True, help="Train file with data.")
    parser.add_argument("--val_file", type=str, required=True, help="Val file with data.")
    parser.add_argument("--tokenizer_path", type=str, required=True, help="Path to tokenizer.")
    parser.add_argument("--num_workers", type=int, required=False, default=40, help="Number of workers for processing.")
    parser.add_argument("--save_path", type=str, required=True, help="Save path for the datasets.")
    args = parser.parse_args()
    data_files = {"train": args.train_file, "val": args.val_file}
    print("Loading dataset...")
    datasets = load_dataset("text", data_files=data_files)
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
    column_names = datasets["train"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    def tokenize_function(examples):
        # Remove empty lines
        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
        return tokenizer(
            examples["text"],
            padding="max_length",
            truncation=True,
            max_length=512,
            # We use this option because DataCollatorForLanguageModeling (see below) is more efficient when it
            # receives the `special_tokens_mask`.
            return_special_tokens_mask=True,
        )
    print("Tokenizing dataset...")
    tokenized_datasets = datasets.map(
        tokenize_function,
        batched=True,
        num_proc=args.num_workers,
        remove_columns=[text_column_name],
        load_from_cache_file=False,
    )
    print("Saving to disk...")
    tokenized_datasets.save_to_disk(args.save_path)
```

## Expected behavior

It's expected that the tokenized texts occupy less space than pure texts, however it uses approx 2 orders of magnitude more disk, making it unfeasible to pre-train a model using Datasets library and therefore examples scripts from Transformers.


Comments:
Yes, it does save the tokenization on disk to avoid loading everything in RAM.

> It's expected that the tokenized texts occupy less space than pure texts

You want the script to destroy your dataset?

The previous version of the scripts were loading every text in RAM so I would argue using Datasets makes it actually more possible to train with big datasets.
The examples are also just examples for quick fine-tuning/pre-training. If you are at the stage were your datasets don't even fit in disk space, some tweaks inside them are expected.

@sgugger The thing is that it uses soooo much disk. When it had processed 18.7GB of texts it was using 2.1TB of disk... My dataset fits in the disk, and I'm sure the tokenized dataset should fit too, as it's actually lighter than text (it's only a list of integers per text), what doesn't fit in disk are the pyarrow objects created by datasets. 
Sorry if I didn't explain myself clearly before, but it's not a problem with RAM or Memory, the main problem is that even when trying to pre-tokenize the whole dataset, saving it to disk for further use, it's not possible because the objects stored in disk by datasets library use 2 orders of magnitude more disk space than the original texts.

Mmm, pinging @lhoestq as this seems indeed a huge bump in memory now that I see the numbers (I understood twice the space, not 100 times more, sorry!)

Same discussion on the `datasets` repo: https://github.com/huggingface/datasets/issues/1825

> tokenizing a dataset using map takes a lot of space since it can store input_ids but also token_type_ids, attention_mask and special_tokens_mask. Moreover if your tokenization function returns python integers then by default they'll be stored as int64 which can take a lot of space. Padding can also increase the size of the tokenized dataset.

To go forward it would be nice to optimize what we actually need to be stored on disk. If some columns are not useful maybe they can be dropped (and possibly recreated on-the-fly if needed). We should also tweak the tensors precisions.

Another approach would be to tokenize on-the-fly for pretraining.

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/master/CONTRIBUTING.md) are likely to be ignored.