5. Architecture LLM

Reading time: 18 minutes

Architecture LLM

Exemple d'architecture LLM de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb :

Une représentation de haut niveau peut être observée dans :

1. Entrée (Texte Tokenisé) : Le processus commence par du texte tokenisé, qui est converti en représentations numériques.
2. Couche d'Embedding de Token et Couche d'Embedding Positionnel : Le texte tokenisé est passé par une couche d'embedding de token et une couche d'embedding positionnel, qui capture la position des tokens dans une séquence, critique pour comprendre l'ordre des mots.
3. Blocs Transformer : Le modèle contient 12 blocs transformer, chacun avec plusieurs couches. Ces blocs répètent la séquence suivante :

• Attention Multi-Tête Masquée : Permet au modèle de se concentrer sur différentes parties du texte d'entrée en même temps.
• Normalisation de Couche : Une étape de normalisation pour stabiliser et améliorer l'entraînement.
• Couche Feed Forward : Responsable du traitement des informations de la couche d'attention et de la prédiction du prochain token.
• Couches de Dropout : Ces couches empêchent le surapprentissage en supprimant aléatoirement des unités pendant l'entraînement.

4. Couche de Sortie Finale : Le modèle produit un tenseur de dimension 4x50,257, où 50,257 représente la taille du vocabulaire. Chaque ligne de ce tenseur correspond à un vecteur que le modèle utilise pour prédire le prochain mot dans la séquence.
5. Objectif : L'objectif est de prendre ces embeddings et de les convertir à nouveau en texte. Plus précisément, la dernière ligne de la sortie est utilisée pour générer le prochain mot, représenté comme "forward" dans ce diagramme.

Représentation du Code

import torch
import torch.nn as nn
import tiktoken

class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
return self.layers(x)

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

# Shortcut connection for feed forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

return x


class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(
cfg["emb_dim"], cfg["vocab_size"], bias=False
)

def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
out = model(batch)
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

Fonction d'activation GELU

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

Objectif et Fonctionnalité

