Transforming Datasets with map and filter

Batched map is the workhorse

Dataset.map(fn) applies a function to every example and adds or replaces columns with the dictionary it returns. With batched=True the function receives a batch of rows as a dictionary of lists, for example {"text": [...1000 strings...]}. That matters because fast tokenizers process a list of texts far faster than one text at a time. remove_columns drops columns the model does not need, such as raw text, and num_proc runs the work in several processes. filter(fn) keeps rows for which the function returns True, which is handy for removing empty or overly long samples. Results are cached by a fingerprint of the function and its inputs, so a second run reuses them. Finally, set_format("torch") or with_format("torch") makes indexing return PyTorch tensors. A common pattern is to tokenize without padding in map and let a data collator pad each batch dynamically during training, so short batches stay short.

Tokenize, clean and format

Padding is left to the data collator at training time.

from datasets import load_dataset
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
ds = load_dataset("stanfordnlp/imdb", split="train")

ds = ds.filter(lambda ex: len(ex["text"]) > 20)

def tokenize(batch):
    return tok(batch["text"], truncation=True, max_length=256)

ds = ds.map(tokenize, batched=True, remove_columns=["text"], num_proc=2)
ds = ds.rename_column("label", "labels")
ds.set_format("torch")
print(ds.column_names)  # ['labels', 'input_ids', 'attention_mask']

Forgetting batched=True

Without batched=True, map calls your function once per row and the tokenizer loses most of its speed advantage. If tokenizing a few hundred thousand reviews takes ages, this is the first thing to check.

त्वरित जाँच: With `ds.map(fn, batched=True)`, what does `fn` receive?

• One row as a dictionary of single values
• A PyTorch tensor of token ids
• A dictionary whose values are lists, one list per column for a batch of rows
• The whole dataset as a pandas DataFrame