# Using load\_dataset.set\_transform() function along with Trainer class

**URL:** <https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758>\
**Category:** 🤗Datasets\
**Created:** [April 26, 2021, 12:48am UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758 "2021-04-26T00:48:12Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![tkon3](https://avatars.discourse-cdn.com/v4/letter/t/9dc877/32.png) [@tkon3](https://discuss.huggingface.co/u/tkon3)\
**Post date:** [April 26, 2021, 12:48am UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758/1 "2021-04-26T00:48:12Z")

</div>

Hello,

I’m trying to use the load\_dataset.set\_transform(…) function along with [DataCollatorForLanguageModeling](https://github.com/huggingface/transformers/blob/master/src/transformers/data/data_collator.py) and the [Trainer](https://huggingface.co/transformers/_modules/transformers/trainer.html) class from the transformers library to pretrain a model. Since I have a large dataset, tokenization does not fit in RAM and using the .map() function uses way too much disk space (\> 500Go) which is limited in my case. So I need to tokenize on the fly.

While the set\_transform works as expected if I index the dataset, I don’t know why it fails when I plug it with a DataCollatorForLanguageModeling and a Trainer.

```auto
tokenizer = ...

    def encode(batch):
        return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=512, return_tensors="pt")

    train_dataset = load_dataset('text', data_files={'train': txt_train_dataset})
    train_dataset.set_transform(encode)

    validation_dataset = load_dataset('text', data_files={'validation': txt_validation_dataset})
    validation_dataset.set_transform(encode)

    print(train_dataset["train"][:3]) # works as expected: {'input_ids': ..., 'attention_mask': ...}

    data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=True, mlm_probability=0.15
        )

    trainer = Trainer(
            model=...,
            args=...,
            data_collator=data_collator,
            train_dataset=train_dataset["train"],
            eval_dataset=validation_dataset["validation"],
            compute_metrics=...,
        )

```

Im probably missing something.

---

<div class="post-metadata">

**Author:** ![lewtun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lewtun/32/30995_2.png) [@lewtun](https://discuss.huggingface.co/u/lewtun)\
**Post date:** [April 26, 2021, 11:10am UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758/2 "2021-04-26T11:10:44Z")

</div>

hey @tkon3, when you say “it fails” what do you mean exactly? there’s a related thread on `set_transform` here that might be useful: [Understanding set\_transform](https://discuss.huggingface.co/t/understanding-set-transform/3740)

---

<div class="post-metadata">

**Author:** ![tkon3](https://avatars.discourse-cdn.com/v4/letter/t/9dc877/32.png) [@tkon3](https://discuss.huggingface.co/u/tkon3)\
**Post date:** [April 26, 2021, 11:44am UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758/3 "2021-04-26T11:44:35Z")

</div>

There is a type change. The trainer class fails to loop over it.

```auto
File "/home/transformers_datasets_test/model/trainer.py", line 88, in encode
        return tokenizer(batch["text"], truncation=True, max_length=tokenizer.model_max_length)
    KeyError: 'text'

```

When I check batch instead of batch[“text”], I get an empty dict {} which is not expected

However:

```auto
train_dataset["train"][:3]

```

Works as expected outside the trainer class

---

<div class="post-metadata">

**Author:** ![lewtun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lewtun/32/30995_2.png) [@lewtun](https://discuss.huggingface.co/u/lewtun)\
**Post date:** [April 26, 2021, 1:48pm UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758/4 "2021-04-26T13:48:08Z")

</div>

ah i think that’s because the trainer removes unused columns like `text` by default (see [docs](https://huggingface.co/transformers/main_classes/trainer.html?highlight=label_names#trainingarguments)). what happens if you set `remove_unused_columns=False` in your `TrainingArguments`?

fyi there is a related issue here that might contain additional info for your use case: [ERROR WHEN USING SET\_TRANSFORM() · Issue #1867 · huggingface/datasets · GitHub](https://github.com/huggingface/datasets/issues/1867)

---

<div class="post-metadata">

**Author:** ![tkon3](https://avatars.discourse-cdn.com/v4/letter/t/9dc877/32.png) [@tkon3](https://discuss.huggingface.co/u/tkon3)\
**Post date:** [April 26, 2021, 2:34pm UTC](https://discuss.huggingface.co/t/using-load-dataset-set-transform-function-along-with-trainer-class/5758/5 "2021-04-26T14:34:50Z")

</div>

Thank you, its working now.

For those wondering how to train a large dataset with on the fly tokenization and a line by line .txt setup 🙂 :

```
model = ...
training_args = ...
max_length = ...
tokenizer = ...
txt_dataset = ...

def encode(batch):
    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=max_length)

train_dataset = load_dataset('text', data_files={'train': txt_dataset})
train_dataset.set_transform(encode)

data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm=True, mlm_probability=0.15
    )

training_args.remove_unused_columns = False

trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset["train"],
        compute_metrics=...,
    )
```