5. LLM Architecture

LLM架构

第五阶段的目标非常简单：开发完整LLM的架构。将所有内容整合在一起，应用所有层并创建所有功能以生成文本或将文本转换为ID及其反向操作。

该架构将用于训练和训练后的文本预测。

LLM架构示例来自 https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

可以观察到高层次的表示：

https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31

1. 输入（标记化文本）：该过程以标记化文本开始，该文本被转换为数值表示。

2. 标记嵌入和位置嵌入层：标记化文本通过标记嵌入层和位置嵌入层，后者捕捉序列中标记的位置，这对理解单词顺序至关重要。

3. 变换器块：模型包含12个变换器块，每个块有多个层。这些块重复以下序列：

• 掩蔽多头注意力：允许模型同时关注输入文本的不同部分。

• 层归一化：一个归一化步骤，以稳定和改善训练。

• 前馈层：负责处理来自注意力层的信息并对下一个标记进行预测。

• 丢弃层：这些层通过在训练期间随机丢弃单元来防止过拟合。

1. 最终输出层：模型输出一个4x50,257维的张量，其中50,257表示词汇表的大小。该张量中的每一行对应于模型用于预测序列中下一个单词的向量。

2. 目标：目标是将这些嵌入转换回文本。具体来说，输出的最后一行用于生成下一个单词，在该图中表示为“前进”。

代码表示

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)

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))
))

目的和功能

• GELU (高斯误差线性单元): 一种激活函数，向模型引入非线性。

• 平滑激活: 与ReLU不同，ReLU将负输入归零，GELU平滑地将输入映射到输出，允许负输入的小非零值。

• 数学定义:

在FeedForward层内部的线性层之后使用此函数的目的是将线性数据转换为非线性，以便模型能够学习复杂的非线性关系。

前馈神经网络

已添加形状作为注释，以更好地理解矩阵的形状:

# 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)

目的和功能

• 位置-wise 前馈网络： 对每个位置分别且相同地应用一个两层全连接网络。

• 层详细信息：

• 第一线性层： 将维度从 emb_dim 扩展到 4 * emb_dim。

• GELU 激活： 应用非线性。

• 第二线性层： 将维度减少回 emb_dim。

如您所见，前馈网络使用了 3 层。第一层是一个线性层，它将维度乘以 4，使用线性权重（模型内部训练的参数）。然后，在所有这些维度中使用 GELU 函数应用非线性变化，以捕捉更丰富的表示，最后再使用另一个线性层返回到原始维度大小。

多头注意力机制

这在前面的部分已经解释过。

目的和功能

• 多头自注意力： 允许模型在编码一个标记时关注输入序列中的不同位置。

• 关键组件：

• 查询、键、值： 输入的线性投影，用于计算注意力分数。

• 头： 多个并行运行的注意力机制（num_heads），每个具有减少的维度（head_dim）。

• 注意力分数： 作为查询和键的点积计算，经过缩放和掩蔽。

• 掩蔽： 应用因果掩蔽，以防止模型关注未来的标记（对像 GPT 这样的自回归模型很重要）。

• 注意力权重： 掩蔽和缩放后的注意力分数的 Softmax。

• 上下文向量： 根据注意力权重的值的加权和。

• 输出投影： 线性层以组合所有头的输出。

该网络的目标是找到同一上下文中标记之间的关系。此外，标记被分配到不同的头中，以防止过拟合，尽管每个头找到的最终关系在该网络的末尾被组合在一起。

此外，在训练期间应用 因果掩蔽，以便在查看特定标记的关系时不考虑后来的标记，并且还应用了一些 dropout 以 防止过拟合。

层 归一化

# 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

目的和功能

• 层归一化： 一种用于对每个批次中单个示例的特征（嵌入维度）进行归一化的技术。

• 组件：

• eps： 一个小常数（1e-5），在归一化过程中添加到方差中以防止除以零。

• scale 和 shift： 可学习参数（nn.Parameter），允许模型对归一化输出进行缩放和偏移。它们分别初始化为1和0。

• 归一化过程：

• 计算均值（mean）： 计算输入 x 在嵌入维度（dim=-1）上的均值，同时保持维度以便广播（keepdim=True）。

• 计算方差（var）： 计算 x 在嵌入维度上的方差，同样保持维度。unbiased=False 参数确保方差使用有偏估计量计算（除以 N 而不是 N-1），这在对特征而非样本进行归一化时是合适的。

• 归一化（norm_x）： 从 x 中减去均值，并除以方差加 eps 的平方根。

• 缩放和偏移： 将可学习的 scale 和 shift 参数应用于归一化输出。

目标是确保同一标记的所有维度的均值为0，方差为1。这样做的目的是通过减少内部协变量偏移来稳定深度神经网络的训练，内部协变量偏移是指由于训练过程中参数更新而导致的网络激活分布的变化。

Transformer 块

已添加形状作为注释，以更好地理解矩阵的形状：

# 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)

目的和功能

• 层的组成： 结合多头注意力、前馈网络、层归一化和残差连接。

