From a dataset to batches a trainer can eat

The three things the model needs

Whatever you are training, the loop wants tensors of the right shape with the right labels. For text that means: input_ids, attention_mask, and labels. Everything in this lesson is getting from a table of strings to those three.

Classification: labels are just a column

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

ds = ds.map(prep, batched=True, remove_columns=["text"])
ds = ds.rename_column("label", "labels")
ds.set_format("torch")

Three specific points. The column must be called labels, plural, because that is what the model’s forward signature expects: label silently becomes an unused column and the loss is never computed. Do not pad here; padding at map time pads everything to 512 and wastes compute. And set_format("torch") makes rows come back as tensors without copying the underlying Arrow data.

Dynamic padding, with a collator

from transformers import DataCollatorWithPadding
collator = DataCollatorWithPadding(tokenizer=tok)

The collator pads each batch to that batch’s longest item rather than to a global maximum. Combined with length-grouped batching (group_by_length=True in TrainingArguments) this is frequently a twofold speed-up for free, for the reasons the batching lesson gave.

Causal language modelling: labels are the inputs

from transformers import DataCollatorForLanguageModeling
collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)

With mlm=False, the collator copies input_ids into labels and sets padded positions to -100, which the loss ignores. The model shifts internally, so you do not shift yourself: doing it manually as well is a classic off-by-one that produces a model which predicts the current token instead of the next one and appears to learn beautifully while being useless.

With mlm=True it masks 15% of tokens instead, for BERT-style training.

Instruction tuning: mask the prompt

For supervised fine-tuning on instruction data, you usually do not want loss on the prompt: only on the response. Otherwise the model spends its capacity learning to generate the questions people ask it.

The mechanism is setting prompt positions in labels to -100:

def prep(example):
    prompt = tok.apply_chat_template(example["messages"][:-1],
                                     tokenize=False, add_generation_prompt=True)
    full = tok.apply_chat_template(example["messages"], tokenize=False)

    p_ids = tok(prompt, add_special_tokens=False)["input_ids"]
    f_ids = tok(full, add_special_tokens=False)["input_ids"]

    labels = list(f_ids)
    labels[:len(p_ids)] = [-100] * len(p_ids)
    return {"input_ids": f_ids, "labels": labels,
            "attention_mask": [1] * len(f_ids)}

Note add_special_tokens=False on both encodes: the chat template has already inserted the special tokens, and letting the tokenizer add a second beginning-of-sequence marker shifts every offset by one and breaks the masking silently. TRL’s SFTTrainer does all of this for you, which is why the training module reaches for it; doing it once by hand is how you understand what it is doing.

Packing, for pretraining

When training on plain text, padding wastes a lot. Packing concatenates documents and cuts fixed-length blocks:

def group(batch, block=1024):
    ids = sum(batch["input_ids"], [])
    n = (len(ids) // block) * block
    chunks = [ids[i:i+block] for i in range(0, n, block)]
    return {"input_ids": chunks, "labels": [c[:] for c in chunks]}

lm_ds = tokenised.map(group, batched=True, remove_columns=tokenised.column_names)

No padding at all, so every token in every batch is doing work. The cost is that documents bleed across block boundaries and attention crosses between unrelated texts. For pretraining this is standard and the effect is small; for instruction tuning it needs attention masking between packed examples, which the libraries handle and hand-rolled code usually does not.

Check the batch before you train

The cheapest hour you will ever save:

batch = collator([ds[i] for i in range(4)])
print({k: v.shape for k, v in batch.items()})
print(tok.decode(batch["input_ids"][0]))
lab = batch["labels"][0]
print(tok.decode([t for t in lab if t != -100]))

Read that decoded text. It is exactly what the model will see. Three things show up here and nowhere else: a doubled beginning-of-sequence token, a chat template that did not render, and prompt masking that covered the wrong span. Each of those costs a full training run if you find it afterwards.

Do this now

Print one collated batch for whatever you plan to train, decode both the input and the unmasked labels, and read them out loud. If they are not what you intended, you have just saved the cost of the run.