Fine-tune Small Models with Hugging Face Trainer

1. Overview

This guide walks you through fine-tuning a small model (DistilBERT) to classify the sentiment of product reviews using the Trainer API from the transformers library. By the end, you’ll know how to prepare datasets, tokenize, configure TrainingArguments, train, evaluate, and save a model for inference.

Throughout this article, we’ll use the example of OrderReview — a customer’s (Customer) review of an order (Order), labeled as positive or negative.

2. Setting up the environment

Install the necessary libraries. You don’t need a powerful GPU; a small model like DistilBERT runs fine on CPU with a small dataset, though it will be slower.

pip install transformers datasets evaluate accelerate torch scikit-learn

Note: If your machine has a GPU, install the corresponding CUDA version of PyTorch for significantly faster training. With CPU, you should limit the number of data rows and epochs to keep the demo from taking too long.

3. Preparing the dataset

Instead of downloading a huge dataset, we’ll create a small one that mimics order reviews to keep things simple. In practice, you could replace it with load_dataset("imdb") or your own CSV data.

from datasets import Dataset

reviews = [
    {"text": "Order delivered quickly, careful packaging, very satisfied", "label": 1},
    {"text": "Defective product, slow delivery, doesn't match description", "label": 0},
    {"text": "Good quality, reasonable price, will buy again", "label": 1},
    {"text": "Poor service, slow customer support", "label": 0},
    {"text": "Very happy with this order, 5 stars", "label": 1},
    {"text": "Item was dented when received, disappointed", "label": 0},
]

dataset = Dataset.from_list(reviews)
dataset = dataset.train_test_split(test_size=0.33, seed=42)
print(dataset)

DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 4})
    test: Dataset({features: ['text', 'label'], num_rows: 2})
})

Dataset.from_list converts a list of dicts into a Hugging Face Dataset. train_test_split splits it into train/test sets — with real data you should have several thousand rows or more; this is just an example to illustrate the workflow.

4. Tokenizing the data

Models don’t understand text directly; you need a tokenizer to convert it to input_ids and attention_mask.

from transformers import AutoTokenizer

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

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=32)

tokenized_dataset = dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0].keys())

dict_keys(['text', 'label', 'input_ids', 'attention_mask'])

• truncation=True: truncate sequences longer than max_length.
• padding="max_length": pad sequences to the same length for easy batch processing.
• map(..., batched=True): apply the tokenize function by batch, which is faster than row-by-row processing.

Note: English DistilBERT is not suitable for fine-tuning Vietnamese reviews in practice. We use it here for technical demonstration; for Vietnamese data, choose a checkpoint like vinai/phobert-base.

5. Loading the model and defining labels

from transformers import AutoModelForSequenceClassification

label2id = {"negative": 0, "positive": 1}
id2label = {v: k for k, v in label2id.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)

num_labels=2 creates a new classification head with 2 outputs (positive/negative) on top of the pretrained encoder. id2label/label2id helps the model output return human-readable labels instead of just 0/1.

6. Configuring TrainingArguments

TrainingArguments bundles all hyperparameters and training loop configuration.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./order-review-model",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=1,
)

• eval_strategy="epoch": evaluate after each epoch instead of every step.
• load_best_model_at_end=True: after training finishes, Trainer automatically loads the best checkpoint according to metric_for_best_model.
• per_device_train_batch_size=4: small batch size because the demo dataset is small; with real data, increase to 16/32 depending on your VRAM.

7. Defining the metrics function

Trainer needs a compute_metrics function so it knows which metrics to use for evaluation.

import numpy as np
import evaluate

accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels)
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

eval_pred is a tuple (logits, labels) that Trainer passes in during each evaluation. np.argmax selects the label with the highest probability (logit) as the prediction.

8. Initializing Trainer and training

from transformers import Trainer, DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

Epoch   Training Loss   Validation Loss   Accuracy   F1
1       0.693100        0.671204          0.500000   0.500000
2       0.654200        0.598311          1.000000   1.000000
3       0.581900        0.520447          1.000000   1.000000

