Entraîner un modèle de langage causal à partir de zéro

Jusqu’à présent, nous avons surtout réutilisé des modèles pré-entraînés et les avons finetunés sur de nouveaux cas d’usage. Comme nous l’avons vu dans le chapitre 1, ceci est communément appelé apprentissage par transfert, et il s’agit d’une stratégie très efficace pour appliquer les transformers à la plupart des applications du monde réel où les données étiquetées sont rares. Dans ce chapitre, nous allons adopter une approche différente consistant à entraîner un modèle complètement nouveau à partir de zéro. C’est une bonne démarche à adopter si vous avez beaucoup de données et qu’elles sont très différentes des données de pré-entraînement utilisées par les modèles disponibles. Cependant, le pré-entraînement d’un modèle de langue nécessite beaucoup plus de ressources informatiques que le simple finetuning d’un modèle existant. Parmi les exemples où il peut être utile d’entraîner un nouveau modèle, citons les jeux de données constitués de notes de musique, de séquences moléculaires telles que l’ADN, ou de langages de programmation. Ces derniers ont récemment gagné en popularité grâce à des outils tels que TabNine et Copilot de GitHub (alimentés par le modèle Codex d’OpenAI) qui peuvent générer de longues séquences de code. Cette tâche de génération de texte est mieux abordée avec des modèles de langage autorégressifs ou causaux tels que le GPT-2.

Dans cette section, nous allons construire une version réduite d’un modèle de génération de code Python. Nous nous concentrerons sur la complétion d’une ligne de code au lieu de fonctions ou de classes complètes. Lorsque vous travaillez sur des projets de science des données en Python, vous êtes souvent en contact avec les bibliothèques matplotlib, seaborn, pandas et scikit-learn. Lors de l’utilisation de ces frameworks, il est fréquent d’avoir besoin de rechercher des commandes spécifiques. Il serait donc bien d’utiliser un modèle pour compléter ces appels pour nous.

Dans le chapitre 6, nous avons créé un tokenizer efficace pour traiter du code Python. Nous avons besoin d’un jeu de données à grande échelle pour pré-entraîner un modèle. Ici, nous allons appliquer notre tokenizer à un corpus de code Python provenant des dépôts GitHub. Nous utiliserons ensuite l’API Trainer et 🤗 Accelerate pour entraîner le modèle. C’est parti !

Il s’agit d’une présentation du modèle qui a été entraîné à l’aide du code présenté dans cette section et qui a ensuité été téléchargé sur le Hub. Vous pouvez le trouver ici. Notez qu’étant donné qu’il y a un certains aléat dans la génération du texte, vous obtiendrez probablement un résultat légèrement différent.

Collecte des données

On peut trouver du code Python en abondance dans les dépôts de code tels que GitHub, que nous pouvons utiliser pour créer un jeu de données en récupérant chaque dépôt Python. C’est l’approche adoptée dans le livre Natural Language Processing with Transformers pour pré-entraîner un grand GPT-2. En utilisant un dépôt GitHub d’environ 180 Go contenant approximativement 20 millions de fichiers Python, les auteurs du livre ont construit un jeu de données appelé codeparrot qu’ils ont ensuite partagé sur le Hub.

Cependant, entraîner sur l’ensemble du corpus prend beaucoup de temps et demande beaucoup de ressources de calculs. Dans notre cas, nous n’avons besoin que du sous-ensemble du jeu de données qui est relatif aux codes portant sur la science des données. Commençons donc par filtrer le jeu de données codeparrot en ne gardant que les fichiers incluant l’une des bibliothèques de science des données énumérées précédemment. En raison de la taille du jeu de données, nous voulons éviter de le télécharger. Nous utiliserons donc la fonctionnalité de streaming de 🤗 Datasets afin de le filtrer à la volée. Pour nous aider à filtrer les échantillons de code utilisant les bibliothèques que nous avons mentionnées précédemment, nous utilisons la fonction suivante :

def any_keyword_in_string(string, keywords):
    for keyword in keywords:
        if keyword in string:
            return True
    return False

Testons-le sur deux exemples :

