TopolearnEfficient, correct batching is half the work of fine-tuning and a common source of shape errors.
Before training, turn text columns into model inputs. The idiom is a batched map with the tokenizer:
def tokenize(batch):
return tok(batch["text"], truncation=True, max_length=256)
ds = ds.map(tokenize, batched=True, remove_columns=["text"])Notice there is no padding here. Padding every example to a global maximum wastes compute. Instead, a data collator pads each batch to its own longest sequence at load time (dynamic padding):
DataCollatorWithPadding(tok)for classification and similar tasks.DataCollatorForTokenClassification(tok)also padslabelswith-100.DataCollatorForLanguageModeling(tok, mlm=False)builds causal LM labels frominput_ids(setmlm=Truefor masked LM with random masking).DataCollatorForSeq2Seq(tok, model=model)for encoder-decoder models.
The label value -100 is special: PyTorch's cross-entropy (and Transformers' loss functions) ignore it, so padding positions and prompt tokens do not contribute to the loss.
Name the target column labels (or label, which collators rename) because models compute the loss when they receive labels. Remove raw string columns; the Trainer drops unused columns automatically, but explicit removal avoids surprises.
Group similar lengths together (train_sampling_strategy="group_by_length" in TrainingArguments) to reduce padding further.
Example
from datasets import Dataset
from transformers import AutoTokenizer, DataCollatorWithPadding
tok = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
ds = Dataset.from_dict({"text": ["short", "a much longer example sentence"],
"label": [0, 1]})
def tokenize(batch):
return tok(batch["text"], truncation=True, max_length=128)
ds = ds.map(tokenize, batched=True, remove_columns=["text"])
print([len(x) for x in ds["input_ids"]]) # [3, 7] (no padding yet)
collator = DataCollatorWithPadding(tok)
batch = collator([ds[0], ds[1]])
print(batch["input_ids"].shape) # torch.Size([2, 7]) padded per batch
print(batch["labels"]) # tensor([0, 1]) 'label' renamedRemember
- Tokenize with a batched
mapand truncation, but do not pad there - Data collators pad each batch dynamically to its longest sequence
- Label -100 is ignored by the loss (padding, prompt tokens)
Watch out for
- Padding everything to max_length in map, which makes training far slower than dynamic padding.