"""
=============================================================================
第 3 讲 · 代码 02：任务的多样性 vs 数据的数量 —— 指令微调学到的是什么
=============================================================================

对应讲义：lectures/lecture03_instruction_tuning.md §3–§5

【这个脚本要回答的问题】
    FLAN（Wei et al. 2022）的核心实验：在很多种任务上做指令微调，模型在**没见过的任务**上的零样本表现会变好，
    而且训练任务越多越好。LIMA（Zhou et al. 2023）又说：对齐只需要很少的高质量数据。在迷你 GPT 上：
    1. 训练任务从 1 种增加到 6 种（总样本数固定），一个**没训练过的任务**（领导人）能答对多少？（演示 1）
    2. 预训练里从没被"问"过的任务（方位、反查首都），指令微调能不能把它"开出来"？（演示 2）
    3. 训练样本从 24 条增加到 1536 条，效果怎样变化？（演示 3）

【怎么读输出】
    每个数字是 8 个测试国家（只在预训练百科里出现过）× 4 种问法 = 32 道题的准确率（情感任务是 60 条评论 × 4）。
    演示 1 每种设置用 3 个不同的数据种子各训练一次，报告平均和每次的结果——单次训练的结果会有波动。

【运行】python3 02_task_diversity.py     （第一次约 2 分钟：训练 17 个 SFT 模型，每个约 6 秒；之后读缓存，约 15 秒）
【依赖】PyTorch；import 同目录的 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 instruct_lab import EVAL_COUNTRIES, evaluate, get_sft_model, question, render_chat, answer  # noqa: E402

TASK_SETS = [
    ["capital"],
    ["capital", "sentiment"],
    ["capital", "language", "export", "sentiment"],
    ["capital", "language", "export", "sentiment", "region", "revcap"],
]
SEEDS = [1, 2, 3]


def mean(xs):
    return sum(xs) / len(xs)


def demo1():
    """演示 1：训练任务的种类数 → 没训练过的任务（领导人）。"""
    print_header("演示 1：训练任务越多，没见过的任务越好吗？（总共 1200 条样本；测试任务 = 领导人，从没出现在微调数据里）")
    rows = []
    res = {}
    for tasks in TASK_SETS:
        lead, cap, examples = [], [], None
        for seed in SEEDS:
            m, tok, _, _ = get_sft_model(tasks=tasks, n=1200, data_seed=seed)
            lead.append(evaluate(m, tok, "leader")[0])
            cap.append(evaluate(m, tok, "capital")[0])
            if examples is None:
                c = EVAL_COUNTRIES[0]
                examples = [answer(m, tok, render_chat(question("leader", tpl, c))) for tpl in range(4)]
        res[len(tasks)] = (lead, cap, examples)
        rows.append([f"{len(tasks)} 种：{', '.join(tasks)}", f"{100 * mean(lead):.1f}%",
                     " / ".join(f"{100 * x:.0f}" for x in lead), f"{100 * mean(cap):.1f}%", " | ".join(examples)])
    print_table(["训练任务", "领导人（平均）", "3 个种子", "首都（平均）", f"问 {EVAL_COUNTRIES[0].name} 的领导人（种子 1，4 种问法）"],
                rows, aligns=["<", ">", ">", ">", "<"])
    one, two, four, six = (res[k] for k in (1, 2, 4, 6))
    empty4 = sum(1 for x in four[2] if x == "")
    print_observations([
        f"只用首都一个任务做指令微调，问领导人时平均 {100 * mean(one[0]):.1f}% 答对。看回答：几乎都是一个首都名"
        f"（{EVAL_COUNTRIES[0].name} 的首都是 {EVAL_COUNTRIES[0].capital}，有时是别国的首都），而且常常停不下来——"
        "它学到的是'这个格式 → 写一个首都'，没有真的读问题。",
        f"加上情感任务（2 种）没有帮助（{100 * mean(two[0]):.1f}%），回答里开始出现 positive；4 种任务 {100 * mean(four[0]):.1f}%，"
        f"6 种 {100 * mean(six[0]):.1f}%，平均都明显高于 1 种、2 种。任务一多，'猜一个任务、套它的答案'不再管用，模型更多地去读问题——"
        "这是 FLAN 核心发现的一个缩影。",
        f"但提升有限，离'训练过这个任务'差得很远（演示 3 里领导人在训练数据中时是 72%–91%）。4 种任务的模型对领导人问题"
        f"常常给出**空回答**（示例里 {empty4}/4 次）：它认得这是一个请求，却不知道该用哪种答案。",
        "注意 3 个种子之间的波动（例如 4 种任务：19 / 22 / 12）：只训练一次就下结论，很容易把运气当成规律（前置讲 P3 §6）。",
        "Wei et al.（2022）还发现，这种好处只在大模型上出现：在他们约 1000 亿参数的模型上明显，在 80 亿参数及以下的模型上，"
        "指令微调反而让没见过的任务变差。我们的迷你模型之所以能看到一点效果，是因为测试任务（领导人）的知识和问法都在预训练里出现过。",
    ])


def demo2():
    """演示 2：预训练里没被问过的任务，能不能靠指令微调'开出来'？"""
    print_header("演示 2：预训练里从没被'问'过的任务——方位（region）和反查首都（revcap）")
    rows = []
    res = {}
    for tasks in TASK_SETS[2:]:
        m, tok, _, _ = get_sft_model(tasks=tasks, n=1200, data_seed=1)
        accs = {t: evaluate(m, tok, t)[0] for t in ["region", "revcap", "leader"]}
        res[len(tasks)] = accs
        rows.append([", ".join(tasks), f"{100 * accs['region']:.1f}%", f"{100 * accs['revcap']:.1f}%", f"{100 * accs['leader']:.1f}%"])
    print_table(["训练任务", "方位", "反查首都", "领导人（对照）"], rows)
    print_observations([
        f"没训练过方位和反查首都时，这两种问题 {100 * res[4]['region']:.0f}% / {100 * res[4]['revcap']:.0f}%；"
        f"训练里加上它们（只用 12 个 SFT 国家的例子），测试国家就答到 {100 * res[6]['region']:.0f}% / {100 * res[6]['revcap']:.0f}%。",
        "和领导人不同，方位和反查首都在预训练里**从没以问答的形式出现过**，只在百科的陈述句里出现。所以模型没法从'别的任务'"
        "泛化到它们，必须在指令数据里见过这种任务。可一旦见过，知识就能迁移到 SFT 数据里没有的国家——这些国家的方位、首都，"
        "模型只在预训练的百科里读到过。",
        "这正是'预训练学知识、指令微调学怎样用知识'这句话的具体含义：指令微调为每**一类**问题建立'从问题到知识'的通道；"
        "通道建好之后，通道另一头是什么，由预训练决定。",
    ])


def demo3():
    """演示 3：数据量。"""
    print_header("演示 3：训练样本的数量（7 种任务；同样 400 步，样本少时每条会被看很多遍）")
    sizes = [24, 48, 96, 192, 384, 1536]
    tasks = ["capital", "language", "export", "sentiment", "leader", "region", "revcap"]
    rows = []
    res = {}
    for n in sizes:
        m, tok, _, _ = get_sft_model(n=n)
        accs = {t: evaluate(m, tok, t)[0] for t in tasks}
        res[n] = mean(list(accs.values()))
        rows.append([n, f"{n // 7}", f"{100 * res[n]:.1f}%"] + [f"{100 * accs[t]:.0f}" for t in tasks])
    print_table(["样本数", "每种任务约", "7 种任务平均"] + tasks, rows)
    print_observations([
        f"24 条样本（每种任务 3 条左右）时平均 {100 * res[24]:.1f}%；96 条 {100 * res[96]:.1f}%；384 条 {100 * res[384]:.1f}%；"
        f"1536 条 {100 * res[1536]:.1f}%。几百条之后，再多的数据几乎不再带来提升。",
        "原因和演示 2 一样：指令微调要学的是'怎样回答'，这件事本身不复杂；知识早就在预训练里了。"
        "LIMA（Zhou et al. 2023）用 1000 条精心挑选的例子微调 65B 的 LLaMA，人类评价接近当时的商用模型，"
        "并据此提出'表面对齐假说'（superficial alignment hypothesis）：模型的知识和能力几乎全部来自预训练，"
        "对齐主要是教它用哪种格式、哪种风格和用户交流。",
        "但'少'有下限：24 条时连格式都没学稳。数据的**覆盖面**（每种任务、每种问法都有例子）比总量更要紧（03 号脚本）。",
    ])


def main():
    demo1()
    demo2()
    demo3()


if __name__ == "__main__":
    main()