filters = ["pandas", "sklearn", "matplotlib", "seaborn"]
example_1 = "import numpy as np"
example_2 = "import pandas as pd"

print(
    any_keyword_in_string(example_1, filters), any_keyword_in_string(example_2, filters)
)

False True

Nous pouvons l’utiliser pour créer une fonction qui va streamer le jeu de donner et filtrer les éléments que nous voulons :

def filter_streaming_dataset(dataset, filters):
    filtered_dict = defaultdict(list)
    total = 0
    for sample in tqdm(iter(dataset)):
        total += 1
        if any_keyword_in_string(sample["content"], filters):
            for k, v in sample.items():
                filtered_dict[k].append(v)
    print(f"{len(filtered_dict['content'])/total:.2%} of data after filtering.")
    return Dataset.from_dict(filtered_dict)

Ensuite, nous pouvons simplement appliquer cette fonction :

# Cette cellule prendra beaucoup de temps à s'exécuter, donc vous devriez la sauter et aller à la suivante !
from datasets import load_dataset

split = "train"  # "valid"
filters = ["pandas", "sklearn", "matplotlib", "seaborn"]

data = load_dataset(f"transformersbook/codeparrot-{split}", split=split, streaming=True)
filtered_data = filter_streaming_dataset(data, filters)

3.26% of data after filtering.

Cela nous laisse avec environ 3 % du jeu de données original, ce qui est tout de même assez important puisqu’il fait 6 Go et se compose de 600 000 scripts Python !

Le filtrage peut prendre de 2 à 3 heures, selon votre machine et votre bande passante. Si vous ne voulez pas passer par ce long processus, nous fournissons sur le Hub le jeu de données filtré pour que vous puissiez le télécharger :

from datasets import load_dataset, DatasetDict

ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="train")

raw_datasets = DatasetDict(
    {
        "train": ds_train,  # .shuffle().select(range(50000)),
        "valid": ds_valid,  # .shuffle().select(range(500))
    }
)

raw_datasets

DatasetDict({
    train: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 606720
    })
    valid: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 3322
    })
})

Examinons un exemple tiré du jeu de données. Nous ne montrerons que les 200 premiers caractères de chaque champ :

for key in raw_datasets["train"][0]:
    print(f"{key.upper()}: {raw_datasets['train'][0][key][:200]}")

'REPO_NAME: kmike/scikit-learn'
'PATH: sklearn/utils/__init__.py'
'COPIES: 3'
'SIZE: 10094'
'''CONTENT: """
The :mod:`sklearn.utils` module includes various utilites.
"""

from collections import Sequence

import numpy as np
from scipy.sparse import issparse
import warnings

from .murmurhash import murm
LICENSE: bsd-3-clause'''

Nous pouvons voir que le champ content contient le code sur lequel nous voulons que notre modèle s’entraîne. Maintenant que nous avons un jeu de données, nous devons préparer les textes afin qu’ils soient dans un format approprié pour le pré-entraînement.

Préparation du jeu de données

La première étape est de tokeniser les données afin de pouvoir les utiliser pour l’entraînement. Puisque notre objectif est d’autocompléter de courts appels de fonctions, nous pouvons garder la taille du contexte relativement petite. L’avantage est que nous pouvons entraîner le modèle beaucoup plus rapidement et qu’il nécessite beaucoup moins de mémoire. Si c’est important pour votre application d’avoir davantage de contexte (par exemple, si vous voulez que le modèle écrive des tests unitaires basés sur un fichier avec la définition de la fonction), assurez-vous d’augmenter ce nombre. Gardez néanmoins à l’esprit que cela s’accompagne d’une plus grande empreinte mémoire du GPU. Pour l’instant, fixons la taille du contexte à 128 tokens, par opposition aux 1 024 ou 2 048 utilisés respectivement dans le GPT-2 et le GPT-3.

La plupart des documents contiennent beaucoup plus de 128 tokens, donc le fait de tronquer les entrées à la longueur maximale éliminerait une grande partie de notre jeu de données. A la place, nous allons utiliser l’option return_overflowing_tokens pour tokeniser l’entrée entière et la diviser en plusieurs morceaux, comme nous l’avons fait dans le chapitre 6. Nous utiliserons également l’option return_length pour retourner automatiquement la longueur de chaque morceau créé. Souvent, le dernier morceau est plus petit que la taille du contexte et nous nous en débarrasserons pour éviter les problèmes de padding. Nous n’en avons pas vraiment besoin puisque de toute façon nous avons beaucoup de données.

