"""
=============================================================================
第 5 讲 · 代码 1：n-gram 语言模型 —— 计数、稀疏、平滑、困惑度、长距离依赖
=============================================================================
对应讲义：lectures/lecture05_language_models_rnn.md 第 1 ~ 3、8 节
-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
语言模型给一个句子打分：P(w_1, ..., w_T)。用链式法则拆开：
P(w_1 ... w_T) = P(w_1) · P(w_2 | w_1) · P(w_3 | w_1 w_2) · ... · P(w_T | w_1 ... w_{T-1})
每一项都是"给定前文，预测下一个词"。问题是前文可以任意长，没法直接估计。
n-gram 模型做了一个简化（马尔可夫假设）：**只看前面 n-1 个词**。
P(w_t | w_1 ... w_{t-1}) ≈ P(w_t | w_{t-n+1} ... w_{t-1})
然后用**数数**来估计：P(w | 上下文) = count(上下文, w) / count(上下文)。
演示 1：数数和链式法则
演示 2：稀疏性 —— n 越大，测试集里没见过的 n-gram 越多，最大似然估计给它们概率 0
演示 3：平滑与插值，用困惑度（perplexity）评估
演示 4：从 n-gram 模型里采样生成句子
演示 5：长距离依赖 —— 主语和动词隔得远了，n-gram 还能选对动词形式吗？
演示 6：存储代价
-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
python 01_ngram_lm.py
只用 Python 标准库。运行时间约 20 秒。
"""
import math
import random
from collections import Counter
from toy_language import Vocab, agreement_test_set, generate_corpus
BOS, EOS = Vocab.BOS, Vocab.EOS
def print_header(title):
"""打印一个带分隔线的标题。"""
print()
print("=" * 70)
print(f" {title}")
print("=" * 70)
# =============================================================================
# 第 1 部分：n-gram 语言模型
# =============================================================================
class NGramLM:
"""
n-gram 语言模型：同时保存 1-gram 到 n-gram 的计数，支持加 k 平滑和线性插值。
【用法】
lm = NGramLM(n=3).fit(train_sentences)
lm.prob("runs", ("the", "dog")) # 加 k 平滑的 P(runs | the dog)
lm.prob_interp("runs", ("the", "dog"), lambdas) # 插值的概率
lm.perplexity(test_sentences, method="interp", lambdas=...)
【边界】
- 句子开头用 n-1 个 填充，结尾加一个 （这样"句子在这里结束"也是要预测的一个事件）。
- 训练时没见过的词按 处理（我们的玩具语言词表是封闭的，测试集不会有新词）。
"""
def __init__(self, n=3, k=0.01):
"""
【参数】
n : int >= 1，阶数（看前面 n-1 个词）
k : float > 0，加 k 平滑中的 k
"""
self.n, self.k = n, k
def pad(self, sentence):
"""在句子前面加 n-1 个 ，后面加 。"""
return [BOS] * (self.n - 1) + list(sentence) + [EOS]
def fit(self, sentences):
"""
数出所有阶数的 n-gram 次数。
【思路】
self.counts[m] 是 m-gram 的计数：键是 (上下文元组, 词)，上下文长度为 m-1。
self.context_counts[m] 是对应上下文出现的次数（作为分母）。
【返回】self
"""
self.vocab = sorted(set(w for s in sentences for w in s) | {EOS})
self.V = len(self.vocab)
self.counts = {m: Counter() for m in range(1, self.n + 1)}
self.context_counts = {m: Counter() for m in range(1, self.n + 1)}
for s in sentences:
toks = self.pad(s)
for t in range(self.n - 1, len(toks)):
for m in range(1, self.n + 1):
ctx = tuple(toks[t - m + 1:t])
self.counts[m][(ctx, toks[t])] += 1
self.context_counts[m][ctx] += 1
return self
def prob_mle(self, word, context, m):
"""
m 阶的最大似然估计 P(word | 最后 m-1 个上下文词) = count / count(上下文)。
【边界】
- 上下文从没出现过时，分母为 0，返回 0.0（这正是稀疏性问题之一）。
"""
ctx = tuple(context[len(context) - (m - 1):]) if m > 1 else ()
denom = self.context_counts[m][ctx]
return self.counts[m][(ctx, word)] / denom if denom else 0.0
def prob(self, word, context):
"""
加 k 平滑的 n 阶概率：(count(上下文, w) + k) / (count(上下文) + k · V)。
【思路】（讲义第 2.3 节）
给每个可能的词都"假装"多见过 k 次，没见过的 n-gram 也能得到一点概率。
k = 1 叫加一平滑（拉普拉斯平滑）；k 越大，概率分布越"平"（越接近均匀分布）。
"""
ctx = tuple(context[len(context) - (self.n - 1):]) if self.n > 1 else ()
return (self.counts[self.n][(ctx, word)] + self.k) / (self.context_counts[self.n][ctx] + self.k * self.V)
def prob_interp(self, word, context, lambdas):
"""
线性插值：P = λ_0 · (1/V) + λ_1 · P_1(w) + λ_2 · P_2(w | 1 个词) + ... + λ_n · P_n(w | n-1 个词)。
【思路】（讲义第 2.3 节）
高阶模型细致但稀疏（经常是 0），低阶模型粗糙但可靠。把它们按权重混合，
高阶模型有信息时就用它，没有信息（概率为 0）时由低阶模型兜底。
λ_0 · (1/V) 保证任何词的概率都大于 0。
【参数】
lambdas : list[float]，长度 n+1，非负且和为 1
【边界】
- λ 的取值应该在**开发集**上调，而不是在训练集上（训练集上高阶模型永远最好）。
"""
p = lambdas[0] / self.V
for m in range(1, self.n + 1):
p += lambdas[m] * self.prob_mle(word, context, m)
return p
def sentence_logprob(self, sentence, method="addk", lambdas=None):
"""
一个句子的对数概率之和（自然对数），以及被预测的词数（包括 ）。
【返回】(logprob, n_tokens)；如果某个词概率为 0，logprob 为 -inf
"""
toks = self.pad(sentence)
lp, cnt = 0.0, 0
for t in range(self.n - 1, len(toks)):
ctx = toks[:t]
if method == "addk":
p = self.prob(toks[t], ctx)
elif method == "mle":
p = self.prob_mle(toks[t], ctx, self.n)
else:
p = self.prob_interp(toks[t], ctx, lambdas)
lp += math.log(p) if p > 0 else -math.inf
cnt += 1
return lp, cnt
def perplexity(self, sentences, method="addk", lambdas=None):
"""
困惑度 = exp(-平均每个词的对数概率)。
【思路】（讲义第 3 节）
困惑度越低越好。直觉：困惑度为 10，相当于模型在每一步平均"在 10 个词里犹豫"。
如果任何一个词的概率是 0，困惑度就是无穷大。
"""
total_lp, total_n = 0.0, 0
for s in sentences:
lp, n = self.sentence_logprob(s, method, lambdas)
total_lp += lp
total_n += n
return math.exp(-total_lp / total_n) if total_lp > -math.inf else math.inf
def next_word_distribution(self, context, lambdas):
"""给定上下文，返回 (词列表, 插值概率列表)。"""
probs = [self.prob_interp(w, context, lambdas) for w in self.vocab]
return self.vocab, probs
def generate(self, rng, lambdas, max_len=30):
"""
从模型里一个词一个词地采样，直到生成 或达到最大长度。
【参数】rng : random.Random；lambdas : 插值权重；max_len : 最大长度
【返回】list[str]
"""
out = [BOS] * (self.n - 1)
for _ in range(max_len):
words, probs = self.next_word_distribution(out, lambdas)
w = rng.choices(words, weights=probs)[0]
if w == EOS:
break
out.append(w)
return out[self.n - 1:]
def tune_lambdas(lm, dev, iters=30):
"""
在开发集上用 EM 算法（期望最大化）估计插值权重 λ_0, ..., λ_n。
【思路】
把插值模型看成"每个词是由某一个子模型（均匀分布、1-gram、……、n-gram）生成的，
但我们不知道是哪一个"。EM 交替做两步：
E 步：对开发集里的每个词，算出它"由第 m 个子模型生成"的后验概率
r_m = λ_m · P_m(w) / Σ_j λ_j · P_j(w)；
M 步：新的 λ_m = 所有词的 r_m 的平均值。
每一轮都保证开发集的似然不下降，通常十几轮就收敛了。
（Jelinek 和 Mercer 在 1980 年代就是这样估计插值权重的。）
【参数】
lm : 训练好的 NGramLM
dev : 开发集句子
iters : EM 迭代次数
【返回】(lambdas, 开发集困惑度)
【边界】
- λ 必须在开发集（而不是训练集）上估计：在训练集上，最高阶的模型永远"最好"，λ_n 会趋于 1。
"""
rows = []
for sent in dev:
toks = lm.pad(sent)
for t in range(lm.n - 1, len(toks)):
rows.append([1.0 / lm.V] + [lm.prob_mle(toks[t], toks[:t], m) for m in range(1, lm.n + 1)])
lambdas = [1.0 / (lm.n + 1)] * (lm.n + 1)
for _ in range(iters):
acc = [0.0] * (lm.n + 1)
for r in rows:
z = sum(l * q for l, q in zip(lambdas, r))
for m in range(lm.n + 1):
acc[m] += lambdas[m] * r[m] / z
lambdas = [a / len(rows) for a in acc]
log_lik = sum(math.log(sum(l * q for l, q in zip(lambdas, r))) for r in rows)
return lambdas, math.exp(-log_lik / len(rows))
def agreement_accuracy(score_fn, tests):
"""
主谓一致准确率：score_fn(prefix, word) 返回 log P(word | prefix)，正确形式得分更高就算对。
【返回】(准确率, 题数)
"""
right = sum(score_fn(t["prefix"], t["correct"]) > score_fn(t["prefix"], t["wrong"]) for t in tests)
return right / len(tests), len(tests)
# =============================================================================
# 演示
# =============================================================================
def main():
train = generate_corpus(20000, seed=0)
dev = generate_corpus(2000, seed=1)
test = generate_corpus(2000, seed=2)
print(f"训练集 {len(train)} 句，开发集 {len(dev)} 句，测试集 {len(test)} 句（玩具语言，词表 {len(Vocab(train)) - 4} 个词）")
# ---------------- 演示 1：数数 ----------------
print_header("演示 1：数数，以及用链式法则计算句子的概率")
lm3 = NGramLM(n=3).fit(train)
ctx = ("the", "dog")
total = lm3.context_counts[3][ctx]
nxt = Counter({w: c for (c_, w), c in lm3.counts[3].items() if c_ == ctx})
print(f"\n 上下文 'the dog' 在训练集中出现 {total} 次，后面跟着的词：")
for w, c in nxt.most_common(6):
print(f" {w:<10s}{c:>5d} 次 P_MLE({w} | the dog) = {c}/{total} = {c / total:.3f}")
sent = ["the", "dog", "runs", "."]
toks = lm3.pad(sent)
print(f"\n P('the dog runs .') 按三元模型（最大似然）拆开：")
lp = 0.0
for t in range(2, len(toks)):
p = lm3.prob_mle(toks[t], toks[:t], 3)
lp += math.log(p)
print(f" P({toks[t]:<5s}| {toks[t - 2]} {toks[t - 1]}) = {p:.4f}")
print(f" 乘起来 = {math.exp(lp):.6f}")
# ---------------- 演示 2：稀疏性 ----------------
print_header("演示 2：稀疏性 —— 测试集中有多少 n-gram 在训练集里从没出现过？")
lm5 = NGramLM(n=5).fit(train)
print(f"\n {'n':>3s}{'不同 n-gram 的种数':>18s}{'测试集中没见过的比例':>20s}{'最大似然的困惑度':>16s}")
for m in range(1, 6):
unseen = total_ = 0
for s in test:
toks = [BOS] * 4 + s + [EOS]
for t in range(4, len(toks)):
ctx = tuple(toks[t - m + 1:t]) if m > 1 else ()
total_ += 1
unseen += lm5.counts[m][(ctx, toks[t])] == 0
ppl = NGramLM(n=m).fit(train).perplexity(test, "mle")
print(f" {m:>3d}{len(lm5.counts[m]):>18,d}{unseen / total_:>20.2%}{ppl:>16.2f}")
print("""
观察：n 越大，没见过的 n-gram 越多。最大似然估计给没见过的 n-gram 概率 0，
只要测试集里有一个这样的 n-gram，整个测试集的概率就是 0、困惑度就是无穷大（inf）。
我们的玩具语言很小、很规则，真实语言里这个问题严重得多。""")
# ---------------- 演示 3：平滑与插值 ----------------
print_header("演示 3：平滑与插值（开发集上选参数，测试集上报告困惑度）")
print(f"\n {'n':>3s}{'加 1':>10s}{'加 0.1':>10s}{'加 0.01':>10s}{'插值':>10s} 插值权重 λ_0..λ_n")
for m in [1, 2, 3, 4]:
row = []
for k in [1.0, 0.1, 0.01]:
row.append(NGramLM(n=m, k=k).fit(train).perplexity(test, "addk"))
lm = NGramLM(n=m).fit(train)
lambdas, _ = tune_lambdas(lm, dev)
row.append(lm.perplexity(test, "interp", lambdas))
print(f" {m:>3d}" + "".join(f"{v:>10.2f}" for v in row) + " " + " ".join(f"{x:.2f}" for x in lambdas))
print(f"\n （参考：如果模型对所有词一视同仁，困惑度等于词表大小 {lm3.V}。）")
print("""
观察：
- 加 1 平滑分给"没见过的事件"的概率太多了（每个上下文都给所有 V 个词各加 1），n 越大越糟；
k 取小一些好得多。
- 插值让高阶模型在有信息时发挥作用、没信息时由低阶模型兜底，是这几种方法里最好的。
- 从 1-gram 到 3-gram 困惑度大幅下降，再往上收益变小：更远的上下文能提供的信息变少，稀疏性却越来越严重。""")
# ---------------- 演示 4：生成 ----------------
print_header("演示 4：从 n-gram 模型采样生成句子")
rng = random.Random(0)
for m in [1, 2, 3, 4]:
lm = NGramLM(n=m).fit(train)
lambdas, _ = tune_lambdas(lm, dev)
print(f"\n {m}-gram：")
for _ in range(3):
print(" " + " ".join(lm.generate(rng, lambdas)))
print("""
观察：1-gram 是一堆随机词；2-gram 开始有局部的搭配；3-gram、4-gram 的句子在局部看相当通顺，
但长一点的句子经常出现主谓不一致（比如复数主语配单数动词）——模型只看得到最近的几个词。""")
# ---------------- 演示 5：主谓一致 ----------------
print_header("演示 5：主谓一致 —— 主语和动词之间隔得越远，n-gram 越不行")
lm4 = NGramLM(n=4).fit(train)
lambdas4, _ = tune_lambdas(lm4, dev)
def score(prefix, word):
return math.log(lm4.prob_interp(word, [BOS] * 3 + prefix, lambdas4))
print(f"\n 4-gram（插值）模型：比较 P(正确动词 | 前文) 和 P(错误动词 | 前文)")
print(f" {'修饰语个数':<10s}{'平均距离':>8s}{'准确率':>8s} {'无干扰名词':>10s}{'有干扰名词':>10s}")
for k in range(5):
tests = agreement_test_set(k, n=400)
acc, n = agreement_accuracy(score, 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(score, no_att)[0] if no_att else float("nan")
a1 = agreement_accuracy(score, att)[0] if att else float("nan")
dist = sum(t["distance"] for t in tests) / n
print(f" {k:<12d}{dist:>8.1f}{acc:>8.1%} {a0:>10.1%}{a1:>10.1%}" + (" ← 训练集里从没有过 4 个修饰语" if k == 4 else ""))
print("""
观察：
- 没有修饰语时，主语就在动词前面（距离 0），4-gram 全对。
- 有了修饰语，主语跑出了 3 个词的窗口，准确率一下子掉到 60% 左右（随机猜是 50%）。
剩下的那点优势来自窗口里的间接线索（比如 "that likes the cats ___" 中的 likes 暴露了主语是单数）。
- 如果修饰语里有和主语单复数相反的"干扰名词"，模型往往更容易被带偏。
- 这就是马尔可夫假设的代价：窗口外的信息，模型完全看不到。""")
# ---------------- 演示 6：存储 ----------------
print_header("演示 6：存储代价")
for m in [1, 2, 3, 4, 5]:
print(f" {m}-gram：需要存储 {len(lm5.counts[m]):>8,d} 个不同的 n-gram 计数")
print("""
观察：n 每加 1，需要存的计数就多很多。真实语料（几十亿词、几十万词表）中，5-gram 模型要存几十亿个计数，
这是 n-gram 模型的另一个瓶颈。神经语言模型（代码 02、04）只需要存固定数量的参数。""")
if __name__ == "__main__":
main()