"""
=============================================================================
第 5 讲 · 代码 2：固定窗口的神经语言模型（Bengio 等，2003）
=============================================================================
对应讲义：lectures/lecture05_language_models_rnn.md 第 4 节
-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
n-gram 模型的两个问题：稀疏（没见过的 n-gram 概率为 0）、存储（要存海量计数）。
Bengio 等人（2003）提出用神经网络来做语言模型，思路和第 3 讲的窗口分类器一模一样：
取前面 n-1 个词的词向量，拼接起来 → 隐藏层 → softmax，输出下一个词在整个词表上的概率分布。
x = [e(w_{t-4}); e(w_{t-3}); e(w_{t-2}); e(w_{t-1})] （窗口大小 4）
h = tanh(W x + b)
P(w_t | 前 4 个词) = softmax(U h + c)
好处：没有稀疏问题（任何窗口都能算出概率），不用存计数；相似的词有相似的向量，
"the dog runs" 学到的东西能自动迁移到 "the cat runs" 上。
问题：窗口仍然是固定的 —— 窗口外的信息还是看不到；窗口越大，W 越大，而且每个位置有自己的一组权重。
演示 1：训练窗口大小为 4 的神经语言模型，和 4-gram 比较困惑度
演示 2：学到的词向量 —— 相似的词靠在一起
演示 3：主谓一致 —— 和 4-gram 一样，主语出了窗口就看不到
演示 4：把窗口开大会怎样？参数量随窗口线性增长
-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
python 02_fixed_window_lm.py # 完整运行
python 02_fixed_window_lm.py --quick # 冒烟测试（少量数据、1 轮）
需要 PyTorch。完整运行约 1~3 分钟（CPU）。
"""
import importlib
import math
import sys
import time
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
ngram = importlib.import_module("01_ngram_lm")
def print_header(title):
"""打印一个带分隔线的标题。"""
print()
print("=" * 70)
print(f" {title}")
print("=" * 70)
def make_windows(sentences, vocab, window):
"""
把句子变成 (前 window 个词, 下一个词) 的训练样本。
【思路】
句子编码成 [, w1, ..., wT, ]，对每个要预测的位置 t（从 w1 到 ），
取它前面的 window 个编号作为输入；不够的部分在左边用 填充。
【参数】
sentences : list[list[str]]
vocab : Vocab
window : int，窗口大小（n-gram 中的 n-1）
【返回】(X, y)：X 是 torch.LongTensor (N, window)，y 是 (N,)
"""
bos = vocab.stoi[Vocab.BOS]
X, y = [], []
for s in sentences:
ids = vocab.encode(s)
padded = [bos] * (window - 1) + ids
for t in range(1, len(ids)):
X.append(padded[t - 1:t - 1 + window])
y.append(ids[t])
return torch.tensor(X, dtype=torch.long), torch.tensor(y, dtype=torch.long)
class FixedWindowLM(nn.Module):
"""
固定窗口的神经语言模型：嵌入 → 拼接 → tanh 隐藏层 → 词表大小的 softmax。
【参数】（构造函数）
vocab_size : 词表大小
window : 窗口大小
embed_dim : 词向量维度
hidden : 隐藏单元数
【边界】
- 隐藏层权重 W 的形状是 (hidden, window × embed_dim)：窗口每大一格，参数就多 hidden × embed_dim 个；
而且 W 中对应不同位置的那几块是**互不共享**的 —— 模型要分别学"前 1 个词是 dog"和"前 3 个词是 dog"。
"""
def __init__(self, vocab_size, window, embed_dim=32, hidden=128):
super().__init__()
self.emb = nn.Embedding(vocab_size, embed_dim)
self.hidden = nn.Linear(window * embed_dim, hidden)
self.out = nn.Linear(hidden, vocab_size)
def forward(self, X):
"""X: (B, window) → (B, vocab_size) 的 logits。"""
h = torch.tanh(self.hidden(self.emb(X).view(X.shape[0], -1)))
return self.out(h)
def train_lm(model, X, y, epochs=3, lr=2e-3, batch_size=256, seed=0):
"""用 Adam 训练，返回每轮的平均交叉熵。"""
torch.manual_seed(seed)
opt = torch.optim.Adam(model.parameters(), lr=lr)
hist = []
for _ in range(epochs):
perm = torch.randperm(len(X))
total = 0.0
for start in range(0, len(X), batch_size):
idx = perm[start:start + batch_size]
opt.zero_grad()
loss = F.cross_entropy(model(X[idx]), y[idx])
loss.backward()
opt.step()
total += loss.item() * len(idx)
hist.append(total / len(X))
return hist
@torch.no_grad()
def perplexity(model, X, y):
"""困惑度 = exp(平均交叉熵)。"""
model.eval()
ce = F.cross_entropy(model(X), y).item()
model.train()
return math.exp(ce)
@torch.no_grad()
def agreement(model, vocab, window, tests):
"""在主谓一致测试题上计算准确率：比较模型给正确 / 错误动词形式的对数概率。"""
model.eval()
bos = vocab.stoi[Vocab.BOS]
X = []
for t in tests:
ids = vocab.encode(t["prefix"], add_eos=False)
X.append(([bos] * window + ids)[-window:])
logp = F.log_softmax(model(torch.tensor(X)), dim=1)
right = sum(logp[i, vocab.stoi[t["correct"]]] > logp[i, vocab.stoi[t["wrong"]]] for i, t in enumerate(tests))
model.train()
return float(right) / len(tests)
def main(quick=False):
n_train = 2000 if quick else 20000
epochs = 1 if quick else 3
train = generate_corpus(n_train, seed=0)
dev = generate_corpus(2000, seed=1)
test = generate_corpus(2000, seed=2)
vocab = Vocab(train)
# ---------------- 演示 1：训练 ----------------
print_header("演示 1：窗口大小 4 的神经语言模型 vs n-gram（窗口 4 = 看前 4 个词，相当于 5-gram）")
window = 4
Xtr, ytr = make_windows(train, vocab, window)
Xte, yte = make_windows(test, vocab, window)
torch.manual_seed(0) # 固定初始化的随机性，保证每次运行结果相同
model = FixedWindowLM(len(vocab), window)
n_params = sum(p.numel() for p in model.parameters())
t0 = time.perf_counter()
hist = train_lm(model, Xtr, ytr, epochs=epochs)
print(f"\n 训练样本 {len(Xtr):,} 个，参数 {n_params:,} 个；每轮平均交叉熵 {[round(h, 4) for h in hist]}"
f"（用时 {time.perf_counter() - t0:.1f} 秒）")
ppl = {}
for n in [4, 5]:
lm = ngram.NGramLM(n=n).fit(train)
lambdas, _ = ngram.tune_lambdas(lm, dev)
ppl[n] = (lm.perplexity(test, "interp", lambdas), sum(len(c) for c in lm.counts.values()))
print(f" 测试集困惑度：4-gram（插值）{ppl[4][0]:.2f}，5-gram（插值）{ppl[5][0]:.2f}，"
f"神经模型（窗口 4）{perplexity(model, Xte, yte):.2f}")
print(f" 存储：5-gram 模型存了 1~5 阶共 {ppl[5][1]:,} 个计数；神经模型 {n_params:,} 个参数")
# ---------------- 演示 2：词向量 ----------------
print_header("演示 2：语言模型学到的词向量")
E = model.emb.weight.detach().numpy()
E = E / np.linalg.norm(E, axis=1, keepdims=True)
for q in ["dog", "dogs", "runs", "run", "near", "old"]:
sims = E @ E[vocab.stoi[q]]
cands = [(w, sims[i]) for i, w in enumerate(vocab.itos) if w != q and w.isalpha()]
top = sorted(cands, key=lambda x: -x[1])[:5]
print(f" {q:<6s}→ " + ", ".join(f"{w}({s:.2f})" for w, s in top))
print("""
观察：近邻里能看到一些"同类"的词（单数名词、复数名词、形容词各自有聚在一起的趋势），
但也混着不少不相关的词，余弦相似度也不高 —— 这个小模型只训练了 3 轮，而且在这门玩具语言里，
预测下一个词并不需要把每个词的语义都分得很细。在大语料上训练的神经语言模型，词向量会清晰得多
（第 1、2 讲的 word2vec 正是 Bengio 这一思路的简化版）。
词向量的相似性让模型能泛化：没见过的 "the old king sleeps" 也能从 "the old dog sleeps" 中借到知识。""")
# ---------------- 演示 3：主谓一致 ----------------
print_header("演示 3：主谓一致")
print(f"\n {'修饰语个数':<10s}{'神经模型（窗口 4）':>18s}")
for k in range(5):
print(f" {k:<12d}{agreement(model, vocab, window, agreement_test_set(k, n=400)):>16.1%}")
print("""
观察：1 个修饰语时（平均距离 3.7），主语还经常落在 4 个词的窗口里，所以准确率还不错；
修饰语再多，主语跑出了窗口，准确率就掉到六成多（随机猜是 50%）。
神经网络解决了**稀疏**问题，但没有解决**窗口有限**的问题。""")
# ---------------- 演示 4：窗口大小 ----------------
print_header("演示 4：把窗口开大")
print(f"\n {'窗口':>4s}{'参数量':>10s}{'困惑度':>8s}{'一致（1 个修饰语）':>18s}{'一致（2 个修饰语）':>18s}")
for w in [2, 4, 8, 12]:
Xw, yw = make_windows(train, vocab, w)
Xwt, ywt = make_windows(test, vocab, w)
torch.manual_seed(0)
m = FixedWindowLM(len(vocab), w)
train_lm(m, Xw, yw, epochs=epochs)
p = sum(q.numel() for q in m.parameters())
a1 = agreement(m, vocab, w, agreement_test_set(1, n=400))
a2 = agreement(m, vocab, w, agreement_test_set(2, n=400))
print(f" {w:>4d}{p:>10,d}{perplexity(m, Xwt, ywt):>8.2f}{a1:>18.1%}{a2:>18.1%}")
print("""
观察：窗口越大，能"看到"的主语越多，主谓一致越好；但参数量随窗口线性增长，
而且每个位置的权重是分开学的（"主语在前 5 个位置"和"在前 9 个位置"要分别学）。
无论开多大，总有更长的句子。我们需要一种能处理**任意长度**输入、并且在所有位置**共享参数**的结构 ——
这就是循环神经网络（代码 03、04）。""")
if __name__ == "__main__":
main(quick="--quick" in sys.argv)