"""
=============================================================================
第 5 讲 · 代码 4：用 PyTorch 训练 RNN 语言模型 —— 困惑度、主谓一致、隐藏状态里有什么
=============================================================================
对应讲义：lectures/lecture05_language_models_rnn.md 第 5 ~ 8 节
-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
代码 03 用 numpy 手写了 RNN，一次只能处理一个句子，很慢。这里用 PyTorch 的 nn.RNN 批量训练，
然后回答本讲最关心的问题：**RNN 真的能利用窗口之外的信息吗？**
演示 1：nn.RNN 计算的就是 h_t = tanh(W_ih x_t + b_ih + W_hh h_{t-1} + b_hh)（手算一步验证）
演示 2：批量训练（填充 + teacher forcing + 梯度裁剪），和 4-gram 比较困惑度
演示 3：主谓一致 —— 按修饰语个数、按干扰名词个数统计准确率，和 4-gram 对比
演示 4：探针 —— 在动词之前的隐藏状态里，能不能线性地读出主语是单数还是复数？
演示 5：从 RNN 语言模型采样生成句子
训练好的模型会保存到 checkpoints/rnn_lm.pt，供 05 号文件和第 6 讲使用。
-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
python 04_rnn_lm.py # 完整运行（训练 8 轮）
python 04_rnn_lm.py --quick # 冒烟测试（2000 句、1 轮，不保存模型）
需要 PyTorch。完整运行约 2~5 分钟（CPU；有 GPU 会更快，但这个小模型用 CPU 就够了）。
"""
import importlib
import math
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from toy_language import Vocab, agreement_test_set, generate_corpus
CKPT_DIR = Path(__file__).parent / "checkpoints"
def print_header(title):
"""打印一个带分隔线的标题。"""
print()
print("=" * 70)
print(f" {title}")
print("=" * 70)
# =============================================================================
# 第 1 部分：模型与数据
# =============================================================================
class RNNLM(nn.Module):
"""
RNN 语言模型：嵌入 → 循环层 → 线性层（输出词表上的 logits）。
【参数】（构造函数）
vocab_size : 词表大小
embed_dim : 词向量维度
hidden : 隐藏状态维度
cell : "rnn"（tanh RNN）、"lstm" 或 "gru"（后两个是第 6 讲的内容，这里先留好接口）
num_layers : 堆叠几层
【Python 小知识】
nn.RNN(..., batch_first=True)：输入形状是 (批大小, 序列长度, 特征维度)。
它返回 (所有时间步的输出, 最后的隐藏状态)；对单层 RNN，"输出"就是每个时间步的 h_t。
"""
def __init__(self, vocab_size, embed_dim=32, hidden=128, cell="rnn", num_layers=1):
super().__init__()
self.cell = cell
self.emb = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
rnn_cls = {"rnn": nn.RNN, "lstm": nn.LSTM, "gru": nn.GRU}[cell]
self.rnn = rnn_cls(embed_dim, hidden, num_layers=num_layers, batch_first=True)
self.out = nn.Linear(hidden, vocab_size)
def forward(self, X, state=None):
"""
X: (B, T) 的词编号 → (logits (B, T, V), 隐藏状态序列 (B, T, H), 最后的状态)。
"""
h, state = self.rnn(self.emb(X), state)
return self.out(h), h, state
def make_batches(sentences, vocab, batch_size, shuffle=True, seed=0):
"""
把句子分成批，每批填充到同样长度。
【思路】
每个句子编码成 [, w1, ..., wT, ]。输入是去掉最后一个的部分，目标是去掉第一个的部分：
输入： the dog runs .
目标：the dog runs .
这就是 **teacher forcing**：训练时，每一步的输入都是**真实的**上一个词，而不是模型自己预测的词。
句子长短不一，短的在末尾用 （编号 0）补齐；计算损失时忽略 位置。
【返回】list[(X, Y)]，X、Y 都是 torch.LongTensor (B, T_max)
"""
rng = np.random.default_rng(seed)
order = rng.permutation(len(sentences)) if shuffle else np.arange(len(sentences))
batches = []
for start in range(0, len(order), batch_size):
seqs = [vocab.encode(sentences[i]) for i in order[start:start + batch_size]]
T = max(len(s) for s in seqs) - 1
X = torch.zeros(len(seqs), T, dtype=torch.long)
Y = torch.zeros(len(seqs), T, dtype=torch.long)
for b, s in enumerate(seqs):
X[b, :len(s) - 1] = torch.tensor(s[:-1])
Y[b, :len(s) - 1] = torch.tensor(s[1:])
batches.append((X, Y))
return batches
def train_rnnlm(train, vocab, cell="rnn", epochs=8, lr=2e-3, batch_size=64, clip=1.0, hidden=128, seed=0,
verbose=True, dev=None):
"""
训练 RNN 语言模型。
【参数】
train, vocab : 训练句子、词表
cell : "rnn" / "lstm" / "gru"
epochs, lr, batch_size : 训练超参数
clip : 梯度裁剪的阈值（梯度的总范数超过它就按比例缩小；第 6 讲详细讲）
hidden : 隐藏状态维度
seed : 随机种子
verbose : 是否打印每轮的损失
dev : None 或开发集句子，每轮结束报告开发集困惑度
【返回】RNNLM
"""
torch.manual_seed(seed)
model = RNNLM(len(vocab), hidden=hidden, cell=cell)
opt = torch.optim.Adam(model.parameters(), lr=lr)
for epoch in range(epochs):
t0, total, n_tok = time.perf_counter(), 0.0, 0
for X, Y in make_batches(train, vocab, batch_size, seed=seed + epoch):
opt.zero_grad()
logits, _, _ = model(X)
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y.reshape(-1), ignore_index=0)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
opt.step()
n = int((Y != 0).sum())
total += loss.item() * n
n_tok += n
if verbose:
msg = f" 第 {epoch + 1} 轮：训练集平均交叉熵 {total / n_tok:.4f}（困惑度 {math.exp(total / 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(所有非填充位置的平均交叉熵)。"""
model.eval()
total, n_tok = 0.0, 0
for X, Y in make_batches(sentences, vocab, 256, shuffle=False):
logits, _, _ = model(X)
total += F.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y.reshape(-1), ignore_index=0,
reduction="sum").item()
n_tok += int((Y != 0).sum())
model.train()
return math.exp(total / n_tok)
@torch.no_grad()
def prefix_states(model, vocab, prefixes):
"""
把一批前缀（词列表）送进模型，返回每个前缀最后一个位置的 (下一个词的对数概率, 隐藏状态)。
【返回】(logp: (N, V), h: (N, H))
"""
model.eval()
seqs = [vocab.encode(p, add_eos=False) for p in prefixes]
T = max(len(s) for s in seqs)
X = torch.zeros(len(seqs), T, dtype=torch.long)
for b, s in enumerate(seqs):
X[b, :len(s)] = torch.tensor(s)
logits, h, _ = model(X)
last = torch.tensor([len(s) - 1 for s in seqs])
idx = torch.arange(len(seqs))
model.train()
return F.log_softmax(logits[idx, last], dim=-1), h[idx, last]
def agreement_accuracy(model, vocab, tests):
"""主谓一致准确率：正确动词形式的概率高于错误形式的比例。"""
logp, _ = prefix_states(model, vocab, [t["prefix"] for t in tests])
right = sum(float(logp[i, vocab.stoi[t["correct"]]] > logp[i, vocab.stoi[t["wrong"]]]) for i, t in enumerate(tests))
return right / len(tests)
@torch.no_grad()
def sample(model, vocab, rng, temperature=1.0, max_len=25):
"""从 开始逐词采样（温度见 05 号文件）。返回词列表。"""
model.eval()
x = torch.tensor([[vocab.stoi[Vocab.BOS]]])
state, out = None, []
for _ in range(max_len):
logits, _, state = model(x, state)
z = logits[0, -1].clone()
z[[vocab.stoi[Vocab.PAD], vocab.stoi[Vocab.BOS], vocab.stoi[Vocab.UNK]]] = -float("inf") # 特殊符号不参与采样
p = torch.softmax(z / temperature, dim=-1).numpy().astype(float)
w = int(rng.choice(len(p), p=p / p.sum()))
if vocab.itos[w] == Vocab.EOS:
break
out.append(vocab.itos[w])
x = torch.tensor([[w]])
model.train()
return out
def load_or_train(vocab, train, cell="rnn", path=None, **kw):
"""如果检查点存在就加载，否则训练并保存。供 05 号文件和第 6 讲使用。"""
path = Path(path) if path else CKPT_DIR / f"{cell}_lm.pt"
model = RNNLM(len(vocab), cell=cell, hidden=kw.get("hidden", 128))
if path.exists():
model.load_state_dict(torch.load(path))
return model
model = train_rnnlm(train, vocab, cell=cell, **kw)
path.parent.mkdir(exist_ok=True)
torch.save(model.state_dict(), path)
return model
# =============================================================================
# 演示
# =============================================================================
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)) # 词表固定用完整训练集建立
# ---------------- 演示 1：nn.RNN 在算什么 ----------------
print_header("演示 1：nn.RNN 在算什么？")
torch.manual_seed(0)
rnn = nn.RNN(3, 4, batch_first=True)
x = torch.randn(1, 2, 3)
out, _ = rnn(x)
h0 = torch.zeros(4)
h1 = torch.tanh(rnn.weight_ih_l0 @ x[0, 0] + rnn.bias_ih_l0 + rnn.weight_hh_l0 @ h0 + rnn.bias_hh_l0)
h2 = torch.tanh(rnn.weight_ih_l0 @ x[0, 1] + rnn.bias_ih_l0 + rnn.weight_hh_l0 @ h1 + rnn.bias_hh_l0)
print(f"\n 手算 h_1 = {h1.detach().numpy().round(4)}")
print(f" nn.RNN = {out[0, 0].detach().numpy().round(4)}")
print(f" 手算 h_2 与 nn.RNN 的第 2 步一致？{torch.allclose(h2, out[0, 1])}")
print(" （PyTorch 用了两个偏置 b_ih 和 b_hh，加起来就是讲义公式里的一个 b。）")
# ---------------- 演示 2：训练 ----------------
print_header(f"演示 2：训练 RNN 语言模型（{len(train)} 句，{epochs} 轮，隐藏状态 128 维）")
t0 = time.perf_counter()
model = train_rnnlm(train, vocab, epochs=epochs, dev=dev)
print(f"\n 训练总用时 {time.perf_counter() - t0:.1f} 秒（取决于你的电脑）")
ngram = importlib.import_module("01_ngram_lm")
lm4 = ngram.NGramLM(n=4).fit(train)
lambdas, _ = ngram.tune_lambdas(lm4, dev)
print(f" 测试集困惑度：4-gram（插值）{lm4.perplexity(test, 'interp', lambdas):.2f}，"
f"RNN {perplexity(model, test, vocab):.2f}")
print(f" RNN 的参数量：{sum(p.numel() for p in model.parameters()):,}")
if not quick:
CKPT_DIR.mkdir(exist_ok=True)
torch.save(model.state_dict(), CKPT_DIR / "rnn_lm.pt")
print(f" （模型已保存到 checkpoints/rnn_lm.pt）")
# ---------------- 演示 3：主谓一致 ----------------
print_header("演示 3：主谓一致 —— RNN vs 4-gram")
def ngram_score(prefix, word):
return math.log(lm4.prob_interp(word, [Vocab.BOS] * 3 + prefix, lambdas))
print(f"\n (a) 用完整训练集（0~3 个修饰语）训练的 RNN：")
print(f" {'修饰语':<6s}{'平均距离':>8s}{'4-gram':>9s}{'RNN':>8s} RNN：{'无干扰名词':>10s}{'有干扰名词':>10s}")
for k in [0, 1, 2, 3, 4, 6, 8]:
tests = agreement_test_set(k, n=400)
a_ng = ngram.agreement_accuracy(ngram_score, tests)[0]
a_rnn = agreement_accuracy(model, vocab, tests)
no_att = [t for t in tests if t["attractors"] == 0]
att = [t for t in tests if t["attractors"] > 0]
a0 = agreement_accuracy(model, vocab, no_att) if no_att else float("nan")
a1 = agreement_accuracy(model, vocab, att) if att else float("nan")
dist = np.mean([t["distance"] for t in tests])
print(f" {k:<8d}{dist:>8.1f}{a_ng:>9.1%}{a_rnn:>8.1%} {a0:>14.1%}{a1:>10.1%}" + (" ← 比训练时更长" if k >= 4 else ""))
print("""
观察：主语跑出 n-gram 的窗口之后，4-gram 只剩六成左右；RNN 却几乎全对，
甚至在比训练时长得多的句子上（8 个修饰语，距离约 30 个词）也是如此。
RNN 确实把"主语是单数还是复数"这一信息带过了中间的所有修饰语 —— 这是固定窗口的模型做不到的。""")
short_train = generate_corpus(n_train, seed=0, modifier_probs=(0.6, 0.4))
model_short = train_rnnlm(short_train, vocab, epochs=max(1, epochs - 2), verbose=False)
print(f"\n (b) 如果训练语料里**只有 0~1 个修饰语**（最远距离约 5 个词），RNN 能推广到更远的距离吗？")
print(f" {'修饰语':<6s}{'平均距离':>8s}{'RNN（只见过短依赖）':>20s}")
for k in [1, 2, 3, 4, 6, 8]:
tests = agreement_test_set(k, n=400)
print(f" {k:<8d}{np.mean([t['distance'] for t in tests]):>8.1f}{agreement_accuracy(model_short, vocab, tests):>20.1%}")
print("""
观察：训练时只见过短依赖，RNN 在稍长的距离上仍然不错，但距离越远准确率越低。
它学到的"记住主语单复数"的机制并不稳健 —— 信息每传一步就损耗一点（代码 03 演示 4 的梯度消失是同一件事的另一面）。
第 6 讲会用同样的实验比较 LSTM（它几乎不受影响）。
（Linzen 等人 2016 年在真实英语上发现，即使是 LSTM，在有干扰名词时也会明显犯错 ——
真实语言比我们的玩具语言复杂得多。具体数字以你运行的结果为准。）""")
# ---------------- 演示 4：探针 ----------------
print_header("演示 4：探针 —— 隐藏状态里有没有'主语是单数还是复数'的信息？")
print("""
做法：取主动词之前那个位置的隐藏状态 h（128 维），训练一个逻辑回归（线性分类器）预测主语的单复数。
一半题目训练，另一半测试。如果线性分类器就能读出来，说明这个信息被**线性地编码**在隐藏状态里。
""")
print(f" {'修饰语':<6s}{'完整训练的 RNN':>16s}{'只见过短依赖的 RNN':>20s}")
singular_forms = {f for f, _ in __import__("toy_language").INTRANSITIVE} | {"is"}
for k in [0, 1, 2, 4, 8]:
tests = agreement_test_set(k, n=400, seed=7)
y = torch.tensor([float(t["correct"] in singular_forms) for t in tests])
accs = []
for m in (model, model_short):
_, H = prefix_states(m, vocab, [t["prefix"] for t in tests])
probe = nn.Linear(H.shape[1], 1)
opt = torch.optim.Adam(probe.parameters(), lr=0.05)
for _ in range(300):
opt.zero_grad()
F.binary_cross_entropy_with_logits(probe(H[:200]).squeeze(1), y[:200]).backward()
opt.step()
accs.append(((probe(H[200:]).squeeze(1) > 0).float() == y[200:]).float().mean().item())
print(f" {k:<8d}{accs[0]:>16.1%}{accs[1]:>20.1%}")
print("""
观察：完整训练的 RNN 的隐藏状态里，主语的单复数信息无论隔多远都能被线性地读出来；
只见过短依赖的 RNN，离得越远越难读出 —— 信息在一步步传递中逐渐"模糊"了。
探针（probing）是研究"模型内部学到了什么"的常用工具，第 17 讲会再讨论。""")
# ---------------- 演示 5：生成 ----------------
print_header("演示 5：从 RNN 语言模型采样")
rng = np.random.default_rng(0)
for _ in range(6):
print(" " + " ".join(sample(model, vocab, rng)))
if __name__ == "__main__":
main(quick="--quick" in sys.argv)