实现完整Bigram模型

import torch
import torch.nn as nn
from torch.nn import functional as F

# 超参数设置
batch_size = 32
block_size = 8
max_iters = 3000
eval_interval = 300
learning_rate = 1e-3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200

# 数据加载与预处理
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# 构建词汇表
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# 划分训练集和验证集
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data))
train_data = data[:n]
val_data = data[n:]

# 数据加载器
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

# 评估时禁用梯度计算
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

# Bigram 语言模型核心
class BigramLanguageModel(nn.Module):

    def __init__(self, vocab_size):
        super().__init__()
        # 词嵌入表：每个词元直接查下一个词元的logits
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # idx 和 targets 的形状都是 (B, T)
        # logits 的形状是 (B, T, C)
        logits = self.token_embedding_table(idx)

        if targets is None:
            loss = None
        else:
            # 为了计算交叉熵，需要调整形状
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx 是 (B, T) 形状的当前上下文
        for _ in range(max_new_tokens):
            # 获取预测
            logits, loss = self(idx)
            # 只关注最后一个时间步
            logits = logits[:, -1, :] # 变成 (B, C)
            # 计算概率
            probs = F.softmax(logits, dim=-1) # (B, C)
            # 采样
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            # 拼接
            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)
        return idx

# 初始化模型并训练
model = BigramLanguageModel(vocab_size)
model = model.to(device)

# 初始化优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# 训练循环
for iter in range(max_iters):

    # 定期评估损失
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # 获取一个 batch 的数据
    xb, yb = get_batch('train')

    # 前向传播，计算损失
    logits, loss = model(xb, yb)

    # 反向传播并更新参数
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# 生成文本
print(decode(model.generate(idx = torch.zeros((1, 1), dtype=torch.long).to(device), max_new_tokens=500)[0].tolist()))

加入单头自注意力机制：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Head(nn.Module):  # 定义一个单头注意力模块，继承nn.Module
    """ one head of self-attention """
    # head_size：每个注意力头的向量维度
    def __init__(self, head_size):  
        super().__init__()
        # 三个线性层，分别把输入特征映射成 K、Q、V
        # 输入n_embd → 输出key：head_size
        self.key = nn.Linear(n_embd, head_size, bias=False)
        # 输入n_embd → 输出query：head_size
        self.query = nn.Linear(n_embd, head_size, bias=False)  
        # 输入n_embd → 输出value：head_size
        self.value = nn.Linear(n_embd, head_size, bias=False) 
        # 注册下三角掩码tril，不是模型参数，是固定常量，所以用register_buffer
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):  # x 是输入：形状 (B, T, C) → B批次，T序列长度，C通道=n_embd
        B, T, C = x.shape  # 把输入形状解包，方便后面操作

        # 步骤1：把输入x分别投影成 K、Q、V
        k = self.key(x)    # (B, T, head_size) —— 每个位置的key向量
        q = self.query(x)  # (B, T, head_size) —— 每个位置的query向量

        # 步骤2：计算注意力分数（相似度）Q @ K^T
        # 公式：注意力分数 = Q @ K.T / sqrt(head_size)
        wei = q @ k.transpose(-2, -1) * C**-0.5  # (B,T,T) —— 相似度矩阵
        # transpose(-2,-1)：把最后两维转置，(B,T,hs) → (B,hs,T)
        # 除以sqrt(C)：缩放，防止方差爆炸，softmax不饱和

        # 步骤3：decoder专用掩码，遮住未来位置（看不到后面的词）
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        # tril[:T,:T]：动态适配当前序列长度T
        # 把上三角全部填成 -∞，softmax后就变成0，看不到未来

        # 步骤4：softmax归一化，变成概率权重（每行和为1）
        wei = F.softmax(wei, dim=-1)  # (B,T,T) —— 注意力权重矩阵

        # 步骤5：计算V，然后加权求和
        v = self.value(x)  # (B, T, head_size) —— 每个位置真实要传递的信息
        out = wei @ v      # (B,T,T) @ (B,T,hs) → (B,T,hs) —— 加权聚合信息

        return out  # 返回单头注意力的输出

