Hello,
I’m trying to use the load_dataset.set_transform(…) function along with DataCollatorForLanguageModeling  3 and the Trainer 1 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.
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.