Voyons comment cela fonctionne en examinant les deux premiers exemples :

from transformers import AutoTokenizer

context_length = 128
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")

outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True,
    max_length=context_length,
    return_overflowing_tokens=True,
    return_length=True,
)

print(f"Input IDs length: {len(outputs['input_ids'])}")
print(f"Input chunk lengths: {(outputs['length'])}")
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")

Input IDs length: 34
Input chunk lengths: [128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 117, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 41]
Chunk mapping: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

Nous pouvons voir que nous obtenons 34 morceaux à partir de ces deux exemples. En regardant leurs longueurs, nous pouvons voir qu’ils se terminent avec moins de 128 tokens (117 et 41, respectivement). Ils ne représentent qu’une petite fraction du total des morceaux que nous avons (2/34), donc nous pouvons les jeter sans risque. Avec le champ overflow_to_sample_mapping, nous pouvons aussi reconstruire quels morceaux appartenaient à quels échantillons d’entrée.

Avec cette opération, nous utilisons une fonctionnalité pratique de la fonction Dataset.map() de 🤗 Datasets. En effet, celle-ci ne nécessite pas une correspondance un à un comme nous l’avons vu dans la section 3. Nous pouvons créer des batchs avec plus ou moins d’éléments que le batch d’entrée. C’est utile lorsque l’on effectue des opérations telles que l’augmentation ou le filtrage des données qui modifient le nombre d’éléments. Dans notre cas, lors de la tokenisation de chaque élément en morceaux de longeur de la taille de contexte spécifiée, nous créons de nombreux échantillons de chaque document. Nous devons juste nous assurer de supprimer les colonnes existantes, car elles ont une taille conflictuelle. Si nous voulions les garder, nous pourrions les répéter de manière appropriée et les retourner dans l’appel Dataset.map() :

def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length:
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
tokenized_datasets

DatasetDict({
    train: Dataset({
        features: ['input_ids'],
        num_rows: 16702061
    })
    valid: Dataset({
        features: ['input_ids'],
        num_rows: 93164
    })
})

Nous avons maintenant 16,7 millions d’exemples avec 128 tokens chacun, ce qui correspond à environ 2,1 milliards de tokens au total. A titre de comparaison, les modèles GPT-3 et Codex d’OpenAI sont entraînés sur 300 et 100 milliards de tokens, respectivement. Les modèles Codex étant initialisés à partir des checkpoints GPT-3. Notre objectif dans cette section n’est pas de rivaliser avec ces modèles, qui peuvent générer des textes longs et cohérents, mais de créer une version réduite fournissant une fonction d’autocomplétion rapide.

Maintenant que le jeu de données est prêt, configurons le modèle !

Initialisation d’un nouveau modèle

Notre première étape consiste à initialiser un GPT-2. Pour notre modèle, nous utiliserons la même configuration que pour le petit modèle GPT-2. Ainsi nous chargeons la configuration pré-entraînée, nous nous assurons que la taille du tokenizer correspond à la taille du vocabulaire du modèle et nous passons les identifiants des tokens bos et eos (début et fin de séquence) :

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

Avec cette configuration, nous pouvons charger un nouveau modèle. Notez que c’est la première fois que nous n’utilisons pas la fonction from_pretrained() puisque nous initialisons nous-mêmes un modèle :

model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

GPT-2 size: 124.2M parameters

Notre modèle comporte 124 millions de paramètres que nous devrons régler. Avant de commencer l’entraînement, nous devons configurer un assembleur de données qui se chargera de créer les batchs. Nous pouvons utiliser le assembleur DataCollatorForLanguageModeling, qui est conçu spécifiquement pour la modélisation du langage (comme son nom le suggère subtilement). En plus de l’empilage et du rembourrage des batchs, il s’occupe aussi de la création des étiquettes du modèle de langage. Dans la modélisation causale du langage, les entrées servent aussi d’étiquettes (juste décalées d’un élément) et que le assembleur de données crée à la volée pendant l’entraînement pour ne pas avoir à dupliquer les input_ids.

