Hugging Face Fine-Tuning with Trainer

Fine-tuning with Hugging Face means taking a pretrained model and training it a little more on your own labeled data. Because the model already understands language, a few thousand examples and a few minutes on a GPU are often enough to get a strong classifier.

The Trainer class handles the training loop, evaluation, mixed precision, checkpoints and logging so you can focus on the data.

Step 1: Prepare the Data

from datasets import load_dataset
from transformers import AutoTokenizer

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

ds = load_dataset("stanfordnlp/imdb")
ds["train"] = ds["train"].shuffle(seed=42).select(range(4000))
ds["test"] = ds["test"].shuffle(seed=42).select(range(1000))

ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=256), batched=True)

Map: 100%|██████████| 4000/4000
Map: 100%|██████████| 1000/1000

Step 2: Load a Model with a Classification Head

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    name, num_labels=2, id2label={0: "neg", 1: "pos"}, label2id={"neg": 0, "pos": 1},
)

Some weights of DistilBertForSequenceClassification were not initialized ... ['classifier.bias', 'classifier.weight', ...]
You should probably TRAIN this model on a down-stream task ...

This warning is expected: the classification head is brand new and random. Training will fill it in.

Step 3: Metrics and TrainingArguments

import numpy as np
from transformers import TrainingArguments, DataCollatorWithPadding

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return {"accuracy": float((preds == labels).mean())}

args = TrainingArguments(
    output_dir="imdb-distilbert",
    num_train_epochs=2,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    bf16=True,            # use mixed precision on supported GPUs; remove on CPU
    report_to="none",
)
collator = DataCollatorWithPadding(tok)

(no output — objects created)

DataCollatorWithPadding pads each batch only to its longest example, which is faster than padding everything to max_length. Note the argument name eval_strategy; older tutorials use evaluation_strategy, which has been removed.

Step 4: Train and Evaluate

from transformers import Trainer

trainer = Trainer(
    model=model, args=args,
    train_dataset=ds["train"], eval_dataset=ds["test"],
    processing_class=tok, data_collator=collator,
    compute_metrics=compute_metrics,
)
trainer.train()
print(trainer.evaluate())

{'eval_loss': 0.29, 'eval_accuracy': 0.887, 'epoch': 2.0, ...}
(sample output — exact numbers vary per run)

processing_class replaces the older tokenizer= argument. It is saved alongside the model so the checkpoint is self-contained.

Step 5: Save and Use the Model

trainer.save_model("imdb-distilbert/final")

from transformers import pipeline
clf = pipeline("text-classification", model="imdb-distilbert/final")
print(clf("A beautiful, moving film."))

[{'label': 'pos', 'score': 0.98}]   (sample output)

Fine-Tuning Tips

• Start with a small subset to check everything runs, then scale up.
• Learning rates between 1e-5 and 5e-5 are typical for full fine-tuning.
• Watch eval loss: if it rises while training loss falls, you are overfitting.
• For LLMs, full fine-tuning is expensive; use LoRA (next lesson) instead.

Frequently asked questions

What is the Hugging Face Trainer?

A high-level class in transformers that runs the PyTorch training and evaluation loop for you, with support for mixed precision, multiple GPUs, checkpointing and logging.

How much data do I need to fine-tune a model?

For text classification, a few hundred to a few thousand labeled examples often give good results because the model is already pretrained.

Why do I get "Some weights were not initialized" when loading a model?

You added a new task head that was not in the checkpoint. It starts random and is learned during fine-tuning, so the warning is normal.

What replaced evaluation_strategy in TrainingArguments?

It is now called eval_strategy. Similarly, pass the tokenizer as processing_class instead of tokenizer.