dataset.map(): Transformations that Cache

~30 min · datasets, map

Level 0Scout

0 XP0/50 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

Why map() is the workhorse

ds.map(fn) applies a function row-by-row (or batch-by-batch) and returns a new dataset. It looks like a Pandas .apply() but with two crucial differences: it's cached (results memoized to disk keyed by function source) and it's multi-process (num_proc=N).

Batched and unbatched

Default is per-row. For tokenization, NLP preprocessing, and most heavy ops, batched=True is dramatically faster — the function receives a dict of column lists, returns the same shape. The library splits batches automatically.

The remove_columns gotcha

If your map function adds new columns (typical: tokenize the text column to input_ids + attention_mask), the original column stays unless you set remove_columns=['text']. Trainers don't accept extra string columns, so this is the most common cause of "the trainer is broken" tickets.

Code

Batched tokenization·python

from datasets import load_dataset
from transformers import AutoTokenizer

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

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

tokenized = ds.map(
    tokenize,
    batched=True,
    batch_size=1000,
    num_proc=4,
    remove_columns=["text"],   # critical for trainer compatibility
)
print(tokenized)
print(tokenized[0].keys())  # input_ids, attention_mask, label

Cache invalidation when you change the function·python

# datasets caches keyed on function source.
# If you tweak the function and don't change kwargs, results are STILL stale.
# Force a re-run with load_from_cache_file=False:

tokenized = ds.map(tokenize, batched=True, load_from_cache_file=False)

# Or inspect cache state:
print(ds.cache_files)  # list of (filename, range) tuples

External links

Exercise

Tokenize stanfordnlp/imdb train split with both batched=False and batched=True, batch_size=1000. Time both. Then run again to verify cache hits (should be near-instant). Now modify the tokenizer max_length, re-run, and confirm the cache invalidates.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousExploring with Pandas, Arrow, and the Viewer Next →Streaming Datasets and IterableDataset

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.