class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 词嵌入层：把词索引 → 词向量 (vocab_size → n_embd)
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 位置嵌入层：给每个位置0~T-1一个向量，让模型知道顺序
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 单头自注意力层（这节课新加的核心）
        self.sa_head = Head(n_embd)
        # 语言模型头：把最终特征 → 词汇表概率logits (n_embd → vocab_size)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape  # idx形状(B,T)：一批句子，每句T个词

        # 1）词嵌入：拿到每个词的向量
        tok_emb = self.token_embedding_table(idx)  # (B, T, n_embd)

        # 2）位置嵌入：给每个位置0~T-1编码
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, n_embd)

        # 3）词向量 + 位置向量 = 模型真正的输入特征
        x = tok_emb + pos_emb  # (B, T, n_embd)

        # 4）过单头自注意力
        x = self.sa_head(x)  # (B, T, n_embd) —— 注意力做完后的特征

        # 5）语言模型头：输出每个位置下一个词的预测分数
        logits = self.lm_head(x)  # (B, T, vocab_size)

        # 下面是计算损失（训练用）
        if targets is None:
            loss = None  # 生成模式，不需要loss
        else:
            B, T, C = logits.shape
            # cross_entropy要求输入形状：(N, C)，目标：(N)
            logits = logits.view(B*T, C)      # 展平 (B*T, vocab_size)
            targets = targets.view(B*T)       # 展平 (B*T)
            loss = F.cross_entropy(logits, targets)  # 计算交叉熵损失

        return logits, loss

    # 生成函数：自回归生成文本
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):  # 循环生成max_new_tokens个词
            # 防止输入超过block_size，只取最后block_size个
            idx_cond = idx[:, -block_size:]
            # 前向传播，拿到logits
            logits, loss = self(idx_cond)
            # 只看最后一个时间步，预测下一个词
            logits = logits[:, -1, :]  # (B, vocab_size)
            # 转概率
            probs = F.softmax(logits, dim=-1)
            # 按概率采样一个词
            idx_next = torch.multinomial(probs, num_samples=1)
            # 拼接到序列后面
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

改成多头自注意力机制：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 单头注意力 Head（保留，因为多头要用到）
class Head(nn.Module):
    """ one head of self-attention """
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # 下三角掩码，注册为常量 buffer，不是参数
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape          # x 形状 (B, T, C)，C=n_embd
        k = self.key(x)            # (B, T, head_size)
        q = self.query(x)          # (B, T, head_size)
        
        # 计算相似度 Q@K.T / sqrt(head_size)
        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5   # (B, T, T)
        # decoder 掩码：看不到未来的词
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1) # 变成权重，每行和为1
        
        v = self.value(x)          # (B, T, head_size)
        out = wei @ v              # (B, T, head_size)
        return out

# 添加内容：多头注意力 MultiHeadAttention
class MultiHeadAttention(nn.Module):
    """多个注意力头拼在一起，就是多头注意力"""
    # 参数：num_heads 头数，head_size 每个头的维度
    def __init__(self, num_heads, head_size):
        super().__init__()
        # 把多个 Head 放进 ModuleList，让 PyTorch 能识别参数
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        # 最后一个投影层：把拼起来的特征再映射回 n_embd
        self.proj = nn.Linear(num_heads * head_size, n_embd)

    def forward(self, x):
        # 1. 每个头独立算一遍自注意力，结果拼在最后一维
        # 每个头输出 (B, T, head_size)
        # 拼接后 → (B, T, num_heads * head_size)
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        # 2. 过一个线性层投影，把维度变回 n_embd
        out = self.proj(out)
        return out

FeedForward前馈网络：

# 单头注意力 Head
class Head(nn.Module):
    """ 单头自注意力 """
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)

        # 计算注意力分数 + 缩放
        wei = q @ k.transpose(-2, -1) * C**-0.5
        # 掩码，看不到未来
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)

        v = self.value(x)
        out = wei @ v
        return out

# 多头注意力 MultiHeadAttention
class MultiHeadAttention(nn.Module):
    """ 多头注意力 = 多个头拼接 """
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(num_heads * head_size, n_embd)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.proj(out)
        return out

# 添加内容：FeedForward 前馈网络
class FeedForward(nn.Module):
    """ 前馈网络：对每个token单独做非线性变换 """
    def __init__(self, n_embd):
        super().__init__()
        # 串行层：线性→激活→线性
        self.net = nn.Sequential(
            # 先把维度放大4倍：n_embd → 4*n_embd
            nn.Linear(n_embd, 4 * n_embd),
            # ReLU非线性，让模型能学复杂关系
            nn.ReLU(),
            # 再投影回n_embd，保持维度不变
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        # 输入形状 (B, T, n_embd) → 输出形状不变
        return self.net(x)

# 主模型：把多头注意力 + FeedForward 用起来
class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size, n_embd, block_size, n_head):
        super().__init__()
        self.block_size = block_size
        # 词嵌入
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 位置嵌入
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 多头自注意力（之前已有）
        self.sa = MultiHeadAttention(n_head, n_embd // n_head)
        # 新增：前馈网络 FeedForward
        self.ffwd = FeedForward(n_embd)
        # 语言模型头，输出词汇概率
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        # 词嵌入 + 位置嵌入
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb

        # 第一步：过多头注意力
        x = self.sa(x)
        # 第二步：过前馈网络（本节课新增的使用位置）
        x = self.ffwd(x)

        # 最后输出logits
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

Block将自注意力self-attention和feed forward打包并且使用层归一化使训练更稳定，使用残差连接防止梯度消失：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 1. 单头注意力 Head
class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)

        wei = q @ k.transpose(-2, -1) * C**-0.5
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)

        v = self.value(x)
        out = wei @ v
        return out