• 层归一化： 在注意力和前馈层之前应用，以确保稳定的训练。

• 残差连接（快捷方式）： 将层的输入添加到其输出，以改善梯度流并使深层网络的训练成为可能。

• 丢弃法： 在注意力和前馈层之后应用，以进行正则化。

逐步功能

1. 第一个残差路径（自注意力）：

• 输入（shortcut）： 保存原始输入以用于残差连接。

• 层归一化（norm1）： 归一化输入。

• 多头注意力（att）： 应用自注意力。

• 丢弃法（drop_shortcut）： 应用丢弃法以进行正则化。

• 添加残差（x + shortcut）： 与原始输入结合。

1. 第二个残差路径（前馈）：

• 输入（shortcut）： 保存更新后的输入以用于下一个残差连接。

• 层归一化（norm2）： 归一化输入。

• 前馈网络（ff）： 应用前馈变换。

• 丢弃法（drop_shortcut）： 应用丢弃法。

• 添加残差（x + shortcut）： 与第一个残差路径的输入结合。

变换器块将所有网络组合在一起，并应用一些归一化和丢弃法以改善训练的稳定性和结果。 注意丢弃法是在每个网络使用后进行的，而归一化是在之前应用的。

此外，它还使用快捷方式，即将网络的输出与其输入相加。这有助于防止消失梯度问题，确保初始层的贡献与最后一层“相同”。

GPTModel

形状已作为注释添加，以更好地理解矩阵的形状：

# 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)

目的和功能

• 嵌入层:

• 令牌嵌入 (tok_emb): 将令牌索引转换为嵌入。作为提醒，这些是赋予词汇中每个令牌每个维度的权重。

• 位置嵌入 (pos_emb): 向嵌入添加位置信息，以捕捉令牌的顺序。作为提醒，这些是根据令牌在文本中的位置赋予的权重。

• 丢弃 (drop_emb): 应用于嵌入以进行正则化。

• 变换器块 (trf_blocks): 一组 n_layers 变换器块，用于处理嵌入。

• 最终归一化 (final_norm): 在输出层之前进行层归一化。

• 输出层 (out_head): 将最终隐藏状态投影到词汇大小，以生成预测的 logits。

该类的目标是使用所有其他提到的网络来预测序列中的下一个令牌，这对于文本生成等任务至关重要。

注意它将使用尽可能多的变换器块，并且每个变换器块使用一个多头注意力网络、一个前馈网络和几个归一化。因此，如果使用12个变换器块，则将其乘以12。

此外，在输出之前添加了一个归一化层，并在最后应用一个最终线性层以获得适当维度的结果。注意每个最终向量的大小与使用的词汇相同。这是因为它试图为词汇中的每个可能令牌获取一个概率。

训练的参数数量

定义了GPT结构后，可以找出要训练的参数数量：

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

逐步计算

1. 嵌入层：标记嵌入和位置嵌入

• 层： nn.Embedding(vocab_size, emb_dim)

• 参数： vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

• 层: nn.Embedding(context_length, emb_dim)

• 参数: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

总嵌入参数

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

2. Transformer Blocks

有12个变换器块，因此我们将计算一个块的参数，然后乘以12。

每个变换器块的参数

a. 多头注意力

• 组件：

• 查询线性层 (W_query): nn.Linear(emb_dim, emb_dim, bias=False)

• 键线性层 (W_key): nn.Linear(emb_dim, emb_dim, bias=False)

• 值线性层 (W_value): nn.Linear(emb_dim, emb_dim, bias=False)

• 输出投影 (out_proj): nn.Linear(emb_dim, emb_dim)

• 计算：

• 每个 W_query、W_key、W_value:

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

由于有三个这样的层：

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

• 输出投影 (out_proj):

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

• 总多头注意力参数：

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

b. 前馈网络

• 组件：

• 第一线性层: nn.Linear(emb_dim, 4 * emb_dim)

• 第二线性层: nn.Linear(4 * emb_dim, emb_dim)

• 计算：

• 第一线性层：

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

• 第二线性层：

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

• 总前馈参数：

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

c. 层归一化

• 组件：

• 每个块有两个 LayerNorm 实例。

• 每个 LayerNorm 有 2 * emb_dim 参数（缩放和偏移）。

• 计算：

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

d. 每个变换器块的总参数

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

所有变换器块的总参数

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

3. 最终层

a. 最终层归一化

• 参数： 2 * emb_dim（缩放和偏移）

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. 输出投影层 (out_head)

• 层: nn.Linear(emb_dim, vocab_size, bias=False)

• 参数: emb_dim * vocab_size

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

4. 汇总所有参数

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

生成文本

拥有一个像之前那样预测下一个标记的模型，只需从输出中获取最后一个标记的值（因为它们将是预测标记的值），这将是词汇表中每个条目的值，然后使用softmax函数将维度归一化为总和为1的概率，然后获取最大条目的索引，这将是词汇表中单词的索引。

来自 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]))

References

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