DataCollatorWithPadding handles dynamic padding per batch (more efficient than fixed max_length padding during tokenization, but since we already padded, it mainly handles batch tensor creation). trainer.train() runs the loop of forward → backward → optimizer step → log loss, and automatically evaluates according to the configured eval_strategy.

8.1. Monitoring training progress

If you want more detailed logging to TensorBoard, add report_to="tensorboard" to TrainingArguments and run tensorboard --logdir ./order-review-model/runs. With a small dataset like this demo, console logs are enough to check if loss is decreasing.

9. Evaluating and saving the model

metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./order-review-model/final")
tokenizer.save_pretrained("./order-review-model/final")

{'eval_loss': 0.5204, 'eval_accuracy': 1.0, 'eval_f1': 1.0, 'eval_runtime': 0.12, ...}

trainer.evaluate() runs again on eval_dataset and returns a metric dict. save_model saves both config and weights; remember to save_pretrained the tokenizer as well so vocab doesn’t shift when loading later.

10. Inference with the fine-tuned model

Use pipeline to quickly test the model you just trained.

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="./order-review-model/final",
    tokenizer="./order-review-model/final",
)

result = classifier("On-time delivery, product looks better than expected")
print(result)

[{'label': 'positive', 'score': 0.9123}]

pipeline automatically handles tokenization, forward pass, and softmax to produce a label with a score (confidence probability). This is the quickest way to integrate the model into a small service, for example a PaymentService that needs to automatically classify customer feedback.

11. Writing tests to verify model behavior

With the demo dataset, we can write simple tests with pytest to ensure the pipeline works correctly and doesn’t regress when changing checkpoints.

import pytest
from transformers import pipeline

@pytest.fixture(scope="module")
def classifier():
    return pipeline(
        "text-classification",
        model="./order-review-model/final",
        tokenizer="./order-review-model/final",
    )

def test_positive_review_is_classified_correctly(classifier):
    result = classifier("Great product, fast delivery")[0]

    assert result["label"] == "positive"
    assert result["score"] > 0.5

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("Defective product, late delivery, very disappointed")[0]

    assert result["label"] == "negative"
    assert result["score"] > 0.5

scope="module" ensures the pipeline loads the model only once for all tests in the file, avoiding expensive reloads for each test. This test only checks high-level behavior (correct label direction) and shouldn’t be used to assess model quality — that’s what the evaluation step above is for.

12. Common mistakes

• Forgot padding/truncation: sequences longer than max_length will cause runtime errors because tensors will have mismatched sizes in a batch.
• Wrong num_labels: if your dataset has 3 classes but you set num_labels=2, Trainer will still run but produce completely wrong results because the classification head doesn’t have enough outputs.
• Tiny dataset and high F1 expectations on real data: the example in this article is only for code flow illustration; 4 training rows are not enough for the model to generalize. Don’t use accuracy=1.0 here as a benchmark.
• No seed set: training results will differ between runs, making it hard to compare when tuning hyperparameters. Use set_seed(42) from transformers before training.
• Using large (base/large) models on weak hardware: start with distilbert or tiny versions so development cycles are fast; only move to larger models once you’ve finalized your pipeline.

When NOT to use Trainer: if you need to fully customize the training loop (e.g., complex multi-task loss, custom schedulers hard to express via callbacks), writing a pure PyTorch loop or using accelerate directly will be more flexible than trying to fit everything into Trainer.

13. Conclusion

The Trainer API lets you fine-tune small models without writing train/eval loops yourself; you just prepare your dataset, tokenizer, model, and TrainingArguments. For simple classification tasks like order reviews, the standard flow is: tokenize → configure TrainingArguments → define compute_metrics → initialize Trainer → train() → evaluate() → save_model(). When your real data gets more complex, you only need to swap the dataset and checkpoint; the rest of the pipeline stays mostly the same.