Fine-tune model nhỏ với Hugging Face Trainer

1. Tổng quan

Bài này hướng dẫn bạn fine-tune một model nhỏ (DistilBERT) để phân loại cảm xúc review sản phẩm, dùng Trainer API của thư viện transformers. Đọc xong, bạn sẽ biết cách chuẩn bị dataset, tokenize, cấu hình TrainingArguments, huấn luyện, đánh giá và lưu model để dùng cho inference.

Ví dụ xuyên suốt bài: OrderReview – review của khách hàng (Customer) về đơn hàng (Order), gắn nhãn positive/negative.

2. Chuẩn bị môi trường

Cài các thư viện cần thiết. Bạn không cần GPU mạnh, model nhỏ như DistilBERT chạy ổn trên CPU với tập dữ liệu nhỏ, tuy chậm hơn.

pip install transformers datasets evaluate accelerate torch scikit-learn

Lưu ý: Nếu máy có GPU, cài thêm bản torch tương ứng CUDA để train nhanh hơn đáng kể. Với CPU, nên giới hạn số dòng dữ liệu và epoch để demo không mất quá nhiều thời gian.

3. Chuẩn bị dataset

Thay vì tải dataset khổng lồ, ta tạo một tập nhỏ mô phỏng review đơn hàng để dễ hình dung. Trong thực tế bạn có thể thay bằng load_dataset("imdb") hoặc dữ liệu CSV của riêng mình.

from datasets import Dataset

