Fine-tune un petit modèle avec Hugging Face Trainer

1. Aperçu général

Cet article vous guide pour fine-tuner un petit modèle (DistilBERT) afin de classifier les sentiments des reviews de produits, en utilisant l’API Trainer de la bibliothèque transformers. À la fin, vous saurez comment préparer un dataset, tokenizer, configurer TrainingArguments, entraîner, évaluer et sauvegarder un modèle pour l’inférence.

Exemple tout au long de l’article : OrderReview – review d’un client (Customer) concernant une commande (Order), étiquetée positive/negative.

2. Préparer l’environnement

Installez les bibliothèques nécessaires. Vous n’avez pas besoin d’un GPU puissant, un petit modèle comme DistilBERT fonctionne bien sur CPU avec un petit dataset, bien que plus lentement.

pip install transformers datasets evaluate accelerate torch scikit-learn

Note : Si votre machine a un GPU, installez la version de torch correspondant à votre CUDA pour un entraînement beaucoup plus rapide. Avec un CPU, limitez le nombre de lignes de données et d’epochs pour que la démo ne prenne pas trop de temps.

3. Préparer le dataset

Au lieu de charger un dataset énorme, on crée un petit ensemble imitant des reviews de commandes pour une meilleure compréhension. En pratique, vous pouvez le remplacer par load_dataset("imdb") ou vos propres données CSV.

from datasets import Dataset

