"""
=============================================================================
第 7 讲 · 代码 5：位置编码与长度泛化 —— 四种做法的比较
=============================================================================

对应讲义：lectures/lecture07_self_attention_transformer.md 第 5.5、11.3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
04 号文件发现：用**可学习的位置嵌入**的迷你 GPT，在训练时见过的句子长度上主谓一致几乎全对，
但句子一长（超过训练时的最大长度），准确率就明显下降；RNN 却几乎不受影响。

问题出在位置信息上吗？换一种位置编码会不会好一些？这里用完全相同的数据、模型大小和训练设置，
只改变位置信息的做法：

    learned     可学习的位置嵌入（BERT、GPT-2）
    sinusoidal  正弦位置编码（原始 Transformer）
    rope        旋转位置编码（LLaMA 等今天的大部分模型）
    none        不加任何位置信息（NoPE）—— 因果掩码本身会"泄露"一些位置信息

    演示 1：四种做法在训练分布内的困惑度
    演示 2：主谓一致准确率：训练时见过的长度 vs 更长的句子（每种 2 个随机种子）

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 05_length_generalization.py           # 完整运行（训练 8 个模型）
    python 05_length_generalization.py --quick   # 冒烟测试

需要 PyTorch。完整运行约 6~10 分钟（CPU）。导入本目录的 04_mini_gpt.py 和第 5 讲的代码。
"""

import os

os.environ.setdefault("OMP_NUM_THREADS", "1")
import importlib
import sys
import time

import numpy as np
import torch

gpt = importlib.import_module("04_mini_gpt")      # 会顺便把第 5 讲的文件夹加进搜索路径
from toy_language import Vocab, agreement_test_set, generate_corpus  # noqa: E402

torch.set_num_threads(1)

POS_KINDS = ["learned", "sinusoidal", "rope", "none"]
POS_NAMES = {"learned": "可学习", "sinusoidal": "正弦", "rope": "RoPE", "none": "无位置信息"}


def main(quick=False):
    n_train, epochs, seeds = (2000, 1, [0]) if quick else (20000, 8, [0, 1])
    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))
    ks = [0, 1, 2, 3, 4, 6, 8, 10]
    tests = {k: agreement_test_set(k, n=400) for k in ks}
    longest_train = max(len(s) for s in generate_corpus(20000, seed=0)) + 2      # 加上 <s>、</s>

    models, t0 = {}, time.perf_counter()
    for pos in POS_KINDS:
        for seed in seeds:
            models[(pos, seed)] = gpt.train_gpt(train, vocab, pos=pos, epochs=epochs, seed=seed, verbose=False)
        print(f"  （{POS_NAMES[pos]}：训练完成，累计用时 {time.perf_counter() - t0:.0f} 秒）")

    # ---------------- 演示 1 ----------------
    gpt.print_header("演示 1：训练分布内的困惑度（测试集，和训练语料同分布）")
    print()
    for pos in POS_KINDS:
        ppl = [gpt.perplexity(models[(pos, s)], test, vocab) for s in seeds]
        print(f"    {POS_NAMES[pos]:<8s}" + "  ".join(f"{p:.3f}" for p in ppl) + "   （每个随机种子一个数）")
    print("""
  观察：四种做法在训练分布内几乎一样好 —— 包括完全不加位置信息的模型。
    因果掩码让第 t 个位置只能看到 t 个词，模型可以从"能看到多少东西"里推断出自己大概在第几个位置
    （例如关注 <s> 的权重会随位置变小）。在这门玩具语言上，这点位置信息已经够用了。""")

    # ---------------- 演示 2 ----------------
    gpt.print_header(f"演示 2：主谓一致准确率（训练语料中最长的句子是 {longest_train} 个词元，含 <s>、</s>）")
    print(f"\n    {'修饰语':<10s}" + "".join(f"{k:>10d}" for k in ks))
    print(f"    {'最长前缀':<8s}" + "".join(f"{max(len(t['prefix']) for t in tests[k]) + 1:>10d}" for k in ks))
    print("    " + "-" * (10 + 10 * len(ks)))
    for pos in POS_KINDS:
        accs = np.array([[gpt.agreement_accuracy(models[(pos, s)], vocab, tests[k]) for k in ks] for s in seeds])
        print(f"    {POS_NAMES[pos]:<8s}" + "".join(f"{a:>10.1%}" for a in accs.mean(axis=0)))
        if len(seeds) > 1:
            print(f"    {'':<10s}" + "".join(f"{f'{lo:.0%}~{hi:.0%}':>10s}" for lo, hi in zip(accs.min(0), accs.max(0))))
    print("""
  （修饰语 0~3 个是训练时见过的；4 个及以上的句子比训练时的任何句子都长。第二行是各随机种子的最低~最高。）

  观察（编写讲义时的结果；具体数字以你运行的结果为准）：
    - 训练时见过的长度上，四种做法都全对。
    - 句子变长之后差别很大，而且和"直觉"不太一样：
        · 无位置信息：几乎全对，推广得最好。找主语只需要看"内容"（哪个词是名词、单数还是复数），
          不依赖位置的模型也就不存在"没见过的位置"这个问题。
        · 可学习的位置嵌入：轻微下降。没训练过的位置向量还是初始化时很小的随机值，
          这些位置上的模型几乎相当于"无位置信息"，所以影响不大。
        · 正弦编码：6 个修饰语以内全对，更长时开始下降，而且两个随机种子差别很大。
        · RoPE：最差，一超过训练长度就明显下降。它的打分取决于相对距离，
          而训练时见过的距离最多二十个词左右；更远的距离对应的旋转角度（尤其是转得慢的那些维度）
          模型从没见过，打分变得不可预测。
    - 这和文献的发现一致：RoPE 本身不能外推到比训练时更长的序列，所以今天的大模型想支持更长的上下文时，
      要对 RoPE 做位置插值、NTK 缩放、YaRN 等调整（把没见过的角度"压缩"回见过的范围），再用长文本继续训练；
      Kazemnejad 等（2023）在小规模实验中也发现不加位置信息的解码器长度泛化最好。
    - 但不要因此得出"位置编码没用"的结论：真实语言里，精确的位置信息对复制、计数、局部词序等都很重要，
      真实的大模型几乎都使用 RoPE。我们的玩具任务只考察了"找主语"这一件事 —— 玩具实验能说明的事情是有边界的。
    - 对比：第 5 讲的 RNN 在 8 个修饰语时仍然全对。RNN 每一步用同一套参数，天生没有"位置"的问题。
      但 RNN 的困惑度并不比 Transformer 低，训练也无法并行 —— 没有哪种结构在所有方面都更好。""")


if __name__ == "__main__":
    main(quick="--quick" in sys.argv)