"""
=============================================================================
第 6 讲 · 代码 4：序列到序列（seq2seq）模型 —— 编码器、解码器、信息瓶颈
=============================================================================
对应讲义：lectures/lecture06_lstm_seq2seq_attention.md 第 5、6、8 节
对应作业：CS224N Winter 2025 及更早的作业 3（神经机器翻译）
-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
机器翻译的输入和输出都是长度可变的序列，而且长度不同、语序不同。
seq2seq（Sutskever 等 2014；Cho 等 2014）用两个 RNN 解决这个问题：
编码器（encoder）：读完整个源句子，把它"压缩"成一个向量（最后的隐藏状态）；
解码器（decoder）：一个**条件语言模型**，从这个向量出发，一个词一个词地生成目标句子。
源句子 → [编码器 LSTM] → 一个向量 → [解码器 LSTM] → 目标句子
↑
整个源句子的全部信息都必须塞进这个向量 —— "信息瓶颈"
演示 1：模型结构
演示 2：训练（teacher forcing），贪心解码几个例子
演示 3：BLEU 和完全正确的比例；按源句子长度分组 —— 源句子越长，翻译越差（瓶颈的代价）
本文件里的数据处理、训练循环、评估函数也被 05 号文件（加了注意力的模型）复用。
-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
python 04_seq2seq.py # 完整运行（训练 10 轮，保存到 checkpoints/seq2seq.pt）
python 04_seq2seq.py --quick # 冒烟测试（少量数据、1 轮，不保存）
需要 PyTorch。完整运行约 3~6 分钟（CPU）。
"""
import os
# 必须在 import torch 之前设置：限制底层数学库（OpenMP）的线程数，原因见下面 set_num_threads 处的说明
os.environ.setdefault("OMP_NUM_THREADS", "1")
import random
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 torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
# 【Python 小知识】这些模型很小、而且是一步一步算的，多线程反而更慢（线程之间的协调开销大于收益，
# 实测单线程比 4 线程快约 2 倍），还会和同时运行的其他程序争抢 CPU。所以只用 1 个线程。
torch.set_num_threads(1)
from bleu import corpus_bleu
from toy_translation import Vocab, make_dataset, make_pair
CKPT_DIR = Path(__file__).parent / "checkpoints"
def print_header(title):
"""打印一个带分隔线的标题。"""
print()
print("=" * 70)
print(f" {title}")
print("=" * 70)
# =============================================================================
# 第 1 部分：数据
# =============================================================================
def make_batches(pairs, src_vocab, tgt_vocab, batch_size, shuffle=True, seed=0):
"""
把句子对分成批，并填充到相同长度。
【思路】
源句子编码为 [w1, ..., wS, ]；目标句子编码为 [, y1, ..., yT, ]。
解码器的输入是目标句去掉最后一个词（[, y1, ..., yT]），
要预测的是目标句去掉第一个词（[y1, ..., yT, ]）—— 和第 5 讲的语言模型一样，teacher forcing。
【返回】list[dict]，每个 dict 包含：
src (B, S_max)、src_len (B,)、tgt_in (B, T_max)、tgt_out (B, T_max)，以及原始词列表 src_words、tgt_words
"""
order = list(range(len(pairs)))
if shuffle:
random.Random(seed).shuffle(order)
batches = []
for start in range(0, len(order), batch_size):
chunk = [pairs[i] for i in order[start:start + batch_size]]
src = [src_vocab.encode(s, add_bos=False) for s, _ in chunk]
tgt = [tgt_vocab.encode(t) for _, t in chunk]
S, T = max(map(len, src)), max(map(len, tgt)) - 1
b = {"src": torch.zeros(len(chunk), S, dtype=torch.long), "src_len": torch.tensor([len(x) for x in src]),
"tgt_in": torch.zeros(len(chunk), T, dtype=torch.long), "tgt_out": torch.zeros(len(chunk), T, dtype=torch.long),
"src_words": [s for s, _ in chunk], "tgt_words": [t for _, t in chunk]}
for i, (s, t) in enumerate(zip(src, tgt)):
b["src"][i, :len(s)] = torch.tensor(s)
b["tgt_in"][i, :len(t) - 1] = torch.tensor(t[:-1])
b["tgt_out"][i, :len(t) - 1] = torch.tensor(t[1:])
batches.append(b)
return batches
# =============================================================================
# 第 2 部分：模型
# =============================================================================
class Seq2Seq(nn.Module):
"""
不带注意力的 seq2seq：双向 LSTM 编码器 + 单向 LSTM 解码器。
【结构】
编码器：嵌入 → 双向 LSTM（每个方向 hidden 维）。
前向 LSTM 读完整句的最后状态 + 后向 LSTM 读完整句的最后状态，拼接后经线性层
变成解码器的初始状态 (h_0, c_0)。**这就是整个源句子留给解码器的全部信息。**
解码器：每一步输入上一个目标词的嵌入，更新 LSTM 状态，用隐藏状态预测下一个词。
【参数】（构造函数）
src_V, tgt_V : 源 / 目标词表大小
emb, hidden : 嵌入维度；编码器每个方向的隐藏维度（解码器隐藏维度是它的 2 倍）
【接口】（05 号文件的注意力模型也实现同样的接口，所以训练和评估代码可以共用）
encode(src, src_len) → (enc_out, mask, state)
decode_step(y_prev, state, enc_out, mask, extra) → (logits, state, extra, attn)
forward(src, src_len, tgt_in) → logits (B, T, V)
"""
def __init__(self, src_V, tgt_V, emb=64, hidden=128):
super().__init__()
self.src_emb = nn.Embedding(src_V, emb, padding_idx=0)
self.tgt_emb = nn.Embedding(tgt_V, emb, padding_idx=0)
self.encoder = nn.LSTM(emb, hidden, batch_first=True, bidirectional=True)
self.h_proj = nn.Linear(2 * hidden, 2 * hidden)
self.c_proj = nn.Linear(2 * hidden, 2 * hidden)
self.decoder = nn.LSTMCell(emb, 2 * hidden)
self.out = nn.Linear(2 * hidden, tgt_V)
def encode(self, src, src_len):
"""
编码源句子。
【返回】
enc_out : (B, S, 2H)，每个源位置的隐藏状态（这个模型的解码器不使用它，注意力模型会用）
mask : (B, S)，True 表示真实的词，False 表示填充
state : (h_0, c_0)，解码器的初始状态
【Python 小知识】
pack_padded_sequence：告诉 LSTM 每个句子的真实长度，让它跳过填充位置。
这样后向 LSTM 从每个句子**真正的最后一个词**开始读，最后状态也不会被填充污染。
"""
packed = pack_padded_sequence(self.src_emb(src), src_len, batch_first=True, enforce_sorted=False)
out, (h, c) = self.encoder(packed)
enc_out, _ = pad_packed_sequence(out, batch_first=True, total_length=src.shape[1])
h = torch.cat([h[0], h[1]], dim=1) # 前向、后向的最后状态拼接
c = torch.cat([c[0], c[1]], dim=1)
return enc_out, src != 0, (self.h_proj(h), self.c_proj(c))
def decode_step(self, y_prev, state, enc_out, mask, extra=None):
"""
解码一步：输入上一个词，输出下一个词的 logits。
【返回】(logits (B, V), 新状态, extra（这里不用）, 注意力权重（这里没有，返回 None）)
"""
h, c = self.decoder(self.tgt_emb(y_prev), state)
return self.out(h), (h, c), None, None
def forward(self, src, src_len, tgt_in):
"""teacher forcing：逐步把真实的上一个目标词喂给解码器，返回所有时间步的 logits。"""
enc_out, mask, state = self.encode(src, src_len)
extra, logits = None, []
for t in range(tgt_in.shape[1]):
lg, state, extra, _ = self.decode_step(tgt_in[:, t], state, enc_out, mask, extra)
logits.append(lg)
return torch.stack(logits, dim=1)
# =============================================================================
# 第 3 部分：训练、解码、评估（04、05 号文件共用）
# =============================================================================
def train_nmt(model, train, src_vocab, tgt_vocab, epochs=10, lr=2e-3, batch_size=64, clip=1.0, seed=0, verbose=True):
"""
用 Adam + 交叉熵训练翻译模型（teacher forcing，忽略填充位置，梯度裁剪）。
【返回】list[float]，每轮的平均交叉熵
"""
torch.manual_seed(seed)
opt = torch.optim.Adam(model.parameters(), lr=lr)
hist = []
for epoch in range(epochs):
model.train()
t0, total, n_tok = time.perf_counter(), 0.0, 0
for b in make_batches(train, src_vocab, tgt_vocab, batch_size, seed=seed + epoch):
opt.zero_grad()
logits = model(b["src"], b["src_len"], b["tgt_in"])
loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), b["tgt_out"].reshape(-1), ignore_index=0)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
opt.step()
n = int((b["tgt_out"] != 0).sum())
total += loss.item() * n
n_tok += n
hist.append(total / n_tok)
if verbose:
print(f" 第 {epoch + 1:>2d} 轮：平均交叉熵 {hist[-1]:.4f}（用时 {time.perf_counter() - t0:.1f} 秒）")
return hist
@torch.no_grad()
def greedy_decode(model, src, src_len, max_len=40):
"""
贪心解码一批句子：每一步选概率最高的词，直到所有句子都生成了 或达到最大长度。
【返回】(预测的编号列表 list[list[int]], 注意力权重列表（每一步一个 (B, S) 张量；没有注意力时为空）)
"""
model.eval()
enc_out, mask, state = model.encode(src, src_len)
y = torch.full((src.shape[0],), 1, dtype=torch.long) #
extra, outputs, attns = None, [], []
done = torch.zeros(src.shape[0], dtype=torch.bool)
for _ in range(max_len):
logits, state, extra, attn = model.decode_step(y, state, enc_out, mask, extra)
y = logits.argmax(dim=1)
outputs.append(y)
if attn is not None:
attns.append(attn)
done |= y == 2
if done.all():
break
return torch.stack(outputs, dim=1).tolist(), attns
def translate(model, pairs, src_vocab, tgt_vocab, decode_fn=None):
"""用给定的解码函数（默认贪心）翻译一组句子对，返回 (候选译文列表, 参考译文列表)。"""
hyps, refs = [], []
for b in make_batches(pairs, src_vocab, tgt_vocab, 128, shuffle=False):
if decode_fn is None:
ids, _ = greedy_decode(model, b["src"], b["src_len"])
hyps += [tgt_vocab.decode(x) for x in ids]
else:
hyps += [decode_fn(model, b["src"][i:i + 1], b["src_len"][i:i + 1]) for i in range(len(b["src"]))]
refs += b["tgt_words"]
return hyps, refs
def evaluate(model, pairs, src_vocab, tgt_vocab, decode_fn=None):
"""返回 (语料 BLEU, 完全正确的比例)。"""
hyps, refs = translate(model, pairs, src_vocab, tgt_vocab, decode_fn)
bleu, _ = corpus_bleu(hyps, [[r] for r in refs])
exact = np.mean([h == r for h, r in zip(hyps, refs)])
return bleu, exact
def length_buckets(seed=7, n=300):
"""
生成按源句子长度分组的测试集：介词短语 0~4 个（训练集中最多 3 个，所以 4 个是"更长"的句子）。
【返回】dict：介词短语个数 → list[(英语, Zhoric)]
"""
rng = random.Random(seed)
return {k: [make_pair(rng, n_pp=k) for _ in range(n)] for k in range(5)}
def report_by_length(model, src_vocab, tgt_vocab, buckets, decode_fn=None, label="模型"):
"""打印每个长度组的平均源句子长度、BLEU 和完全正确的比例，并返回 BLEU 列表。"""
out = []
for k, pairs in buckets.items():
bleu, exact = evaluate(model, pairs, src_vocab, tgt_vocab, decode_fn)
avg_len = np.mean([len(s) for s, _ in pairs])
out.append(bleu)
print(f" {label:<14s}介词短语 {k} 个（源句平均 {avg_len:4.1f} 词）：BLEU {bleu:.3f}，完全正确 {exact:.1%}"
+ (" ← 比训练时更长" if k == 4 else ""))
return out
def load_or_train(model_cls, train, src_vocab, tgt_vocab, name, epochs=10, verbose=True):
"""有检查点就加载，否则训练并保存（供 05 号文件复用）。"""
model = model_cls(len(src_vocab), len(tgt_vocab))
path = CKPT_DIR / f"{name}.pt"
if path.exists():
model.load_state_dict(torch.load(path))
if verbose:
print(f" （从 {path.name} 加载了训练好的模型）")
return model
train_nmt(model, train, src_vocab, tgt_vocab, epochs=epochs, verbose=verbose)
CKPT_DIR.mkdir(exist_ok=True)
torch.save(model.state_dict(), path)
return model
# =============================================================================
# 演示
# =============================================================================
def main(quick=False):
train, test = make_dataset()
src_vocab, tgt_vocab = Vocab([s for s, _ in train]), Vocab([t for _, t in train])
if quick:
train, test = train[:2000], test[:200]
epochs = 1 if quick else 10
# ---------------- 演示 1：模型结构 ----------------
print_header("演示 1：seq2seq 模型结构（不带注意力）")
torch.manual_seed(0)
model = Seq2Seq(len(src_vocab), len(tgt_vocab))
for name, p in model.named_parameters():
print(f" {name:<28s}{str(tuple(p.shape)):<14s}{p.numel():>8,d}")
print(f" 合计 {sum(p.numel() for p in model.parameters()):,} 个参数；英语词表 {len(src_vocab)}，Zhoric 词表 {len(tgt_vocab)}")
# ---------------- 演示 2：训练 ----------------
print_header(f"演示 2：训练（{len(train)} 对句子，{epochs} 轮）")
t0 = time.perf_counter()
if quick:
train_nmt(model, train, src_vocab, tgt_vocab, epochs=epochs)
else:
path = CKPT_DIR / "seq2seq.pt"
train_nmt(model, train, src_vocab, tgt_vocab, epochs=epochs)
CKPT_DIR.mkdir(exist_ok=True)
torch.save(model.state_dict(), path)
print(" （模型已保存到 checkpoints/seq2seq.pt）")
print(f" 训练总用时 {time.perf_counter() - t0:.1f} 秒（取决于你的电脑）")
hyps, refs = translate(model, test[:5], src_vocab, tgt_vocab)
print("\n 测试集上的几个例子（贪心解码）：")
for (s, _), h, r in zip(test[:5], hyps, refs):
print(f" 英语：{' '.join(s)}\n 参考：{' '.join(r)}\n 模型：{' '.join(h)}{' ✓' if h == r else ''}\n")
# ---------------- 演示 3：按长度分组 ----------------
print_header("演示 3：BLEU 与源句子长度")
bleu, exact = evaluate(model, test, src_vocab, tgt_vocab)
print(f"\n 整个测试集（{len(test)} 句）：语料 BLEU {bleu:.3f}，完全正确 {exact:.1%}\n")
report_by_length(model, src_vocab, tgt_vocab, length_buckets(n=100 if quick else 300), label="seq2seq")
print("""
观察：源句子越长，翻译质量越差；比训练时更长的句子（4 个介词短语）更是明显下降。
原因是**信息瓶颈**：无论源句子多长，解码器得到的只有编码器最后的一个向量（这里 256 维）。
句子越长，要塞进这个向量的信息越多；而且翻译到句末的动词时，它离源句子中动词的位置很远。
下一个文件给解码器加上**注意力**，让它每一步都能直接"回头看"源句子的每个词。
（具体数字以你运行的结果为准。）""")
if __name__ == "__main__":
main(quick="--quick" in sys.argv)