datasets map() handles all data at a stroke and takes long time
State: closed | Author: BruceStayHungry (NONE)

**1. Background**

Huggingface datasets package advises using `map()`  to process data in batches. In the example code on pretraining masked language model, they use `map()` to tokenize all data at a stroke before the train loop. 

The corresponding code:
```
with accelerator.main_process_first():
    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        batched=True,
        num_proc=args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not args.overwrite_cache,
        desc="Running tokenizer on every text in dataset"
    )
```

**2. The problem**

Thus, when I try the same pertaining code with a much larger corpus, it takes quite a long time to tokenize.

Also, we can choose to tokenize data in `data-collator`. In this way, the program only tokenizes one batch in the next training step and avoids getting stuck in tokenization.

**3. My question**

As described above, my questions are:
* **Which is better? Process in `map()` or in `data-collator`**
* **Why huggingface advises `map()` function?** There should be some advantages to using `map()`


Thanks for your answers!

Comments: Hi ! Interesting question ;)

> Which is better? Process in map() or in data-collator

As you said, both can be used in practice: map() if you want to preprocess before training, or a data-collator (or the equivalent `dataset.set_transform`) if you want to preprocess on-the-fly during training. Both options are great and really depend on your case.

To choose between the two, here are IMO the main caveats of each approach:
- if your preprocessing takes too much CPU for example, using a data-collator may slow down your training and your GPUs may not work at full speed
- on the other hand, map() may take a lot of time and disk space to run if your dataset is too big.

> Why huggingface advises map() function? There should be some advantages to using map()

To get the best throughput when training a model, it is often recommended to preprocess your dataset before training. Note that preprocessing may include other steps before tokenization such as data filtering, cleaning, chunking etc. which are often done before training.

Thanks for your clear explanation @lhoestq ! 
> *  if your preprocessing takes too much CPU for example, using a data-collator may slow down your training and your GPUs may not work at full speed
> * on the other hand, map() may take a lot of time and disk space to run if your dataset is too big.

I really agree with you. There should be some trade-off between processing before and during the train loop.
Besides, I find `map()` function can cache the results once it has been executed. Very useful!

I'm closing this issue if you don't mind, feel free to reopen if needed ;)

@lhoestq How to preprocess on-the-fly during training？my data is about 1w hours, when I use map to  preprocess, and It's not finished yet， but all disk space(2T) is full.

Hi ! You can do that using `set_transform`, see https://huggingface.co/docs/datasets/process#format-transform for more info :)

unfortunately , it not work.

Could you share more details ?