5. Arquitectura LLM

Reading time: 20 minutes

Arquitectura LLM

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

Se puede observar una representación de alto nivel en:

1. Entrada (Texto Tokenizado): El proceso comienza con texto tokenizado, que se convierte en representaciones numéricas.
2. Capa de Embedding de Tokens y Capa de Embedding Posicional: El texto tokenizado se pasa a través de una capa de embedding de tokens y una capa de embedding posicional, que captura la posición de los tokens en una secuencia, crítica para entender el orden de las palabras.
3. Bloques de Transformer: El modelo contiene 12 bloques de transformer, cada uno con múltiples capas. Estos bloques repiten la siguiente secuencia:

• Atención Multi-Cabeza enmascarada: Permite que el modelo se enfoque en diferentes partes del texto de entrada a la vez.
• Normalización de Capa: Un paso de normalización para estabilizar y mejorar el entrenamiento.
• Capa Feed Forward: Responsable de procesar la información de la capa de atención y hacer predicciones sobre el siguiente token.
• Capas de Dropout: Estas capas previenen el sobreajuste al eliminar aleatoriamente unidades durante el entrenamiento.

4. Capa de Salida Final: El modelo produce un tensor de 4x50,257 dimensiones, donde 50,257 representa el tamaño del vocabulario. Cada fila en este tensor corresponde a un vector que el modelo utiliza para predecir la siguiente palabra en la secuencia.
5. Objetivo: El objetivo es tomar estos embeddings y convertirlos de nuevo en texto. Específicamente, la última fila de la salida se utiliza para generar la siguiente palabra, representada como "adelante" en este diagrama.

Representación de Código

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)

Función de Activación 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))
))

Propósito y Funcionalidad

• GELU (Unidad Lineal de Error Gaussiano): Una función de activación que introduce no linealidad en el modelo.
• Activación Suave: A diferencia de ReLU, que anula las entradas negativas, GELU mapea suavemente las entradas a salidas, permitiendo valores pequeños y no nulos para entradas negativas.
• Definición Matemática:

Red Neuronal FeedForward

Las formas se han añadido como comentarios para entender mejor las formas de las 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)

Propósito y Funcionalidad

• Red FeedForward por Posición: Aplica una red completamente conectada de dos capas a cada posición de manera separada e idéntica.
• Detalles de la Capa:
• Primera Capa Lineal: Expande la dimensionalidad de emb_dim a 4 * emb_dim.
• Activación GELU: Aplica no linealidad.
• Segunda Capa Lineal: Reduce la dimensionalidad de nuevo a emb_dim.

Mecanismo de Atención Multi-Cabeza

Esto ya fue explicado en una sección anterior.

Propósito y Funcionalidad

• Autoatención Multi-Cabeza: Permite que el modelo se enfoque en diferentes posiciones dentro de la secuencia de entrada al codificar un token.
• Componentes Clave:
• Consultas, Claves, Valores: Proyecciones lineales de la entrada, utilizadas para calcular puntajes de atención.
• Cabezas: Múltiples mecanismos de atención que funcionan en paralelo (num_heads), cada uno con una dimensión reducida (head_dim).
• Puntajes de Atención: Calculados como el producto punto de consultas y claves, escalados y enmascarados.
• Enmascaramiento: Se aplica una máscara causal para evitar que el modelo preste atención a tokens futuros (importante para modelos autorregresivos como GPT).
• Pesos de Atención: Softmax de los puntajes de atención enmascarados y escalados.
• Vector de Contexto: Suma ponderada de los valores, de acuerdo con los pesos de atención.
• Proyección de Salida: Capa lineal para combinar las salidas de todas las cabezas.

Normalización de Capa

# 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

Propósito y Funcionalidad

• Layer Normalization: Una técnica utilizada para normalizar las entradas a través de las características (dimensiones de embedding) para cada ejemplo individual en un lote.
• Componentes:
• eps: Una constante pequeña (1e-5) añadida a la varianza para prevenir la división por cero durante la normalización.
• scale y shift: Parámetros aprendibles (nn.Parameter) que permiten al modelo escalar y desplazar la salida normalizada. Se inicializan en unos y ceros, respectivamente.
• Proceso de Normalización:
• Calcular Media (mean): Calcula la media de la entrada x a través de la dimensión de embedding (dim=-1), manteniendo la dimensión para la difusión (keepdim=True).
• Calcular Varianza (var): Calcula la varianza de x a través de la dimensión de embedding, también manteniendo la dimensión. El parámetro unbiased=False asegura que la varianza se calcule utilizando el estimador sesgado (dividiendo por N en lugar de N-1), lo cual es apropiado al normalizar sobre características en lugar de muestras.
• Normalizar (norm_x): Resta la media de x y divide por la raíz cuadrada de la varianza más eps.
• Escalar y Desplazar: Aplica los parámetros aprendibles scale y shift a la salida normalizada.

Bloque Transformer

Las formas se han añadido como comentarios para entender mejor las formas de las 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)

Propósito y Funcionalidad