Notez que DataCollatorForLanguageModeling supporte à la fois la modélisation du langage masqué (MLM pour masked language modeling) et la modélisation du langage causal (CLM pour causal language modeling). Par défaut, il prépare les données pour la MLM mais nous pouvons passer à la CLM en définissant l’argument mlm=False :

from transformers import DataCollatorForLanguageModeling

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

Prenons un exemple :

out = data_collator([tokenized_dataset["train"][i] for i in range(5)])
for key in out:
    print(f"{key} shape: {out[key].shape}")

input_ids shape: torch.Size([5, 128])
attention_mask shape: torch.Size([5, 128])
labels shape: torch.Size([5, 128])

Nous pouvons voir que les exemples ont été empilés et que tous les tenseurs ont la même forme.

Nous avons maintenant tout ce qu’il faut pour entraîner notre modèle. Ce n’était pas si compliqué ! Avant de commencer l’entraînement, nous devons nous connecter à Hugging Face. Si vous travaillez dans un notebook, vous pouvez le faire avec la fonction utilitaire suivante :

from huggingface_hub import notebook_login

notebook_login()

Cela affichera un widget où vous pourrez entrer vos identifiants de connexion à Hugging Face.

Si vous ne travaillez pas dans un notebook, tapez simplement la ligne suivante dans votre terminal :

huggingface-cli login

Tout ce qu’il reste à faire est de configurer les arguments d’entraînement et de lancer la fonction Trainer. Nous utiliserons un programme de taux d’apprentissage de type cosinus avec un réchauffement et une taille de batch de 256 (per_device_train_batch_size x gradient_accumulation_steps). L’accumulation du gradient est utilisée lorsqu’un seul batch ne tient pas en mémoire, et construit le gradient de manière incrémentale à travers plusieurs passages en avant/en arrière. Nous verrons cela en action lorsque nous créerons la boucle d’entraînement avec 🤗 Accelerate.

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir="codeparrot-ds",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="steps",
    eval_steps=5_000,
    logging_steps=5_000,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=1_000,
    lr_scheduler_type="cosine",
    learning_rate=5e-4,
    save_steps=5_000,
    fp16=True,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["valid"],
)

Maintenant, nous pouvons simplement lancer le Trainer et attendre que l’entraînement se termine. Selon que vous l’exécutez sur la totalité ou sur un sous-ensemble de l’échantillon d’entraînement, cela prendra respectivement 20 ou 2 heures. Alors prenez quelques cafés et un bon livre à lire !

trainer.train()

Une fois l’entraînement terminé, nous pouvons pousser le modèle et le tokenizer vers le Hub :

trainer.push_to_hub()

Génération de code avec le pipeline

C’est maintenant le moment de vérité : voyons comment le modèle entraîné fonctionne réellement ! Nous pouvons voir dans les logs que la perte a diminué régulièrement, mais pour mettre le modèle à l’épreuve, regardons comment il fonctionne sur certains messages. Pour ce faire, nous allons envelopper le modèle dans un pipeline de génération de texte et, s’il y en a un de disponible, utiliser un GPU pour avoir des générations rapidement :

import torch
from transformers import pipeline

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
pipe = pipeline(
    "text-generation", model="huggingface-course/codeparrot-ds", device=device
)

Let’s start with the simple task of creating a scatter plot:

txt = """\
# créer des données
x = np.random.randn(100)
y = np.random.randn(100)

# créer un nuage de points avec x, y
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# créer des données
x = np.random.randn(100)
y = np.random.randn(100)

# créer un nuage de points avec x, y
plt.scatter(x, y)

Le résultat semble correct. Est-ce que cela fonctionne aussi pour une opération pandas ? Voyons si nous pouvons créer un DataFrame à partir de deux tableaux :

txt = """\
# créer des données
x = np.random.randn(100)
y = np.random.randn(100)

# créer un tableau de données à partir de x et y
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# créer des données
x = np.random.randn(100)
y = np.random.randn(100)