• GELU (Unité Linéaire d'Erreur Gaussienne) : Une fonction d'activation qui introduit de la non-linéarité dans le modèle.
• Activation Douce : Contrairement à ReLU, qui annule les entrées négatives, GELU mappe en douceur les entrées aux sorties, permettant des valeurs petites et non nulles pour les entrées négatives.
• Définition Mathématique :

Réseau de Neurones FeedForward

Des formes ont été ajoutées en tant que commentaires pour mieux comprendre les formes des matrices :

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[1](x) # x shape remains: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[2](x) # x shape: (batch_size, seq_len, emb_dim)
return x  # Output shape: (batch_size, seq_len, emb_dim)

Objectif et Fonctionnalité

• Réseau FeedForward par Position : Applique un réseau entièrement connecté à deux couches à chaque position séparément et de manière identique.
• Détails des Couches :
• Première Couche Linéaire : Augmente la dimensionnalité de emb_dim à 4 * emb_dim.
• Activation GELU : Applique une non-linéarité.
• Deuxième Couche Linéaire : Réduit la dimensionnalité à nouveau à emb_dim.

Mécanisme d'Attention Multi-Tête

Cela a déjà été expliqué dans une section précédente.

Objectif et Fonctionnalité

• Auto-Attention Multi-Tête : Permet au modèle de se concentrer sur différentes positions au sein de la séquence d'entrée lors de l'encodage d'un token.
• Composants Clés :
• Requêtes, Clés, Valeurs : Projections linéaires de l'entrée, utilisées pour calculer les scores d'attention.
• Têtes : Plusieurs mécanismes d'attention fonctionnant en parallèle (num_heads), chacun avec une dimension réduite (head_dim).
• Scores d'Attention : Calculés comme le produit scalaire des requêtes et des clés, mis à l'échelle et masqués.
• Masquage : Un masque causal est appliqué pour empêcher le modèle de prêter attention aux tokens futurs (important pour les modèles autorégressifs comme GPT).
• Poids d'Attention : Softmax des scores d'attention masqués et mis à l'échelle.
• Vecteur de Contexte : Somme pondérée des valeurs, selon les poids d'attention.
• Projection de Sortie : Couche linéaire pour combiner les sorties de toutes les têtes.

Normalisation de Couche

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5 # Prevent division by zero during normalization.
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

Objectif et Fonctionnalité

• Layer Normalization : Une technique utilisée pour normaliser les entrées à travers les caractéristiques (dimensions d'embedding) pour chaque exemple individuel dans un lot.
• Composants :
• eps : Une petite constante (1e-5) ajoutée à la variance pour éviter la division par zéro lors de la normalisation.
• scale et shift : Paramètres apprenables (nn.Parameter) qui permettent au modèle de mettre à l'échelle et de décaler la sortie normalisée. Ils sont initialisés à un et zéro, respectivement.
• Processus de Normalisation :
• Calculer la Moyenne (mean) : Calcule la moyenne de l'entrée x à travers la dimension d'embedding (dim=-1), en gardant la dimension pour le broadcasting (keepdim=True).
• Calculer la Variance (var) : Calcule la variance de x à travers la dimension d'embedding, en gardant également la dimension. Le paramètre unbiased=False garantit que la variance est calculée en utilisant l'estimateur biaisé (en divisant par N au lieu de N-1), ce qui est approprié lors de la normalisation sur les caractéristiques plutôt que sur les échantillons.
• Normaliser (norm_x) : Soustrait la moyenne de x et divise par la racine carrée de la variance plus eps.
• Mettre à l'échelle et Décaler : Applique les paramètres apprenables scale et shift à la sortie normalisée.

Bloc Transformer

Des formes ont été ajoutées en tant que commentaires pour mieux comprendre les formes des matrices :

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"]
)
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for attention block
shortcut = x  # shape: (batch_size, seq_len, emb_dim)
x = self.norm1(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.att(x)    # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for feedforward block
shortcut = x       # shape: (batch_size, seq_len, emb_dim)
x = self.norm2(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.ff(x)     # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

return x  # Output shape: (batch_size, seq_len, emb_dim)

Objectif et Fonctionnalité

• Composition des Couches : Combine l'attention multi-têtes, le réseau feedforward, la normalisation de couche et les connexions résiduelles.
• Normalisation de Couche : Appliquée avant les couches d'attention et feedforward pour un entraînement stable.
• Connexions Résiduelles (Raccourcis) : Ajoutent l'entrée d'une couche à sa sortie pour améliorer le flux de gradient et permettre l'entraînement de réseaux profonds.
• Dropout : Appliqué après les couches d'attention et feedforward pour la régularisation.

Fonctionnalité Étape par Étape

1. Premier Chemin Résiduel (Auto-Attention) :

• Entrée (shortcut) : Sauvegarder l'entrée originale pour la connexion résiduelle.
• Norme de Couche (norm1) : Normaliser l'entrée.
• Attention Multi-Têtes (att) : Appliquer l'auto-attention.
• Dropout (drop_shortcut) : Appliquer le dropout pour la régularisation.
• Ajouter Résiduel (x + shortcut) : Combiner avec l'entrée originale.

2. Deuxième Chemin Résiduel (FeedForward) :

• Entrée (shortcut) : Sauvegarder l'entrée mise à jour pour la prochaine connexion résiduelle.
• Norme de Couche (norm2) : Normaliser l'entrée.
• Réseau FeedForward (ff) : Appliquer la transformation feedforward.
• Dropout (drop_shortcut) : Appliquer le dropout.
• Ajouter Résiduel (x + shortcut) : Combiner avec l'entrée du premier chemin résiduel.

GPTModel

Des formes ont été ajoutées en tant que commentaires pour mieux comprendre les formes des matrices :

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
# shape: (vocab_size, emb_dim)

self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
# shape: (context_length, emb_dim)

self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
# Stack of TransformerBlocks

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
# shape: (emb_dim, vocab_size)

def forward(self, in_idx):
# in_idx shape: (batch_size, seq_len)
batch_size, seq_len = in_idx.shape

# Token embeddings
tok_embeds = self.tok_emb(in_idx)
# shape: (batch_size, seq_len, emb_dim)

# Positional embeddings
pos_indices = torch.arange(seq_len, device=in_idx.device)
# shape: (seq_len,)
pos_embeds = self.pos_emb(pos_indices)
# shape: (seq_len, emb_dim)

# Add token and positional embeddings
x = tok_embeds + pos_embeds  # Broadcasting over batch dimension
# x shape: (batch_size, seq_len, emb_dim)

x = self.drop_emb(x)  # Dropout applied
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.trf_blocks(x)  # Pass through Transformer blocks
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.final_norm(x)  # Final LayerNorm
# x shape remains: (batch_size, seq_len, emb_dim)

logits = self.out_head(x)  # Project to vocabulary size
# logits shape: (batch_size, seq_len, vocab_size)

return logits  # Output shape: (batch_size, seq_len, vocab_size)

Objectif et Fonctionnalité

• Couches d'Incorporation :
• Incorporations de Tokens (tok_emb) : Convertit les indices de tokens en incorporations. En rappel, ce sont les poids attribués à chaque dimension de chaque token dans le vocabulaire.
• Incorporations Positionnelles (pos_emb) : Ajoute des informations positionnelles aux incorporations pour capturer l'ordre des tokens. En rappel, ce sont les poids attribués aux tokens selon leur position dans le texte.
• Dropout (drop_emb) : Appliqué aux incorporations pour la régularisation.
• Blocs Transformer (trf_blocks) : Empilement de n_layers blocs transformer pour traiter les incorporations.
• Normalisation Finale (final_norm) : Normalisation de couche avant la couche de sortie.
• Couche de Sortie (out_head) : Projette les états cachés finaux à la taille du vocabulaire pour produire des logits pour la prédiction.

Nombre de Paramètres à Entraîner

Ayant la structure GPT définie, il est possible de déterminer le nombre de paramètres à entraîner :

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

model = GPTModel(GPT_CONFIG_124M)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")
# Total number of parameters: 163,009,536

Calcul de l'étape par étape

1. Couches d'incorporation : Incorporation de jetons et incorporation de position

• Couche : nn.Embedding(vocab_size, emb_dim)
• Paramètres : vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

• Couche : nn.Embedding(context_length, emb_dim)
• Paramètres : context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Paramètres d'Embedding Totaux

embedding_params = token_embedding_params + position_embedding_params
embedding_params = 38,597,376 + 786,432 = 39,383,808

2. Blocs de Transformateur

Il y a 12 blocs de transformateur, donc nous allons calculer les paramètres pour un bloc et ensuite multiplier par 12.

Paramètres par Bloc de Transformateur

a. Attention Multi-Tête

• Composants :

• Couche Linéaire de Requête (W_query) : nn.Linear(emb_dim, emb_dim, bias=False)

• Couche Linéaire de Clé (W_key) : nn.Linear(emb_dim, emb_dim, bias=False)

• Couche Linéaire de Valeur (W_value) : nn.Linear(emb_dim, emb_dim, bias=False)

• Projection de Sortie (out_proj) : nn.Linear(emb_dim, emb_dim)

• Calculs :

• Chacune de W_query, W_key, W_value :

qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824

Puisqu'il y a trois de ces couches :

total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472

• Projection de Sortie (out_proj) :

out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592

• Total des Paramètres d'Attention Multi-Tête :

mha_params = total_qkv_params + out_proj_params
mha_params = 1,769,472 + 590,592 = 2,360,064

b. Réseau FeedForward

• Composants :

• Première Couche Linéaire : nn.Linear(emb_dim, 4 * emb_dim)

• Deuxième Couche Linéaire : nn.Linear(4 * emb_dim, emb_dim)

• Calculs :

• Première Couche Linéaire :

ff_first_layer_params = (emb_dim * 4 * emb_dim) + (4 * emb_dim)
ff_first_layer_params = (768 * 3072) + 3072 = 2,359,296 + 3,072 = 2,362,368

• Deuxième Couche Linéaire :

ff_second_layer_params = (4 * emb_dim * emb_dim) + emb_dim
ff_second_layer_params = (3072 * 768) + 768 = 2,359,296 + 768 = 2,360,064

• Total des Paramètres FeedForward :

ff_params = ff_first_layer_params + ff_second_layer_params
ff_params = 2,362,368 + 2,360,064 = 4,722,432

c. Normalisations de Couche

• Composants :
• Deux instances de LayerNorm par bloc.
• Chaque LayerNorm a 2 * emb_dim paramètres (échelle et décalage).
• Calculs :

layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072

d. Total des Paramètres par Bloc de Transformateur

pythonCopy codeparams_per_block = mha_params + ff_params + layer_norm_params_per_block
params_per_block = 2,360,064 + 4,722,432 + 3,072 = 7,085,568

Paramètres totaux pour tous les blocs de transformateur

pythonCopy codetotal_transformer_blocks_params = params_per_block * n_layers
total_transformer_blocks_params = 7,085,568 * 12 = 85,026,816

3. Couches finales

a. Normalisation de la couche finale

• Paramètres : 2 * emb_dim (échelle et décalage)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Couche de projection de sortie (out_head)

• Couche : nn.Linear(emb_dim, vocab_size, bias=False)
• Paramètres : emb_dim * vocab_size

pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376

4. Résumé de tous les paramètres

pythonCopy codetotal_params = (
embedding_params +
total_transformer_blocks_params +
final_layer_norm_params +
output_projection_params
)
total_params = (
39,383,808 +
85,026,816 +
1,536 +
38,597,376
)
total_params = 163,009,536

Générer du texte

Avoir un modèle qui prédit le prochain token comme le précédent, il suffit de prendre les valeurs du dernier token de la sortie (car ce seront celles du token prédit), ce qui sera une valeur par entrée dans le vocabulaire et ensuite utiliser la fonction softmax pour normaliser les dimensions en probabilités qui s'additionnent à 1 et ensuite obtenir l'index de la plus grande entrée, qui sera l'index du mot dans le vocabulaire.

Code de https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

def generate_text_simple(model, idx, max_new_tokens, context_size):
# idx is (batch, n_tokens) array of indices in the current context
for _ in range(max_new_tokens):

# Crop current context if it exceeds the supported context size
# E.g., if LLM supports only 5 tokens, and the context size is 10
# then only the last 5 tokens are used as context
idx_cond = idx[:, -context_size:]

# Get the predictions
with torch.no_grad():
logits = model(idx_cond)

# Focus only on the last time step
# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
logits = logits[:, -1, :]

# Apply softmax to get probabilities
probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

# Get the idx of the vocab entry with the highest probability value
idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

# Append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

return idx


start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)

encoded_tensor = torch.tensor(encoded).unsqueeze(0)
print("encoded_tensor.shape:", encoded_tensor.shape)

model.eval() # disable dropout

out = generate_text_simple(
model=model,
idx=encoded_tensor,
max_new_tokens=6,
context_size=GPT_CONFIG_124M["context_length"]
)

print("Output:", out)
print("Output length:", len(out[0]))

Références

• https://www.manning.com/books/build-a-large-language-model-from-scratch