Why I try to replace the transformers TextDataset class with datasets Dataset class

I stumbled upon this when I tried to make the train_dataset of the Transformers Trainer class from a text file, see How can you get a Huggingface fine-tuning model with the Trainer class from your own text where you can set the arguments for truncation and padding?.

The TextDataset of the transformers package is

  • buggy (next heading) and
  • outdated (overnext heading).

Transformers TextDataset drops the last block of the split text

The TextDataset class drops the last block of the text that was split into blocks by means of the block_size parameter, in the following example, 512 tokens (~ words and other things) per block:

from transformers import AutoTokenizer, TextDataset

model_name = "dbmdz/german-gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
file_path = './myfile.txt'

train_dataset = TextDataset(

If I check the last block, I see that it cuts the very last block that has the tail of the text. This code shows only the second last block, the last block gets dropped by the TextDataset class:


Instead, the Trainer class does not drop the last batch by default, but you see from this that there is such a parameter also for the Auto dataloader arguments of the Trainer class, see class transformers Training Arguments:

dataloader_drop_last (bool, optional, defaults to False) — Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) or not.

Transformers TextDataset is outdated

When I change the setting of a tokenizer and build the TextDataset object another time, sometimes a warning shows that you should take the Transformers datasets Dataset class instead.

Here is the warning (there are two warnings in it):

Warning 1:

> /srv/home/my_user/.local/lib/python3.9/site-packages/transformers/data/datasets/language_modeling.py:54:
> FutureWarning: This dataset will be removed from the library soon,
> preprocessing should be handled with the 🤗 Datasets library. You can
> have a look at this example script for pointers:
> https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm.py

Warning 2:

> warnings.warn( Token indices sequence length is longer than the
> specified maximum sequence length for this model (31482 > 512).
> Running this sequence through the model will result in indexing errors

Warning 2 is just from changing from one tokenizer to another, it comes from this line in the given link of the warning.

        if data_args.max_seq_length > tokenizer.model_max_length:
                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the "
                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."

It is enough to run the code again to get rid of warning 2. This question is only about warning 1 ("FutureWarning: This dataset will be removed...").


How do I replace the transformers Textdataset class with the datasets Dataset class so that the output is a dataset that can be the argument of the train_dataset parameter of the transformers Trainer class?


1 Answer 1


You need to from datasets import load_dataset, and even though the warning tells you so, you do not seem to need to import the Dataset class if you just want to run the Trainer on your own text file input. For me, the load_dataset module was enough. This code runs through, its output is a dataset that can build a fine-tuned model with the help of the Huggingface Transformers PyTorch Trainer class.

from transformers import AutoTokenizer
from datasets import load_dataset

model_name = "dbmdz/german-gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
file_path = './myfile.txt'
bln_truncation = False

dataset = load_dataset("text", data_files={"train": file_path})

block_size = 512
tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_function(examples):
    return tokenizer(
        examples["text"], padding="max_length", truncation=bln_truncation)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

The dataset must then be passed to the Trainer class like this:

trainer = Trainer(

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.