# créer un tableau de données à partir de x et y
df = pd.DataFrame({'x': x, 'y': y})
df.insert(0,'x', x)
for

Bien, c’est la bonne réponse. Bien qu’il insère ensuite la colonne x à nouveau. Comme le nombre de tokens générés est limité, la boucle for suivante est coupée. Voyons si nous pouvons faire quelque chose d’un peu plus complexe et faire en sorte que le modèle nous aide à utiliser l’opération groupby :

txt = """\
# tableau de données avec profession, revenu et nom
df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

# calculer le revenu moyen par profession
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# tableau de données avec profession, revenu et nom
df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

# calculer le revenu moyen par profession
profession = df.groupby(['profession']).mean()

Pas mal, c’est la bonne façon de faire. Enfin, voyons si nous pouvons aussi l’utiliser pour scikit-learn et utiliser un modèle Random Forest :

txt = """
# import random forest regressor from scikit-learn
from sklearn.ensemble import RandomForestRegressor

# entraînement du modèle de forêt aléatoire avec 300 estimateurs sur X, y :
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# import random forest regressor from scikit-learn
from sklearn.ensemble import RandomForestRegressor

# entraînement du modèle de forêt aléatoire avec 300 estimateurs sur X, y :
rf = RandomForestRegressor(n_estimators=300, random_state=random_state, max_depth=3)
rf.fit(X, y)
rf

Au vu de ces quelques exemples, il semble que le modèle ait appris une partie de la syntaxe des bibliothèques Python de science des données. Bien sûr, nous devrions évaluer le modèle de manière plus approfondie avant de le déployer dans le monde réel, mais il s’agit tout de même d’un prototype impressionnant. Parfois, il est nécessaire de personnaliser davantage l’entraînement du modèle afin d’obtenir les performances nécessaires pour un cas d’utilisation donné. Par exemple, que se passe-t-il si l’on souhaite mettre à jour dynamiquement la taille du batch ou si l’on dispose d’une boucle d’entraînement conditionnelle qui ignore les mauvais exemples à la volée ? Une option serait de sous-classer le Trainer et d’ajouter les changements nécessaires, mais parfois il est plus simple d’écrire la boucle d’entraînement à partir de zéro. C’est là qu’intervient 🤗 Accelerate.

Entraîner avec 🤗 <i> Accelerate </i>

Nous avons vu comment entraîner un modèle avec le Trainer, qui permet une certaine personnalisation. Cependant, parfois nous voulons un contrôle total sur la boucle d’entraînement ou nous souhaitons faire quelques changements exotiques. Dans ce cas, 🤗 Accelerate est un excellent choix, et dans cette section, nous allons suivre les étapes pour l’utiliser pour entraîner notre modèle. Pour rendre les choses plus intéressantes, nous allons également ajouter une touche à la boucle d’entraînement.

Puisque nous sommes principalement intéressés par l’autocomplétion pour les bibliothèques de science des données, il est logique de donner plus de poids aux échantillons d’entraînement qui utilisent davantage ces bibliothèques. Nous pouvons facilement identifier ces exemples grâce à l’utilisation de mots-clés tels que plt, pd, sk, fit, et predict, qui sont les noms d’importation les plus fréquents pour matplotlib.pyplot, pandas, et sklearn ainsi que les fonctions fit et predict de cette dernière. Si chacun d’entre eux est représenté par un seul token, nous pouvons facilement vérifier s’ils apparaissent dans la séquence d’entrée. Les tokens peuvent avoir un préfixe d’espacement, donc nous vérifierons aussi ces versions dans le vocabulaire du tokenizer. Pour vérifier que cela fonctionne, nous ajouterons un token de test qui devrait être divisé en plusieurs tokens :

keytoken_ids = []
for keyword in [
    "plt",
    "pd",
    "sk",
    "fit",
    "predict",
    " plt",
    " pd",
    " sk",
    " fit",
    " predict",
    "testtest",
]:
    ids = tokenizer([keyword]).input_ids[0]
    if len(ids) == 1:
        keytoken_ids.append(ids[0])
    else:
        print(f"Keyword has not single token: {keyword}")

'Keyword has not single token: testtest'

