La classification de tokens est une tâche de traitement du langage naturel (NLP) qui consiste à attribuer une étiquette spécifique à chaque unité textuelle (token) d'une phrase. L'application la plus emblématique est la reconnaissance d'entités nommées (NER - Named Entity Recognition), dont le but est d'identifier des catégories telles que des noms de personnes, des lieux géographiques ou des organisations au sein d'un texte.

Ce guide détaille le processus de réglage fin (fine-tuning) du modèle DistilBERT sur le jeu de données WNUT 17 pour la détection d'entités spécifiques, ainsi que l'utilisation du modèle final pour l'inférence.

Configuration de l'environnement

Avant de commencer, installez les bibliothèques nécessaires à la manipulation des modèles et des métriques :

pip install transformers datasets evaluate seqeval

Chargement et exploration des données

Nous utilisons le jeu de données WNUT 17 via la bibliothèque 🤗 Datasets pour entraîner notre modèle sur des entités parfois complexes ou rares.

from datasets import load_dataset

dataset_wnut = load_dataset("wnut_17")

# Visualisation d'un exemple
print(dataset_wnut["train"][0])

Chaque exemple contient une liste de tokens et une liste correspondante de ner_tags. Les étiquettes numériques correspondent à des catégories précises :

noms_labels = dataset_wnut["train"].features["ner_tags"].feature.names
print(noms_labels)

Le format suit généralement le schéma BIO :

• B- (Begin) indique le début d'une entité.
• I- (Inside) indique que le token appartient à la même entité que le précédent.
• O (Outside) signifie que le token ne correspond à aucune entité.

Prétraitement et alignement des étiquettes

L'utilisation de modèles comme DistilBERT implique une tokenisation par sous-mots (subwords). Cela peut fragmenter un mot original en plusieurs tokens, créant un décalage avec les étiquettes initiales. Nous devons donc aligner les labels en suivant ces règles :

1. Mapper chaque token à son mot d'origine via les word_ids.
2. Attribuer l'indice -100 aux tokens spéciaux (comme [CLS] et [SEP]) pour qu'ils soient ignorés par la fnoction de perte.
3. Étiqueter uniquement le premier sous-mot d'un mot fragmenté et ignorer les suivants (valeur -100).

from transformers import AutoTokenizer

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

def aligner_labels_et_tokens(exemples):
    entrees_tokenisees = bert_tokenizer(exemples["tokens"], truncation=True, is_split_into_words=True)

    nouvelles_etiquettes = []
    for i, etiquette in enumerate(exemples["ner_tags"]):
        identifiants_mots = entrees_tokenisees.word_ids(batch_index=i)
        prec_id_mot = None
        liste_ids_labels = []
        
        for id_mot in identifiants_mots:
            if id_mot is None:
                liste_ids_labels.append(-100)
            elif id_mot != prec_id_mot:
                liste_ids_labels.append(etiquette[id_mot])
            else:
                liste_ids_labels.append(-100)
            prec_id_mot = id_mot
        nouvelles_etiquettes.append(liste_ids_labels)

    entrees_tokenisees["labels"] = nouvelles_etiquettes
    return entrees_tokenisees

donnees_traitees = dataset_wnut.map(aligner_labels_et_tokens, batched=True)

Pour finaliser la préparation, nous utilisons un DataCollatorForTokenClassification qui gère dynamiquement le rembourrage (padding) au sein des lots (batches).

from transformers import DataCollatorForTokenClassification

collecteur_donnees = DataCollatorForTokenClassification(tokenizer=bert_tokenizer)

Évaluation des performances

Pour mesurer l'efficacité du modèle, nous utilisons la métrique seqeval, standard pour les tâches de NER, qui calcule la précision, le rappel et le score F1.

import evaluate
import numpy as np

metrique_seqeval = evaluate.load("seqeval")

def calculer_metriques(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=2)

    pred_reelles = [
        [noms_labels[p] for (p, l) in zip(pred, label) if l != -100]
        for pred, label in zip(predictions, labels)
    ]
    labels_reels = [
        [noms_labels[l] for (p, l) in zip(pred, label) if l != -100]
        for pred, label in zip(predictions, labels)
    ]

    resultats = metrique_seqeval.compute(predictions=pred_reelles, references=labels_reels)
    return {
        "precision": resultats["overall_precision"],
        "recall": resultats["overall_recall"],
        "f1": resultats["overall_f1"],
        "accuracy": resultats["overall_accuracy"],
    }

Entraînement du modèle

Avant de lancer l'entraînement, définissez les correspondances entre idetnifiants et noms d'étiquettes :

id2label = {i: label for i, label in enumerate(noms_labels)}
label2id = {label: i for i, label in enumerate(noms_labels)}

from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer

modele_ner = AutoModelForTokenClassification.from_pretrained(
    "distilbert/distilbert-base-uncased", 
    num_labels=len(noms_labels), 
    id2label=id2label, 
    label2id=label2id
)

config_entrainement = TrainingArguments(
    output_dir="mon_modele_ner_wnut",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

entraineur = Trainer(
    model=modele_ner,
    args=config_entrainement,
    train_dataset=donnees_traitees["train"],
    eval_dataset=donnees_traitees["test"],
    tokenizer=bert_tokenizer,
    data_collator=collecteur_donnees,
    compute_metrics=calculer_metriques,
)

entraineur.train()

Inférence

Une fois le modèle entraîné, la méthode la plus simple pour l'utiliser consiste à faire appel à une pipeline de la bibliothèque Transformers.

from transformers import pipeline

phrase_test = "Le siège de Google se trouve à Mountain View, en Californie."
ner_pipeline = pipeline("ner", model="chemin_vers_votre_modele")
resultats = ner_pipeline(phrase_test)

for entite in resultats:
    print(entite)

Vous pouvez également effectuer l'inférence manuellement en utilisant PyTorch pour obtenir un contrôle total sur les logits et la conversion des IDs en labels.

import torch

entrees = bert_tokenizer(phrase_test, return_tensors="pt")
with torch.no_grad():
    outputs = modele_ner(**entrees).logits

predictions = torch.argmax(outputs, dim=2)
labels_predis = [modele_ner.config.id2label[t.item()] for t in predictions[0]]
print(list(zip(bert_tokenizer.convert_ids_to_tokens(entrees["input_ids"][0]), labels_predis)))