Datasets

The datasets library ships with Transformers: pull Hub data, tokenize with map, then feed Trainer or a DataLoader. This chapter uses public rotten_tomatoes (binary movie reviews), shared with Fine-tuning.


load_dataset

from datasets import load_dataset

ds = load_dataset("rotten_tomatoes")
print(ds)
print(ds["train"][0])

Typical shape:

DatasetDict({
  train: Dataset({ features: ['text', 'label'], num_rows: 8530 })
  validation: Dataset({ ... num_rows: 1066 })
  test: Dataset({ ... num_rows: 1066 })
})

label is 0 / 1 (negative / positive). A slice such as load_dataset("rotten_tomatoes", split="train[:500]") is enough for a CPU smoke test.

Other sources: Hub org/name, local data/*.json, CSV. Gated datasets need hf auth login and a click-through on the website.


Features and labels

print(ds["train"].features)
print(ds["train"].features["label"].names)  # if ClassLabel
print(ds["train"].unique("label"))

Before you write a training script, confirm the text column (text here) and the label column (label). Those two names are the usual copy-paste bugs when you switch corpora.


map: batched tokenization

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True)

tokenized = ds.map(tokenize, batched=True, remove_columns=["text"])
tokenized = tokenized.rename_column("label", "labels")
print(tokenized["train"][0].keys())

Tips:

  • batched=True: many rows per call—much faster than a Python loop.
  • truncation=True: clip to model_max_length so long reviews do not OOM.
  • Leave dynamic padding to a DataCollator; avoid padding="max_length" on the whole set unless you truly want a fixed length.
  • Trainer looks for a labels column. Some versions also accept label; renaming is safer.

DataCollator: build batches

from transformers import DataCollatorWithPadding
from torch.utils.data import DataLoader

collator = DataCollatorWithPadding(tokenizer=tokenizer)
loader = DataLoader(
    tokenized["train"].with_format("torch"),
    batch_size=8,
    collate_fn=collator,
)
batch = next(iter(loader))
print({k: v.shape for k, v in batch.items()})

DataCollatorWithPadding pads to the longest sequence in the current batch, which is cheaper than a global max length. Generation jobs often use DataCollatorForLanguageModeling or DataCollatorForSeq2Seq.

[batch, seq] tensor shapes are covered in the PyTorch tutorial.


Format, cache, and small slices

small = tokenized["train"].shuffle(seed=42).select(range(256))
small.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

map caches on disk. After you change the function, pass load_from_cache_file=False or a new cache_file_name. Do not compute metrics inside map—use evaluate in Trainer’s compute_metrics (next chapter).


Dataset pages on the Hub

Every dataset has a card (e.g. rotten_tomatoes): license, fields, citation. Read the license before production; academic-only corpora do not belong in a commercial product.

You can ds.push_to_hub("your-name/my-reviews") to publish a preprocessed copy (login required).


Next steps

评论