reviews = [
    {"text": "Đơn hàng giao nhanh, đóng gói cẩn thận, rất hài lòng", "label": 1},
    {"text": "Sản phẩm lỗi, giao chậm, không giống mô tả", "label": 0},
    {"text": "Chất lượng tốt, giá hợp lý, sẽ mua lại", "label": 1},
    {"text": "Dịch vụ tệ, hỗ trợ khách hàng chậm trễ", "label": 0},
    {"text": "Rất ưng ý với đơn hàng này, 5 sao", "label": 1},
    {"text": "Hàng bị móp méo khi nhận, thất vọng", "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 biến list dict thành Dataset của Hugging Face. train_test_split chia thành tập train/test — với dataset thật bạn nên có vài nghìn dòng trở lên, đây chỉ là ví dụ minh họa luồng xử lý.

4. Tokenize dữ liệu

Model không hiểu text trực tiếp, cần tokenizer để chuyển thành input_ids và 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: cắt câu dài hơn max_length.
• padding="max_length": đệm thêm token để mọi câu cùng độ dài, giúp batch xử lý song song dễ dàng.
• map(..., batched=True): áp dụng hàm tokenize theo batch, nhanh hơn xử lý từng dòng.

Lưu ý: DistilBERT tiếng Anh không phù hợp để fine-tune review tiếng Việt trong thực tế. Ở đây dùng để demo kỹ thuật; với dữ liệu tiếng Việt, hãy chọn checkpoint như vinai/phobert-base.

5. Load model và định nghĩa nhãn

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 tạo một classification head mới với 2 output (positive/negative) gắn lên trên phần encoder đã pretrain. id2label/label2id giúp output của model trả về nhãn người đọc được thay vì chỉ số 0/1.

6. Cấu hình TrainingArguments

TrainingArguments gom toàn bộ hyperparameter và cấu hình vòng lặp huấn luyện.

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": đánh giá sau mỗi epoch thay vì mỗi step.
• load_best_model_at_end=True: sau khi train xong, Trainer tự nạp lại checkpoint tốt nhất theo metric_for_best_model.
• per_device_train_batch_size=4: batch nhỏ vì dataset demo cũng nhỏ; với dữ liệu thật, tăng lên 16/32 tùy VRAM.

7. Định nghĩa hàm tính metric

Trainer cần một hàm compute_metrics để biết đánh giá model dựa trên tiêu chí nào.

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 là tuple (logits, labels) do Trainer tự truyền vào sau mỗi lần eval. np.argmax chọn nhãn có xác suất (logit) cao nhất làm prediction.

8. Khởi tạo Trainer và huấn luyện

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 xử lý padding động theo batch (hiệu quả hơn pad cố định max_length ở bước tokenize, nhưng ở đây đã pad sẵn nên collator chủ yếu lo việc tạo batch tensor). trainer.train() chạy vòng lặp forward → backward → optimizer step → log loss, và tự động eval theo eval_strategy đã cấu hình.

8.1. Theo dõi quá trình train

Nếu muốn log chi tiết hơn ra TensorBoard, thêm report_to="tensorboard" vào TrainingArguments rồi chạy tensorboard --logdir ./order-review-model/runs. Với dataset nhỏ như demo, log console là đủ để kiểm tra loss có giảm dần hay không.

9. Đánh giá và lưu 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() chạy lại trên eval_dataset và trả về dict metric. save_model lưu cả config, weights; nhớ save_pretrained tokenizer đi kèm để khi load lại không bị lệch vocab.

10. Inference với model đã fine-tune

Dùng pipeline để test nhanh model vừa huấn luyện.

from transformers import pipeline

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

result = classifier("Giao hàng đúng hẹn, sản phẩm đẹp hơn mong đợi")
print(result)

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

pipeline tự lo tokenize, forward, và softmax để ra label kèm score (xác suất tin cậy). Đây là cách nhanh nhất để tích hợp model vào một service nhỏ, ví dụ một PaymentService cần phân loại phản hồi khách hàng tự động.

11. Viết test kiểm tra hành vi model

Với dataset demo, ta có thể viết test đơn giản bằng pytest để đảm bảo pipeline hoạt động đúng và không bị regression khi đổi checkpoint.

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("Sản phẩm tuyệt vời, giao hàng nhanh")[0]

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

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("Hàng lỗi, giao trễ, rất thất vọng")[0]

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

scope="module" giúp pipeline chỉ load model một lần cho tất cả test trong file, tránh tốn thời gian load lại mỗi test case. Test này chỉ kiểm tra hành vi mức cao (label đúng chiều), không nên dùng để đánh giá chất lượng model — việc đó thuộc về bước evaluate ở trên.

12. Lỗi thường gặp

• Quên padding/truncation: câu dài hơn max_length sẽ gây lỗi runtime vì tensor không đồng kích thước trong batch.
• Nhầm num_labels: nếu dataset có 3 lớp nhưng khai báo num_labels=2, Trainer vẫn chạy nhưng kết quả sai lệch hoàn toàn vì classification head không đủ output.
• Dataset quá nhỏ rồi kỳ vọng F1 cao trên dữ liệu thật: ví dụ trong bài chỉ để minh họa luồng code, 4 dòng train không đủ để model học tổng quát. Đừng lấy accuracy=1.0 ở đây làm tiêu chuẩn.
• Không set seed: kết quả train sẽ khác nhau giữa các lần chạy, khó so sánh khi tune hyperparameter. Dùng set_seed(42) từ transformers trước khi train.
• Dùng model lớn (base/large) cho máy yếu: nên bắt đầu với bản distilbert/tiny để vòng lặp phát triển nhanh, chỉ chuyển sang model lớn khi đã chốt pipeline.

Khi nào KHÔNG nên dùng Trainer: nếu bạn cần custom hoàn toàn vòng lặp training (ví dụ multi-task loss phức tạp, custom scheduler khó biểu diễn qua callback), viết loop PyTorch thuần hoặc dùng accelerate trực tiếp sẽ linh hoạt hơn là cố nhét vào Trainer.

13. Kết luận

Trainer API giúp bạn fine-tune model nhỏ mà không cần tự viết vòng lặp train/eval, chỉ cần chuẩn bị dataset, tokenizer, model và TrainingArguments. Với bài toán phân loại đơn giản như review đơn hàng, luồng chuẩn là: tokenize → cấu hình TrainingArguments → định nghĩa compute_metrics → khởi tạo Trainer → train() → evaluate() → save_model(). Khi dữ liệu thật phức tạp hơn, bạn chỉ cần thay dataset và checkpoint, phần còn lại của pipeline gần như giữ nguyên.