• Composición de Capas: Combina atención multi-cabeza, red de avance, normalización de capas y conexiones residuales.
• Normalización de Capas: Aplicada antes de las capas de atención y avance para un entrenamiento estable.
• Conexiones Residuales (Atajos): Añaden la entrada de una capa a su salida para mejorar el flujo de gradientes y permitir el entrenamiento de redes profundas.
• Dropout: Aplicado después de las capas de atención y avance para regularización.

Funcionalidad Paso a Paso

1. Primer Camino Residual (Autoatención):

• Entrada (shortcut): Guarda la entrada original para la conexión residual.
• Norma de Capa (norm1): Normaliza la entrada.
• Atención Multi-Cabeza (att): Aplica autoatención.
• Dropout (drop_shortcut): Aplica dropout para regularización.
• Añadir Residual (x + shortcut): Combina con la entrada original.

2. Segundo Camino Residual (FeedForward):

• Entrada (shortcut): Guarda la entrada actualizada para la siguiente conexión residual.
• Norma de Capa (norm2): Normaliza la entrada.
• Red de Avance (ff): Aplica la transformación de avance.
• Dropout (drop_shortcut): Aplica dropout.
• Añadir Residual (x + shortcut): Combina con la entrada del primer camino residual.

GPTModel

Las formas se han añadido como comentarios para entender mejor las formas de las 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)

Propósito y Funcionalidad

• Capas de Embedding:
• Embeddings de Tokens (tok_emb): Convierte índices de tokens en embeddings. Como recordatorio, estos son los pesos dados a cada dimensión de cada token en el vocabulario.
• Embeddings Posicionales (pos_emb): Agrega información posicional a los embeddings para capturar el orden de los tokens. Como recordatorio, estos son los pesos dados a los tokens de acuerdo a su posición en el texto.
• Dropout (drop_emb): Aplicado a los embeddings para regularización.
• Bloques de Transformer (trf_blocks): Pila de n_layers bloques de transformer para procesar embeddings.
• Normalización Final (final_norm): Normalización de capa antes de la capa de salida.
• Capa de Salida (out_head): Proyecta los estados ocultos finales al tamaño del vocabulario para producir logits para la predicción.

Número de Parámetros a Entrenar

Teniendo la estructura de GPT definida, es posible averiguar el número de parámetros a entrenar:

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

Cálculo Paso a Paso

1. Capas de Embedding: Embedding de Tokens y Embedding de Posición

• Capa: nn.Embedding(vocab_size, emb_dim)
• Parámetros: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

• Capa: nn.Embedding(context_length, emb_dim)
• Parámetros: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Total de Parámetros de Embedding

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

2. Bloques de Transformador

Hay 12 bloques de transformador, así que calcularemos los parámetros para un bloque y luego multiplicaremos por 12.

Parámetros por Bloque de Transformador

a. Atención Multi-Cabeza

• Componentes:

• Capa Lineal de Consulta (W_query): nn.Linear(emb_dim, emb_dim, bias=False)

• Capa Lineal de Clave (W_key): nn.Linear(emb_dim, emb_dim, bias=False)

• Capa Lineal de Valor (W_value): nn.Linear(emb_dim, emb_dim, bias=False)

• Proyección de Salida (out_proj): nn.Linear(emb_dim, emb_dim)

• Cálculos:

• Cada uno de W_query, W_key, W_value:

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

Dado que hay tres capas de este tipo:

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

• Proyección de Salida (out_proj):

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

• Total de Parámetros de Atención Multi-Cabeza:

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

b. Red FeedForward

• Componentes:

• Primera Capa Lineal: nn.Linear(emb_dim, 4 * emb_dim)

• Segunda Capa Lineal: nn.Linear(4 * emb_dim, emb_dim)

• Cálculos:

• Primera Capa Lineal:

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

• Segunda Capa Lineal:

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 de Parámetros FeedForward:

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

c. Normalizaciones de Capa

• Componentes:
• Dos instancias de LayerNorm por bloque.
• Cada LayerNorm tiene 2 * emb_dim parámetros (escala y desplazamiento).
• Cálculos:

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

d. Total de Parámetros por Bloque de Transformador

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

Total de Parámetros para Todos los Bloques de Transformadores

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

3. Capas Finales

a. Normalización de la Capa Final

• Parámetros: 2 * emb_dim (escalar y desplazar)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Capa de Proyección de Salida (out_head)

• Capa: nn.Linear(emb_dim, vocab_size, bias=False)
• Parámetros: emb_dim * vocab_size

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

4. Resumiendo Todos los Parámetros

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

Generar Texto

Teniendo un modelo que predice el siguiente token como el anterior, solo es necesario tomar los valores del último token de la salida (ya que serán los del token predicho), que será un valor por entrada en el vocabulario y luego usar la función softmax para normalizar las dimensiones en probabilidades que sumen 1 y luego obtener el índice de la entrada más grande, que será el índice de la palabra dentro del vocabulario.

Código 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]))

Referencias

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