# 2. 多头注意力 MultiHeadAttention
class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(num_heads * head_size, n_embd)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.proj(out)
        return out

# 3. FeedForward 前馈网络
class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
        )

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

# 4. 新增：Block 块（打包注意力+前馈+残差+层归一化）
class Block(nn.Module):
    """ Transformer块： communication followed by computation """
    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head  # 每个头的维度 = 总维度 ÷ 头数
        self.sa = MultiHeadAttention(n_head, head_size)  # 多头注意力
        self.ffwd = FeedForward(n_embd)                   # 前馈网络
        self.ln1 = nn.LayerNorm(n_embd)  # 第一个层归一化
        self.ln2 = nn.LayerNorm(n_embd)  # 第二个层归一化

    def forward(self, x):
        # 残差1：先LN → 注意力 → 加回x
        x = x + self.sa(self.ln1(x))
        # 残差2：先LN → 前馈 → 加回x
        x = x + self.ffwd(self.ln2(x))
        return x


# 5. 主模型：现在用 Block 替换原来的 sa + ffwd
class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size, n_embd, block_size, n_head):
        super().__init__()
        self.block_size = block_size
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 原来：sa + ffwd
        # 现在：直接用一个 Block（里面已经包含sa+ffwd+残差+LN）
        self.block = Block(n_embd, n_head)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb

        # 原来：x = self.sa(x)
        #      x = self.ffwd(x)
        # 现在：直接过一个 Block
        x = self.block(x)

        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

加入Dropout防止过拟合：

# 导入必要库：torch核心、神经网络模块、函数库
import torch
import torch.nn as nn
import torch.nn.functional as F

# ------------------- 超参数定义（训练配置，可调整） -------------------
batch_size = 64        # 每次训练取64个样本（批次大小）
block_size = 256       # 输入序列最大长度（一次看256个token）
max_iters = 5000       # 总训练迭代次数
eval_interval = 500    # 每500步评估一次损失
learning_rate = 3e-4   # 优化器学习率
# 自动选择设备：有GPU用cuda，没有用cpu
device = 'cuda' if torch.cuda.is_available() else 'cpu'
eval_iters = 200       # 评估时取200次损失求平均，更稳定
n_embd = 384           # 词嵌入/特征维度（每个token用384维向量表示）
n_head = 6             # 多头注意力头数（6个并行头）
n_layer = 6            # 堆叠6个Transformer Block（6层）
dropout = 0.2          # Dropout概率：随机丢弃20%神经元，防过拟合

# ------------------- 数据读取与编码处理 -------------------
torch.manual_seed(1337)  # 固定随机种子，保证结果可复现

# 读取文本数据集（input.txt需放在同目录下）
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

# 构建词汇表：提取文本中所有唯一字符并排序
chars = sorted(list(set(text)))
vocab_size = len(chars)  # 词汇表总大小

# 字符→数字编码字典 stoi；数字→字符解码字典 itos
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }
encode = lambda s: [stoi[c] for c in s]      # 编码函数：字符串→数字列表
decode = lambda l: ''.join([itos[i] for i in l]) # 解码函数：数字列表→字符串

# 将全文编码为张量，划分训练集(90%)、验证集(10%)
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data))
train_data = data[:n]  # 训练数据
val_data = data[n:]    # 验证数据

# 批次采样函数：随机取一段连续序列，构造输入x和标签y（下一个token）
def get_batch(split):
    data = train_data if split == 'train' else val_data
    # 随机生成batch_size个起始位置
    ix = torch.randint(len(data) - block_size, (batch_size,))
    # 构造输入x：从i开始取block_size个token
    x = torch.stack([data[i:i+block_size] for i in ix])
    # 构造标签y：x对应位置的下一个token
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)  # 数据移到指定设备
    return x, y

