"""
=============================================================================
第 7 讲 · 代码 4：迷你 GPT —— 在第 5 讲的玩具语言上训练一个 Transformer 语言模型
=============================================================================
对应讲义：lectures/lecture07_self_attention_transformer.md 第 8、11 节
-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
把 transformer.py 里的零件拼成一个"解码器结构"的 Transformer 语言模型（和 GPT 一样：
词嵌入 + 可学习的位置嵌入 + 2 个 Pre-LN 块 + 因果掩码），在第 5 讲的玩具语言上训练，
然后和第 5 讲的 RNN 语言模型做同样的测试：
演示 1：模型结构与参数量
演示 2：训练（AdamW + 学习率预热 + 余弦下降），和 RNN 比较测试集困惑度
演示 3：主谓一致 —— 训练时见过的距离 vs 比训练时更长的句子
（没训练过的位置嵌入会不会让模型出错？）
演示 4：注意力在看哪里？—— 预测主动词时，有没有哪个头专门"盯着"主语？
演示 5：从迷你 GPT 采样
本文件导入第 5 讲的 toy_language.py 和 04_rnn_lm.py（请保持 code/ 下的目录结构），
并加载第 5 讲保存的 RNN 模型 lecture05/checkpoints/rnn_lm.pt（不存在时会先训练一个，约 1 分钟）。
-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
python 04_mini_gpt.py # 完整运行（训练 8 轮，保存到 checkpoints/mini_gpt.pt）
python 04_mini_gpt.py --quick # 冒烟测试（少量数据、1 轮，不保存）
需要 PyTorch。完整运行约 2~4 分钟（CPU）。
"""
import os
os.environ.setdefault("OMP_NUM_THREADS", "1")
import importlib
import math
import sys
import time
from pathlib import Path
import numpy as np
# 【Python 小知识】把第 5 讲的文件夹加进搜索路径，才能 import 那里的模块
LECTURE05 = Path(__file__).resolve().parent.parent / "lecture05"
if str(LECTURE05) not in sys.path:
sys.path.append(str(LECTURE05))
import torch # noqa: E402
import torch.nn.functional as F # noqa: E402
import toy_language # noqa: E402
from toy_language import Vocab, agreement_test_set, generate_corpus # noqa: E402
from transformer import TransformerLM, count_parameters, lm_loss, warmup_cosine # noqa: E402
rnnlm = importlib.import_module("04_rnn_lm") # 第 5 讲的 RNN 语言模型（数据分批、困惑度等函数也复用）
torch.set_num_threads(1) # 小模型单线程最快（见第 6 讲 04 号文件的说明）
CKPT_DIR = Path(__file__).parent / "checkpoints"
MAX_LEN = 64 # 位置嵌入表的长度：足够放下测试用的最长句子（8 个修饰语时约 36 个词）
NOUN_FORMS = {w for pair in toy_language.NOUNS + toy_language.PLACES for w in pair}
def print_header(title):
"""打印一个带分隔线的标题。"""
print()
print("=" * 70)
print(f" {title}")
print("=" * 70)
# =============================================================================
# 第 1 部分：训练与评估
# =============================================================================
def train_gpt(train, vocab, pos="learned", epochs=8, lr=3e-3, batch_size=64, d_model=64, n_heads=4,
n_layers=2, warmup=200, weight_decay=0.01, seed=0, dev=None, verbose=True):
"""
训练一个迷你 GPT。
【参数】
train, vocab : 训练句子（list[list[str]]）和词表
pos : 位置信息的做法，"learned" / "sinusoidal" / "rope" / "none"（05 号文件比较它们）
epochs, lr, batch_size : 训练轮数、最大学习率、批大小
d_model, n_heads, n_layers : 模型大小
warmup : 学习率预热的步数（讲义第 9 节）
weight_decay : AdamW 的权重衰减
seed : 随机种子（决定初始化和数据顺序）
dev : None 或开发集句子，每轮结束报告开发集困惑度
verbose : 是否打印每轮的信息
【返回】TransformerLM（训练好的模型，处于 train 模式）
【思路】
数据处理和第 5 讲完全一样（rnnlm.make_batches）：输入 = [, w1, ..., wT]，目标 = [w1, ..., wT, ]，
填充位置的目标是 0（），计算损失时忽略。
和 RNN 唯一的区别：RNN 一步一步地读，Transformer 一次读进整个句子，靠因果掩码保证每个位置只看前面。
【Python 小知识】
torch.optim.AdamW：Adam + "解耦的"权重衰减（Loshchilov & Hutter 2019），训练 Transformer 的标准选择。
opt.param_groups[0]["lr"] = ...：手动设置这一步的学习率（实现学习率调度）。
"""
torch.manual_seed(seed)
model = TransformerLM(len(vocab), d_model, n_heads, n_layers, max_len=MAX_LEN, pos=pos)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
steps_per_epoch = math.ceil(len(train) / batch_size)
total, step = epochs * steps_per_epoch, 0
warmup = min(warmup, total // 5) # 总步数很少时（例如 --quick）预热也相应缩短
for epoch in range(epochs):
t0, loss_sum, n_tok = time.perf_counter(), 0.0, 0
for X, Y in rnnlm.make_batches(train, vocab, batch_size, seed=seed + epoch):
opt.param_groups[0]["lr"] = warmup_cosine(step, total, warmup, lr)
logits, _ = model(X)
loss = lm_loss(logits, Y)
opt.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
step += 1
n = int((Y != 0).sum())
loss_sum += loss.item() * n
n_tok += n
if verbose:
msg = f" 第 {epoch + 1} 轮：训练集困惑度 {math.exp(loss_sum / n_tok):.2f}"
if dev is not None:
msg += f"，开发集困惑度 {perplexity(model, dev, vocab):.2f}"
print(msg + f"，用时 {time.perf_counter() - t0:.1f} 秒")
return model
@torch.no_grad()
def perplexity(model, sentences, vocab):
"""
困惑度 = exp(所有真实词位置的平均交叉熵)。
【边界】会临时切换到 eval 模式（关闭 dropout），结束后切回 train 模式。
"""
model.eval()
total, n_tok = 0.0, 0
for X, Y in rnnlm.make_batches(sentences, vocab, 256, shuffle=False):
logits, _ = model(X)
total += lm_loss(logits, Y, reduction="sum").item()
n_tok += int((Y != 0).sum())
model.train()
return math.exp(total / n_tok)
def encode_prefixes(vocab, prefixes):
"""
把一批前缀（词列表）编码成右侧填充的张量。
【返回】(X: (N, T) LongTensor, last: (N,) 每个前缀最后一个词的下标)
【边界】因为用了因果掩码，右侧的填充不会影响前面真实位置的输出，所以不需要传 pad_mask。
"""
seqs = [vocab.encode(p, add_eos=False) for p in prefixes]
X = torch.zeros(len(seqs), max(len(s) for s in seqs), dtype=torch.long)
for b, s in enumerate(seqs):
X[b, :len(s)] = torch.tensor(s)
return X, torch.tensor([len(s) - 1 for s in seqs])
@torch.no_grad()
def agreement_accuracy(model, vocab, tests):
"""
主谓一致准确率：在主动词的位置上，正确形式的概率高于错误形式的比例（和第 5 讲的定义相同）。
【参数】tests：toy_language.agreement_test_set 的输出
【返回】float，0~1
"""
model.eval()
X, last = encode_prefixes(vocab, [t["prefix"] for t in tests])
logits, _ = model(X)
logp = F.log_softmax(logits[torch.arange(len(tests)), last], dim=-1)
model.train()
right = [logp[i, vocab.stoi[t["correct"]]] > logp[i, vocab.stoi[t["wrong"]]] for i, t in enumerate(tests)]
return float(np.mean(right))
@torch.no_grad()
def attention_to_subject(model, vocab, tests):
"""
统计：模型在"预测主动词"的那个位置（前缀的最后一个词），每一层、每个头的注意力有多少落在主语名词上。
【返回】(subj, attractor, uniform)，前两个是 (层数, 头数) 的数组：
subj[l, h] : 第 l 层第 h 个头分给主语名词的平均权重
attractor[l, h] : 分给修饰语里**单复数和主语不同**的名词（干扰名词）的平均权重（每个干扰名词）
uniform : 如果均匀分配，每个位置平均能分到的权重（对照基线）
【思路】
主语名词在前缀中的下标 = len(prefix) - distance - 1；编码后前面多了一个 ，所以再加 1。
干扰名词：主语之后、单复数和主语不同的名词（用"正确的动词是不是单数形式"判断主语的单复数）。
"""
model.eval()
X, last = encode_prefixes(vocab, [t["prefix"] for t in tests])
_, _, attns = model(X, return_attn=True)
model.train()
L, H = len(attns), attns[0].shape[1]
singular_verbs = {s for s, _ in toy_language.INTRANSITIVE} | {"is"}
plural_nouns = {p for _, p in toy_language.NOUNS + toy_language.PLACES}
subj, attr, uni = np.zeros((L, H)), np.zeros((L, H)), []
n_attr = 0
for i, t in enumerate(tests):
q = int(last[i])
s_pos = len(t["prefix"]) - t["distance"] - 1 + 1
subj_plural = t["correct"] not in singular_verbs
attr_pos = [j + 1 for j, w in enumerate(t["prefix"]) if j + 1 > s_pos and w in NOUN_FORMS
and (w in plural_nouns) != subj_plural]
uni.append(1.0 / (q + 1))
for l in range(L):
subj[l] += attns[l][i, :, q, s_pos].numpy()
for j in attr_pos:
attr[l] += attns[l][i, :, q, j].numpy()
n_attr += len(attr_pos)
return subj / len(tests), attr / max(1, n_attr), float(np.mean(uni))
@torch.no_grad()
def show_attention(model, vocab, prefix, layer, head):
"""把某一层某个头在前缀最后一个位置的注意力画成横向条形图。"""
model.eval()
X, last = encode_prefixes(vocab, [prefix])
_, _, attns = model(X, return_attn=True)
model.train()
w = attns[layer][0, head, int(last[0])]
for j, word in enumerate([Vocab.BOS] + prefix):
print(f" {word:>9s} {w[j]:.2f} " + "█" * int(round(float(w[j]) * 40)))
@torch.no_grad()
def sample(model, vocab, rng, temperature=1.0, max_len=30):
"""
从 开始逐词采样。每一步都把**整个前缀**重新送进模型（简单但重复计算；
真实系统会缓存之前位置的键和值，叫 KV 缓存，见 CS336 第 10 讲）。
【返回】词列表（不含 、）
"""
model.eval()
ids = [vocab.stoi[Vocab.BOS]]
banned = [vocab.stoi[Vocab.PAD], vocab.stoi[Vocab.BOS], vocab.stoi[Vocab.UNK]]
for _ in range(min(max_len, MAX_LEN - 1)):
logits, _ = model(torch.tensor([ids]))
z = logits[0, -1].clone()
z[banned] = -float("inf")
p = torch.softmax(z / temperature, dim=-1).numpy().astype(float)
w = int(rng.choice(len(p), p=p / p.sum()))
if w == vocab.stoi[Vocab.EOS]:
break
ids.append(w)
model.train()
return [vocab.itos[i] for i in ids[1:]]
# =============================================================================
# 演示
# =============================================================================
def main(quick=False):
n_train, epochs = (2000, 1) if quick else (20000, 8)
train = generate_corpus(n_train, seed=0)
dev = generate_corpus(2000, seed=1)
test = generate_corpus(2000, seed=2)
vocab = Vocab(generate_corpus(20000, seed=0)) # 和第 5 讲完全相同的词表
# ---------------- 演示 1 ----------------
print_header("演示 1：模型结构与参数量")
torch.manual_seed(0)
gpt = TransformerLM(len(vocab), d_model=64, n_heads=4, n_layers=2, max_len=MAX_LEN)
print(f"\n 词表 {len(vocab)} 个词；d_model = 64，4 个头（每个 16 维），2 层，位置嵌入表 {MAX_LEN} 行")
print(f" 迷你 GPT 的参数：{count_parameters(gpt):,} 个")
print(f" 其中词嵌入 {len(vocab) * 64:,}，位置嵌入 {MAX_LEN * 64:,}，每个块 {12 * 64 * 64 + 13 * 64:,}")
rnn_tmp = rnnlm.RNNLM(len(vocab))
print(f" 第 5 讲的 RNN（词向量 32 维、隐藏状态 128 维）：{count_parameters(rnn_tmp):,} 个")
# ---------------- 演示 2 ----------------
print_header(f"演示 2：训练迷你 GPT（{len(train)} 句，{epochs} 轮）")
t0 = time.perf_counter()
gpt = train_gpt(train, vocab, epochs=epochs, dev=dev)
print(f"\n 训练总用时 {time.perf_counter() - t0:.1f} 秒（取决于你的电脑）")
if quick:
rnn = rnnlm.train_rnnlm(train, vocab, epochs=1, verbose=False)
else:
rnn = rnnlm.load_or_train(vocab, generate_corpus(20000, seed=0)) # 第 5 讲的模型
print(f" 测试集困惑度：RNN（第 5 讲）{rnnlm.perplexity(rnn, test, vocab):.2f}，迷你 GPT {perplexity(gpt, test, vocab):.2f}")
if not quick:
CKPT_DIR.mkdir(exist_ok=True)
torch.save(gpt.state_dict(), CKPT_DIR / "mini_gpt.pt")
print(" （模型已保存到 checkpoints/mini_gpt.pt）")
print("""
观察：两者的困惑度很接近 —— 这门玩具语言的大部分词只依赖近处的上下文，
两种模型都已经接近这门语言本身的"不确定性下限"了（第 5 讲第 6.5 节）。差别要看专门的测试。""")
# ---------------- 演示 3 ----------------
print_header("演示 3：主谓一致 —— 训练语料里最多 3 个修饰语（句子最长约 20 个词）")
print(f"\n {'修饰语':<6s}{'平均距离':>8s}{'最长前缀':>8s}{'RNN':>9s}{'迷你 GPT':>10s}")
for k in [0, 1, 2, 3, 4, 6, 8]:
tests = agreement_test_set(k, n=400)
longest = max(len(t["prefix"]) for t in tests) + 1
print(f" {k:<8d}{np.mean([t['distance'] for t in tests]):>8.1f}{longest:>10d}"
f"{rnnlm.agreement_accuracy(rnn, vocab, tests):>9.1%}{agreement_accuracy(gpt, vocab, tests):>10.1%}"
+ (" ← 比训练时更长" if k >= 4 else ""))
print("""
观察（具体数字以你运行的结果为准）：
- 在训练时见过的长度上（0~3 个修饰语），迷你 GPT 和 RNN 一样全对：
自注意力可以**直接**看到任意远的主语，不需要一步一步地把信息传过来。
- 在比训练时更长的句子上，迷你 GPT 也几乎全对，只比 RNN 略低一点。这有点出乎意料：
位置嵌入表第 20 行以后的向量在训练中从没出现过，还是初始化时很小的随机值。
模型找主语主要靠的是"内容"（哪个词是名词、是单数还是复数），而不是精确的位置；
没训练过的位置向量很小，在那些位置上几乎相当于"没有位置信息"，没有把模型带偏。
- 05 号文件比较四种位置信息的做法 —— 并不是每种做法都这么幸运。""")
# ---------------- 演示 4 ----------------
print_header("演示 4：注意力在看哪里？—— 预测主动词时，每个头分给主语的注意力")
tests = agreement_test_set(2, n=1000, seed=7)
pp_only = [t for t in tests if "that" not in t["prefix"]] # 两个修饰语都是介词短语的句子
subj, attr, uni = attention_to_subject(gpt, vocab, tests)
print(f"\n （2 个修饰语的句子，共 {len(tests)} 句；如果均匀分配，每个位置平均分到 {uni:.2f}）")
print(f" {'':<10s}" + "".join(f"{'头' + str(h):>16s}" for h in range(subj.shape[1])))
for l in range(subj.shape[0]):
print(f" 第 {l + 1} 层 " + "".join(f"{f'主语 {subj[l, h]:.2f}/干扰 {attr[l, h]:.2f}':>18s}" for h in range(subj.shape[1])))
subj_pp, _, _ = attention_to_subject(gpt, vocab, pp_only)
l1, h1 = np.unravel_index(np.argmax(subj_pp), subj_pp.shape)
ex1 = next(t for t in pp_only if t["attractors"] > 0)
print(f"\n 例句 1（只有介词短语，{len(pp_only)} 句里分给主语最多的是第 {l1 + 1} 层的头 {h1}，平均 {subj_pp[l1, h1]:.2f}）。"
f"\n 它在最后一个位置（要预测下一个词，也就是主动词）的注意力：\n")
show_attention(gpt, vocab, ex1["prefix"], l1, h1)
print(f"\n 正确的动词：{ex1['correct']}，错误的形式：{ex1['wrong']}")
transitive = {w for pair in toy_language.TRANSITIVE for w in pair}
ex2 = next(t for t in tests if any(w == "that" and nxt in transitive for w, nxt in zip(t["prefix"], t["prefix"][1:])))
l2, h2 = np.unravel_index(np.argmax(subj), subj.shape)
print(f"\n 例句 2（含主语关系从句 \"that + 动词\"），全部句子里分给主语最多的第 {l2 + 1} 层头 {h2}：\n")
show_attention(gpt, vocab, ex2["prefix"], l2, h2)
print(f"\n 正确的动词：{ex2['correct']}，错误的形式：{ex2['wrong']}")
print("""
观察（哪一层、哪个头负责什么，每次训练都可能不同，具体以你运行的结果为准）：
- 有的头在预测主动词时，把大部分注意力直接放在**主语名词**上，几乎不看中间的干扰名词 ——
模型自己学会了"找主语"这种句法关系，没有人告诉它什么是主语。
- 主语关系从句（"that sees the boys"）里的动词和主语的单复数一致，它的形式同样"泄露"了答案。
所以有的头会去看这些动词，而不是主语本身：模型会利用任何能用的线索。
- 有的头还会把一部分注意力停在 和句首的 the 上，这些位置并不包含主语单复数的信息。
很多模型里都能看到这种现象：头"没事可做"时把注意力停在开头的位置（attention sink）。
- 第 2 层的头看到的不再是原始的词，而是第 1 层已经"加工"过的表示：
第 1 层可能已经把主语的信息搬到了后面的位置上，第 2 层再从那些位置读取 —— 注意力可以**分几步**传递信息。
（Clark 等 2019 在 BERT 里也找到了专门关注"动词的宾语""名词的限定词"等句法关系的头。
注意：注意力权重大不等于这个位置"决定了"输出，第 6 讲第 8.5 节和第 17 讲会讨论这个问题。）""")
# ---------------- 演示 5 ----------------
print_header("演示 5：从迷你 GPT 采样（温度 1.0）")
rng = np.random.default_rng(0)
for _ in range(6):
print(" " + " ".join(sample(gpt, vocab, rng)))
if __name__ == "__main__":
main(quick="--quick" in sys.argv)