Title: Datasets library not suitable for huge text datasets.
Author: avacaondata (NONE)
State: closed
Created: 2021-02-05T11:06:50Z

Hi,

I'm trying to use datasets library to load a 187GB dataset of pure text, with the intention of building a Language Model. The problem is that from the 187GB it goes to some TB when processed by Datasets. First of all, I think the pre-tokenizing step (with tokenizer.map()) is not really thought for datasets this big, but for fine-tuning datasets, as this process alone takes so much time, usually in expensive machines (due to the need of tpus - gpus) which is not being used for training. It would possibly be more efficient in such cases to tokenize each batch at training time (receive batch - tokenize batch - train with batch), so that the whole time the machine is up it's being used for training. 
Moreover, the pyarrow objects created from a 187 GB datasets are huge, I mean, we always receive OOM, or No Space left on device errors when only 10-12% of the dataset has been processed, and only that part occupies 2.1TB in disk, which is so many times the disk  usage of the pure text (and this doesn't make sense, as tokenized texts should be lighter than pure texts).

Any suggestions??

Comments:
Hi ! Looks related to #861 

You are right: 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 make things more convenient, we recently added a "lazy map" feature that allows to tokenize each batch at training time as you mentioned. For example you'll be able to do
```python
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

def encode(batch):
    return tokenizer(batch["text"], padding="longest", truncation=True, max_length=512, return_tensors="pt")

dataset.set_transform(encode)
print(dataset.format)
# {'type': 'custom', 'format_kwargs': {'transform': <function __main__.encode(batch)>}, 'columns': ['idx', 'label', 'sentence1', 'sentence2'], 'output_all_columns': False}
print(dataset[:2])
# {'input_ids': tensor([[  101,  2572,  3217, ... 102]]), 'token_type_ids': tensor([[0, 0, 0, ... 0]]), 'attention_mask': tensor([[1, 1, 1, ... 1]])}

```
In this example the `encode` transform is applied on-the-fly on the "text" column.

This feature will be available in the next release 2.0 which will happen in a few days.
You can already play with it by installing `datasets` from source if you want :)

Hope that helps !

How recently was `set_transform` added? I am actually trying to implement it and getting an error:

`AttributeError: 'Dataset' object has no attribute 'set_transform'
`

I'm on v.1.2.1.

EDIT: Oh, wait I see now it's in the v.2.0. Whoops! This should be really useful.

Yes indeed it was added a few days ago. The code is available on master
We'll do a release next week :)

Feel free to install `datasets` from source to try it out though, I would love to have some feedbacks

For information: it's now available in `datasets` 1.3.0.
The 2.0 is reserved for even cooler features ;)

Hi @alexvaca0 , we have optimized Datasets' disk usage in the latest release v1.5.

Feel free to update your Datasets version
```shell
pip install -U datasets
```
and see if it better suits your needs.