"""
=============================================================================
第 4 讲 · 代码 01：组合性差距 —— 每一跳都会，两跳合起来就不会了
=============================================================================

对应讲义：lectures/lecture04_chain_of_thought.md §2

【这个脚本要回答的问题】
    Press et al.（2023）发现：大模型常常**分别**答得出两个子问题（"X 的首都是哪里？""讲 L 语的是哪个国家？"），
    却答不出把它们合起来的问题——他们把这个差距叫作**组合性差距**（compositionality gap）。在迷你 GPT 上：
    1. 一个只学过一跳问题的指令微调模型，每一跳答得怎么样？两跳合起来呢？（演示 1）
    2. 在训练里加入两跳问题（直接写答案），训练过的组合、没训练过的组合各能答对多少？（演示 2）
    3. 由程序来分解：先问第一跳、把答案填进第二跳的问题里再问（Self-Ask 的外部版本）。（演示 3）

【怎么读输出】
    测试国家是 8 个"只在预训练百科里出现过"的国家：它们的任何事实都没进过微调数据。
    "桥"是用来指代国家的那个属性：capital（用首都指代）、language（用语言指代）、leader（用领导人指代，反转诅咒）。

【运行】python3 01_compositionality_gap.py     （第一次约 25 秒：训练 2 个模型；之后读缓存，约 10 秒）
【依赖】PyTorch；import 同目录的 cot_lab.py（以及第 3 讲的 instruct_lab.py），第 1 讲的 console.py。
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from cot_lab import (BRIDGES, EVAL_COUNTRIES, HOP1_TASK, TEST_COMBOS, TRAIN_COMBOS, attr, chained,  # noqa: E402
                     eval_twohop, get_cot_model, onehop_acc)


def demo1(m1, tok):
    """演示 1：只学过一跳的模型。"""
    print_header("演示 1：只学过一跳问题的模型——每一跳 vs 两跳合起来（测试国家）")
    rows = [["第一跳：用首都反查国家（revcap）", f"{100 * onehop_acc(m1, tok, 'revcap', EVAL_COUNTRIES):.1f}%"],
            ["第一跳：用语言反查国家（langcountry）", f"{100 * onehop_acc(m1, tok, 'langcountry', EVAL_COUNTRIES):.1f}%"],
            ["第一跳：用领导人反查国家（leadcountry，反转诅咒）", f"{100 * onehop_acc(m1, tok, 'leadcountry', EVAL_COUNTRIES):.1f}%"]]
    for t in ["capital", "language", "export", "leader", "region"]:
        rows.append([f"第二跳：已知国家，问 {t}", f"{100 * onehop_acc(m1, tok, t, EVAL_COUNTRIES):.1f}%"])
    print_table(["一跳问题（4 种问法）", "准确率"], rows)
    r = eval_twohop(m1, tok, TRAIN_COMBOS + TEST_COMBOS, EVAL_COUNTRIES, "direct", keep=3)
    print(f"\n  两跳问题（12 种组合 × 8 个国家），直接问这个模型：准确率 {100 * r['acc']:.1f}%；回答举例：")
    for q, out, g in r["examples"]:
        print(f"    {q}  →  {out!r}（标准答案 {g}）")
    print_observations([
        "用首都、用语言反查国家，模型大多答得出；已知国家问它的属性，几乎全对——两跳里的每一跳它都会（领导人做桥的除外）。",
        f"可把两跳合成一个问题直接问，准确率只有 {100 * r['acc']:.1f}%。看回答：问'首都是 grorken 的那个国家讲什么语言'，"
        "它答 volel——只答了第一跳。它把这个没见过的问题当成了它见过的反查问题。",
        "Press et al.（2023）在 GPT-3 上测量过这个差距，发现它**不随模型变大而缩小**：模型越大，一跳问题的正确率提高得比两跳问题快——"
        "更大的模型记住的事实更多，但'把两个事实串起来'的能力没有同步变强。",
    ])


def demo2(md, tok):
    """演示 2：训练里加入两跳问题（直接写答案）。"""
    print_header("演示 2：训练时加入两跳问题（直接写答案，只用 6 种组合）——训练过的组合 vs 没训练过的组合")
    rows, res = [], {}
    for label, combos in [("训练过的 6 种组合", TRAIN_COMBOS), ("没训练过的 6 种组合", TEST_COMBOS)]:
        for bridge in ["capital", "language", "leader"]:
            cs = [x for x in combos if x[0] == bridge]
            r = eval_twohop(md, tok, cs, EVAL_COUNTRIES, "direct")
            res[(label, bridge)] = r["acc"]
            rows.append([label, f"{bridge}（{', '.join(t for _, t in cs)}）", f"{100 * r['acc']:.1f}%"])
    print_table(["组合", "桥（目标）", "测试国家的准确率"], rows)
    seen = (res[("训练过的 6 种组合", "capital")] + res[("训练过的 6 种组合", "language")]) / 2
    unseen = (res[("没训练过的 6 种组合", "capital")] + res[("没训练过的 6 种组合", "language")]) / 2
    print_observations([
        f"训练过的组合（用首都 / 语言做桥）：测试国家平均 {100 * seen:.1f}%。这些国家的事实从没进过微调数据，"
        "可模型照样答出了不少——一部分原因是捷径：百科里同一个国家的首都、语言、出口写在同一篇词条里，"
        "模型可能直接学到了'这个首都 ↔ 这种语言'的关联，不必经过'国家'这一步。",
        f"没训练过的组合：平均只有 {100 * unseen:.1f}%。模型学到的是'这种问法 → 这种答案'，"
        "换一种'桥 × 目标'的搭配，它就不会了。这就是组合泛化的失败。",
        "用领导人做桥的题，训练过、没训练过都接近 0：第一跳本身就做不到（反转诅咒），多少两跳训练也救不回来。",
    ])
    return seen, unseen


def demo3(m1, md, tok):
    """演示 3：由程序来分解。"""
    print_header("演示 3：由程序来分解——先问第一跳，把答出的国名填进第二跳的问题里再问（测试国家）")
    rows = []
    res = {}
    for name, m in [("只学过一跳的模型", m1), ("也学过两跳（直接写答案）的模型", md)]:
        for label, combos in [("训练过的组合", TRAIN_COMBOS), ("没训练过的组合", TEST_COMBOS)]:
            for bridge in ["capital", "language", "leader"]:
                cs = [x for x in combos if x[0] == bridge]
                hop1 = final = n = 0
                for b, t in cs:
                    for c in EVAL_COUNTRIES:
                        x, a = chained(m, tok, b, t, c)
                        hop1 += x == c.name
                        final += a == attr(c, t)
                        n += 1
                res[(name, label, bridge)] = final / n
                rows.append([name, label, bridge, f"{100 * hop1 / n:.1f}%", f"{100 * final / n:.1f}%"])
    print_table(["模型", "组合", "桥", "第一跳答对", "最终答对"], rows)
    a = res[("只学过一跳的模型", "没训练过的组合", "capital")]
    b = res[("只学过一跳的模型", "没训练过的组合", "language")]
    print_observations([
        f"只学过一跳的模型，由程序分解之后，没训练过的组合也能答对（用首都做桥 {100 * a:.1f}%、用语言做桥 {100 * b:.1f}%）——"
        "因为对程序来说根本没有'没训练过的组合'：每一步都是模型会的一跳问题。",
        "最终的准确率大约是两跳各自准确率的乘积：分解之后，错误来自'第一跳答错'和'第二跳答错'，而不是'不会组合'。",
        "这就是 Self-Ask（Press et al. 2023）的想法：让模型（或程序）把问题拆成后续问题（follow-up questions），"
        "每个后续问题单独回答，必要时交给搜索引擎。第 9 讲的任务分解会系统地讲这条路线。",
        "领导人做桥的题仍然几乎全错（偶尔'最终答对'，是因为第一跳答错的国家碰巧和正确的国家有同样的属性，例如同一个方位）："
        "分解只能把模型**会的**东西串起来，模型不知道的一跳，分解也变不出来。",
    ])


def main():
    m1, tok, _ = get_cot_model(fmt="direct", n_twohop=0)        # 只学一跳
    md, _, _ = get_cot_model(fmt="direct")                      # 一跳 + 两跳（直接写答案）
    demo1(m1, tok)
    demo2(md, tok)
    demo3(m1, md, tok)


if __name__ == "__main__":
    main()