Super, ça a l’air de bien fonctionner ! Nous pouvons maintenant écrire une fonction de perte personnalisée qui prend la séquence d’entrée, les logits et les tokens clés que nous venons de sélectionner comme entrées. Tout d’abord, nous devons aligner les logits et les entrées : la séquence d’entrée décalée d’une unité vers la droite forme les étiquettes, puisque le token suivant est l’étiquette du token actuel. Nous pouvons y parvenir en commençant les étiquettes à partir du deuxième token de la séquence d’entrée, puisque le modèle ne fait pas de prédiction pour le premier token de toute façon. Ensuite, nous coupons le dernier logit, car nous n’avons pas d’étiquette pour le token qui suit la séquence d’entrée complète. Avec cela, nous pouvons calculer la perte par échantillon et compter les occurrences de tous les mots-clés dans chaque échantillon. Enfin, nous calculons la moyenne pondérée sur tous les échantillons en utilisant les occurrences comme poids. Comme nous ne voulons pas rejeter tous les échantillons qui ne contiennent pas de mots-clés, nous ajoutons 1 aux poids :

from torch.nn import CrossEntropyLoss
import torch


def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    # Décalage pour que tokens < n prédisent n
    shift_labels = inputs[..., 1:].contiguous()
    shift_logits = logits[..., :-1, :].contiguous()
    # Calcul de la perte par token
    loss_fct = CrossEntropyLoss(reduce=False)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    # Redimensionnement et perte moyenne par échantillon
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1)
    # Calculer et échelonner la pondération
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum(
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    # Calculer la moyenne pondérée
    weighted_loss = (loss_per_sample * weights).mean()
    return weighted_loss

Avant de commencer à entraîner avec cette nouvelle fonction de perte géniale, nous devons préparer quelques éléments :

• Nous avons besoin de chargeurs de données pour charger les données par batch.
• Nous devons définir les paramètres de décroissance des poids.
• De temps en temps, nous voulons évaluer, il est donc logique d’envelopper le code d’évaluation dans une fonction.

Commençons par les chargeurs de données. Nous avons seulement besoin de définir le format du jeu de données à "torch" et ensuite nous pouvons le passer à un PyTorch DataLoader avec la taille de batch appropriée :

from torch.utils.data.dataloader import DataLoader

tokenized_dataset.set_format("torch")
train_dataloader = DataLoader(tokenized_dataset["train"], batch_size=32, shuffle=True)
eval_dataloader = DataLoader(tokenized_dataset["valid"], batch_size=32)

Ensuite, nous regroupons les paramètres de façon à ce que l’optimiseur sache lesquels bénéficieront d’une décroissance de poids supplémentaire. Habituellement, tous les termes de biais et les poids de la LayerNorm en sont exemptés. Voici comment nous pouvons le faire :

weight_decay = 0.1


def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay):
            params_without_wd.append(p)
        else:
            params_with_wd.append(p)
    return [
        {"params": params_with_wd, "weight_decay": weight_decay},
        {"params": params_without_wd, "weight_decay": 0.0},
    ]

Puisque nous voulons évaluer le modèle régulièrement sur l’ensemble de validation pendant l’entraînement, écrivons une fonction pour cela aussi. Elle passe simplement par le dataloader d’évaluation et rassemble toutes les pertes à travers les processus :

def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch["input_ids"], labels=batch["input_ids"])

        losses.append(accelerator.gather(outputs.loss))
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

Avec la fonction evaluate() nous pouvons rapporter la perte et la perplexité à intervalles réguliers. Ensuite, nous redéfinissons notre modèle pour nous assurer que nous entraînons à nouveau à partir de zéro :

model = GPT2LMHeadModel(config)

Nous pouvons ensuite définir notre optimiseur, en utilisant la fonction précédente pour diviser les paramètres de décroissance des poids :

from torch.optim import AdamW

optimizer = AdamW(get_grouped_params(model), lr=5e-4)

Préparons maintenant le modèle, l’optimiseur et les chargeurs de données pour pouvoir commencer l’entraînement :

from accelerate import Accelerator

accelerator = Accelerator(fp16=True)

model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