reviews = [
    {"text": "Commande livrée rapidement, emballage soigné, très satisfait", "label": 1},
    {"text": "Produit défectueux, livraison lente, ne correspond pas à la description", "label": 0},
    {"text": "Bonne qualité, prix raisonnable, j'en rachèterai", "label": 1},
    {"text": "Service mauvais, support client lent", "label": 0},
    {"text": "Très content de cette commande, 5 étoiles", "label": 1},
    {"text": "Colis endommagé à la réception, déçu", "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 convertit une liste de dicts en Dataset de Hugging Face. train_test_split divise en ensembles train/test — avec un vrai dataset, vous devriez avoir plusieurs milliers de lignes, c’est juste un exemple pour illustrer le flux de traitement.

4. Tokenizer les données

Le modèle ne comprend pas directement le texte, il faut un tokenizer pour le convertir en input_ids et 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 : tronque les phrases plus longues que max_length.
• padding="max_length" : ajoute des tokens de padding pour que toutes les phrases aient la même longueur, facilitant le traitement par batch parallèle.
• map(..., batched=True) : applique la fonction tokenize par batch, plus rapide que de traiter ligne par ligne.

Note : DistilBERT en anglais n’est pas adapté pour fine-tuner des reviews en vietnamien en pratique. Ici on l’utilise pour démontrer la technique ; avec des données en vietnamien, choisissez un checkpoint comme vinai/phobert-base.

5. Charger le modèle et définir les 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 crée un nouveau classification head avec 2 sorties (positive/negative) au-dessus de la partie encoder pré-entraînée. id2label/label2id permettent à la sortie du modèle de retourner les labels lisibles par les humains au lieu de simples 0/1.

6. Configurer TrainingArguments

TrainingArguments regroupe tous les hyperparamètres et la configuration de la boucle d’entraînement.

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" : évalue après chaque epoch au lieu de chaque step.
• load_best_model_at_end=True : une fois l’entraînement terminé, Trainer recharge automatiquement le meilleur checkpoint selon metric_for_best_model.
• per_device_train_batch_size=4 : petit batch car le dataset de démo est petit aussi ; avec de vraies données, augmentez à 16/32 selon votre VRAM.

7. Définir la fonction de calcul des métriques

Trainer a besoin d’une fonction compute_metrics pour savoir comment évaluer le modèle selon quels critères.

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 est un tuple (logits, labels) que Trainer passe automatiquement après chaque évaluation. np.argmax sélectionne le label avec la plus haute probabilité (logit) comme prédiction.

8. Initialiser Trainer et entraîner

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 gère le padding dynamique par batch (plus efficace que le padding fixe en max_length lors du tokenization, mais ici on l’a déjà fait donc le collator s’occupe surtout de créer les tensors de batch). trainer.train() exécute la boucle forward → backward → optimizer step → log loss, et évalue automatiquement selon l’eval_strategy configurée.

8.1. Suivre la progression de l’entraînement

Si vous voulez enregistrer plus de détails dans TensorBoard, ajoutez report_to="tensorboard" dans TrainingArguments puis exécutez tensorboard --logdir ./order-review-model/runs. Avec un dataset petit comme la démo, les logs console suffisent pour vérifier que la loss diminue régulièrement.

9. Évaluer et sauvegarder le modèle

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() réexécute sur eval_dataset et retourne un dict de métriques. save_model sauvegarde config et poids ; n’oubliez pas de save_pretrained le tokenizer aussi pour éviter un décalage de vocab au rechargement.

10. Inférence avec le modèle fine-tuné

Utilisez pipeline pour tester rapidement le modèle qu’on vient d’entraîner.

from transformers import pipeline

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

result = classifier("Livraison à l'heure, produit plus beau que prévu")
print(result)

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

pipeline s’occupe automatiquement du tokenization, du forward pass et du softmax pour retourner le label avec le score (probabilité de confiance). C’est le moyen le plus rapide d’intégrer le modèle dans un petit service, par exemple un PaymentService qui a besoin de classifier automatiquement les retours client.

11. Écrire des tests pour vérifier le comportement du modèle

Avec le dataset de démo, on peut écrire des tests simples avec pytest pour s’assurer que le pipeline fonctionne correctement et ne régresse pas en changeant de 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("Produit merveilleux, livraison rapide")[0]

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

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("Produit défectueux, livraison tardive, très déçu")[0]

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

scope="module" fait que le pipeline charge le modèle une seule fois pour tous les tests du fichier, évitant de recharger à chaque test case. Ce test ne vérifie que le haut niveau du comportement (label correct), et ne doit pas être utilisé pour évaluer la qualité du modèle — c’est le rôle de l’étape evaluate ci-dessus.

12. Erreurs courantes

• Oublier padding/truncation : les phrases plus longues que max_length provoquent une erreur runtime car les tensors n’ont pas la même taille dans le batch.
• Confondre num_labels : si le dataset a 3 classes mais qu’on déclare num_labels=2, Trainer fonctionnera quand même mais les résultats seront totalement faux car le classification head n’a pas assez de sorties.
• Dataset trop petit puis attendre une F1 élevée sur les vraies données : l’exemple dans cet article ne sert qu’à illustrer le flux de code, 4 lignes d’entraînement ne suffisent pas pour que le modèle généralise. Ne prenez pas accuracy=1.0 ici comme référence.
• Ne pas définir seed : les résultats d’entraînement varieront à chaque exécution, ce qui complique la comparaison quand vous tunez les hyperparamètres. Utilisez set_seed(42) de transformers avant l’entraînement.
• Utiliser un gros modèle (base/large) sur une machine faible : commencez par les versions distilbert/tiny pour accélérer la boucle de développement, ne passez aux gros modèles qu’après avoir validé le pipeline.

Quand NE PAS utiliser Trainer : si vous avez besoin de personnaliser complètement la boucle d’entraînement (par exemple une loss multi-tâche complexe, un scheduler custom difficile à exprimer par callback), écrire une boucle PyTorch pure ou utiliser accelerate directement sera plus flexible que d’essayer de le forcer dans Trainer.

13. Conclusion

L’API Trainer vous permet de fine-tuner un petit modèle sans avoir à coder une boucle train/eval, il suffit de préparer le dataset, le tokenizer, le modèle et les TrainingArguments. Pour un problème de classification simple comme les reviews de commandes, le flux standard est : tokenize → configurer TrainingArguments → définir compute_metrics → créer Trainer → train() → evaluate() → save_model(). Quand vos données deviennent plus complexes, vous n’avez qu’à changer le dataset et le checkpoint, le reste du pipeline reste presque inchangé.