# 损失评估函数：无梯度计算，用于统计训练/验证损失
@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()  # 切换为评估模式（关闭dropout、batchnorm等）
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()  # 取平均损失
    model.train()  # 切回训练模式
    return out

# ------------------- 模型组件定义：从底层到Block -------------------

# 单头自注意力模块
class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        # 线性层生成Q、K、V，无偏置
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # 注册下三角掩码缓冲区，防止看到未来token
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)  # 注意力权重dropout，第七课新增

    def forward(self, x):
        B, T, C = x.shape  # B批次、T序列长度、C嵌入维度
        k = self.key(x)    # 得到key (B,T,head_size)
        q = self.query(x)  # 得到query (B,T,head_size)

        # 计算注意力分数：Q@K转置 + 缩放（防softmax饱和）
        wei = q @ k.transpose(-2,-1) * C**-0.5
        # 掩码：未来位置置为负无穷，softmax后为0
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)  # 归一化为注意力权重
        wei = self.dropout(wei)       # 对权重dropout
        v = self.value(x)             # 得到value (B,T,head_size)
        out = wei @ v                 # 加权求和得到输出
        return out

# 多头注意力：多个单头拼接+投影
class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        # 并行创建num_heads个单头注意力
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)  # 投影层，映射回原维度
        self.dropout = nn.Dropout(dropout)     # 输出dropout，第七课新增

    def forward(self, x):
        # 拼接所有头的输出（最后一维拼接）
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        # 投影+dropout，映射回残差路径
        out = self.dropout(self.proj(out))
        return out

# 前馈网络：每个token独立做非线性变换
class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),  # 维度放大4倍
            nn.ReLU(),                      # 非线性激活
            nn.Linear(4 * n_embd, n_embd),   # 投影回原维度
            nn.Dropout(dropout),             # 输出dropout，第七课新增
        )

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

# Transformer Block：注意力+前馈+残差+层归一化（pre-norm）
class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head  # 单头维度=总维度÷头数
        self.sa = MultiHeadAttention(n_head, head_size)  # 多头注意力
        self.ffwd = FeedForward(n_embd)                   # 前馈网络
        self.ln1 = nn.LayerNorm(n_embd)  # 第一个层归一化
        self.ln2 = nn.LayerNorm(n_embd)  # 第二个层归一化

    def forward(self, x):
        # 残差1：先LN→注意力→加回原x
        x = x + self.sa(self.ln1(x))
        # 残差2：先LN→前馈→加回原x
        x = x + self.ffwd(self.ln2(x))
        return x

# 主模型：GPT语言模型
class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 词嵌入层：数字token→向量
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 位置嵌入层：给每个位置编码位置信息
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 堆叠n_layer个Block
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)  # 最终层归一化
        self.lm_head = nn.Linear(n_embd, vocab_size)  # 输出层：映射到词汇表概率

    def forward(self, idx, targets=None):
        B, T = idx.shape
        # 词嵌入 + 位置嵌入，相加得到输入特征
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb
        # 过所有Transformer层
        x = self.blocks(x)
        x = self.ln_f(x)
        # 得到输出logits（未softmax的分数）
        logits = self.lm_head(x)

        # 计算损失（有标签时）
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)      # 展平为二维，适配交叉熵
            targets = targets.view(B*T)       # 标签展平
            loss = F.cross_entropy(logits, targets)  # 交叉熵损失

        return logits, loss

    # 生成函数：根据上下文生成新token
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            # 截断上下文，只保留最后block_size个token
            idx_cond = idx[:, -block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]  # 只取最后一个token的预测
            probs = F.softmax(logits, dim=-1)  # 转概率
            idx_next = torch.multinomial(probs, num_samples=1) # 采样下一个token
            idx = torch.cat((idx, idx_next), dim=1)  # 拼接到上下文
        return idx

# ------------------- 模型初始化与训练 -------------------
model = GPTLanguageModel()    # 创建模型
model = model.to(device)      # 模型移到指定设备
# 优化器：AdamW，适配大模型训练，含权重衰减
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# 训练循环
for iter in range(max_iters):
    # 每eval_interval步，评估并打印损失
    if iter % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    # 取一个训练批次
    xb, yb = get_batch('train')
    # 前向传播：得到logits和损失
    logits, loss = model(xb, yb)
    # 梯度清零→反向传播→参数更新
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# ------------------- 文本生成 -------------------
# 初始上下文：1x1的0（代表起始符）
context = torch.zeros((1, 1), dtype=torch.long, device=device)
# 生成500个token并解码打印
print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))