Maintenant que nous avons envoyé notre train_dataloader à accelerator.prepare(), nous pouvons utiliser sa longueur pour calculer le nombre d’étapes d’entraînement. Rappelez-vous que nous devons toujours faire cela après avoir préparé le dataloader car cette méthode modifiera sa longueur. Nous utilisons un programme linéaire classique du taux d’apprentissage à 0 :

num_train_epochs = 1
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    name="linear",
    optimizer=optimizer,
    num_warmup_steps=1_000,
    num_training_steps=num_training_steps,
)

Enfin, pour pousser notre modèle vers le Hub, nous aurons besoin de créer un objet Repository dans un dossier de travail. Tout d’abord, connectez-vous au Hub, si vous n’êtes pas déjà connecté. Nous déterminerons le nom du dépôt à partir de l’identifiant du modèle que nous voulons donner à notre modèle (n’hésitez pas à remplacer le repo_name par votre propre choix. Il doit juste contenir votre nom d’utilisateur, ce que fait la fonction get_full_repo_name()) :

from huggingface_hub import Repository, get_full_repo_name

model_name = "codeparrot-ds-accelerate"
repo_name = get_full_repo_name(model_name)
repo_name

'sgugger/codeparrot-ds-accelerate'

Ensuite, nous pouvons cloner ce dépôt dans un dossier local. S’il existe déjà, ce dossier local doit être un clone existant du dépôt avec lequel nous travaillons :

output_dir = "codeparrot-ds-accelerate"
repo = Repository(output_dir, clone_from=repo_name)

Nous pouvons maintenant télécharger tout ce que nous sauvegardons dans output_dir en appelant la méthode repo.push_to_hub(). Cela nous aidera à télécharger les modèles intermédiaires à la fin de chaque époque.

Avant de nous entraîner, exécutons un test rapide pour voir si la fonction d’évaluation fonctionne correctement :

evaluate()

(10.934126853942871, 56057.14453125)

Ce sont des valeurs très élevées pour la perte et la perplexité, mais ce n’est pas surprenant puisque nous n’avons pas encore entraîné le modèle. Avec cela, nous avons tout préparé pour écrire la partie principale du script d’entraînement : la boucle d’entraînement. Dans celle-ci, nous itérons sur le chargeur de données et transmettons les batchs au modèle. Avec les logits, nous pouvons alors évaluer notre fonction de perte personnalisée. Nous mettons à l’échelle la perte par le nombre d’étapes d’accumulation du gradient afin de ne pas créer de plus grandes pertes en agrégeant plus d’étapes. Avant de procéder à l’optimisation, nous découpons également les gradients pour une meilleure convergence. Enfin, tous les quelques pas, nous évaluons le modèle sur l’ensemble d’évaluation avec notre nouvelle fonction evaluate() :

from tqdm.notebook import tqdm

gradient_accumulation_steps = 8
eval_steps = 5_000

model.train()
completed_steps = 0
for epoch in range(num_train_epochs):
    for step, batch in tqdm(
        enumerate(train_dataloader, start=1), total=len(train_dataloader)
    ):
        logits = model(batch["input_ids"]).logits
        loss = keytoken_weighted_loss(batch["input_ids"], logits, keytoken_ids)
        if step % 100 == 0:
            accelerator.print(
                {
                    "lr": get_lr(),
                    "samples": step * samples_per_step,
                    "steps": completed_steps,
                    "loss/train": loss.item() * gradient_accumulation_steps,
                }
            )
        loss = loss / gradient_accumulation_steps
        accelerator.backward(loss)
        if step % gradient_accumulation_steps == 0:
            accelerator.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            completed_steps += 1
        if (step % (eval_steps * gradient_accumulation_steps)) == 0:
            eval_loss, perplexity = evaluate()
            accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
            model.train()
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
            if accelerator.is_main_process:
                tokenizer.save_pretrained(output_dir)
                repo.push_to_hub(
                    commit_message=f"Training in progress step {step}", blocking=False
                )

Et voilà, vous disposez maintenant de votre propre boucle d’entraînement personnalisée pour les modèles de langage causal tels que le GPT-2. Vous pouvez encore l’adapter à vos besoins.