"""
=============================================================================
第 3 讲 · 代码 05：指令数据从哪里来 —— 让模型自己写（Self-Instruct），以及数据质量 vs 数量
=============================================================================

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

【这个脚本要回答的问题】
    人工写指令数据很贵。Self-Instruct（Wang et al. 2023）和 Alpaca（Taori et al. 2023）让模型自己生成指令数据；
    LIMA（Zhou et al. 2023）说少量高质量的数据就够。在迷你 GPT 上：
    1. 让 base 模型自己写问答页：写出来的有多少能用、多少是对的、有多少重复？（演示 1）
    2. 用它自己写的数据做指令微调，和用"人写的"数据比，差在哪里？（演示 2）
    3. 数据里混进错误答案（30%、60%），和数据少（150 条），哪个更伤？（演示 3）

【运行】python3 05_self_instruct.py     （第一次约 50 秒：生成 300 个问答页 + 训练 7 个 SFT 模型；之后约 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, TASKS, TRAIN_TEMPLATES, UNSEEN_TEMPLATES, evaluate, get_sft_model,  # noqa: E402
                          self_instruct_data)

QA4 = ["capital", "language", "export", "leader"]


def avg(model, tok, tasks, templates=TRAIN_TEMPLATES):
    """若干任务在测试国家上的平均准确率。"""
    return sum(evaluate(model, tok, t, templates=templates)[0] for t in tasks) / len(tasks)


def demo1():
    """演示 1：base 模型自己写的问答对。"""
    print_header("演示 1：让 base 模型自己写问答页（从 '<doc> faq : q :' 开始，温度 1，300 页）")
    data, st = self_instruct_data()
    print_table(["项目", "数量"], [
        ["抽出来的 '问题 ? 答案' 对", st["generated"]],
        ["能解析成 4 种已知问题之一", f"{st['parsed']}（答案正确 {100 * st['correct']:.1f}%）"],
        ["涉及的国家", f"{st['countries']} 个（其中测试国家：{', '.join(st['held_countries']) or '无'}，"
                    f"这 {st['dropped_held']} 条不用于训练）"],
        ["去重之后留下的训练样本", st["kept"]],
    ])
    print("  前 4 条：" + " | ".join(f"{p.split('<user> ')[1].split(' <end>')[0]} → {r.split(' <end>')[0]}" for p, r in data[:4]))
    print_observations([
        f"base 模型写的问答对 {100 * st['correct']:.1f}% 是对的：它在复述自己预训练里学到的东西，所以几乎不会错。",
        f"但去重之后只剩 {st['kept']} 条：{st['generated']} 对里绝大多数是重复的。模型反复写它最熟悉的那几个问题——"
        "Self-Instruct 论文花了很大力气对付这个问题（和已有指令的 ROUGE-L 相似度超过 0.7 就丢掉）。",
        "它写的问题只有一种问法（预训练问答页里的那一种），只涉及 4 种关系；方位、反查首都、评论情感，它根本不会想到去写。"
        "**模型自己生成的数据，覆盖不了它自己不会的东西**。",
        f"它还写了 {len(st['held_countries'])} 个'只在百科里出现过'的国家（{', '.join(st['held_countries']) or '无'}）——"
        "知识从百科迁移到了问答格式（第 1 讲 §8.5）。这些样本我们丢掉了，免得测试题'泄漏'进训练数据。",
    ])
    return st


def demo2(st):
    """演示 2：自己写的数据 vs 人写的数据。"""
    print_header("演示 2：用自己写的数据做指令微调，和人写的数据比（测试国家；4 种关系）")
    n = st["kept"]
    configs = [
        (f"base 自己写的，去重后 {n} 条", dict(source="self")),
        ("base 自己写的，不去重", dict(source="self_nodedup")),
        (f"人写的，同样 4 种关系、只用第 1 种问法、{n} 条", dict(tasks=QA4, templates=[0], n=n)),
        (f"人写的，4 种关系、4 种问法、{n} 条", dict(tasks=QA4, n=n)),
        ("人写的，7 种任务、4 种问法、1200 条（默认）", dict()),
    ]
    rows, res = [], {}
    for name, over in configs:
        m, tok, _, _ = get_sft_model(**over)
        a0 = avg(m, tok, QA4, templates=[0])
        a1 = avg(m, tok, QA4, templates=[1, 2, 3])
        au = avg(m, tok, QA4, templates=list(UNSEEN_TEMPLATES))
        other = avg(m, tok, ["region", "revcap", "sentiment"])
        res[name] = (a0, a1, au, other)
        rows.append([name, f"{100 * a0:.1f}%", f"{100 * a1:.1f}%", f"{100 * au:.1f}%", f"{100 * other:.1f}%"])
    print_table(["训练数据", "第 1 种问法", "第 2–4 种", "没见过的第 5–6 种", "方位 / 反查首都 / 情感"], rows)
    s, nd, g1, g4, d = (res[c[0]] for c in configs)
    print_observations([
        f"只用自己写的 {n} 条数据，第 1 种问法（也就是它自己写的那种）就有 {100 * s[0]:.1f}%：模型'教会了自己'用对话格式回答。"
        f"这和 Alpaca 的故事一样——用 52K 条 text-davinci-003 生成的指令数据微调 LLaMA-7B，就得到了一个像样的指令跟随模型。",
        f"不去重（{st['parsed'] - st['dropped_held']} 条，绝大多数是重复）：第 1 种问法 {100 * nd[0]:.1f}%，其余问法 "
        f"{100 * nd[1]:.1f}% / {100 * nd[2]:.1f}%，和去重后差不多。重复让模型在那一种问法上更熟练一点，却没有带来新的问法、新的任务。",
        f"短板在**多样性**：换成别的问法，自己写的数据只有 {100 * s[1]:.1f}%（第 2–4 种）、{100 * s[2]:.1f}%（没见过的），"
        f"和'人写的、只用一种问法'差不多（{100 * g1[1]:.1f}%、{100 * g1[2]:.1f}%）；同样条数但有 4 种问法的人写数据是 "
        f"{100 * g4[1]:.1f}%、{100 * g4[2]:.1f}%。",
        f"方位、反查首都、情感这三种它没写过的任务，自己写的数据训练出来只有 {100 * s[3]:.1f}%；默认配方（7 种任务）是 {100 * d[3]:.1f}%。",
    ])


def demo3():
    """演示 3：数据质量 vs 数量。"""
    print_header("演示 3：错误答案 vs 数据太少（7 种任务、4 种问法；测试国家）")
    configs = [
        ("1200 条，全对（默认）", dict()),
        ("150 条，全对", dict(n=150)),
        ("1200 条，30% 的答案是错的", dict(noise=0.3)),
        ("1200 条，60% 的答案是错的", dict(noise=0.6)),
    ]
    rows, res = [], {}
    tasks = list(TASKS)
    for name, over in configs:
        m, tok, _, _ = get_sft_model(**over)
        accs = {t: evaluate(m, tok, t)[0] for t in tasks}
        res[name] = sum(accs.values()) / len(accs)
        rows.append([name, f"{100 * res[name]:.1f}%"] + [f"{100 * accs[t]:.0f}" for t in tasks])
    print_table(["训练数据", "7 种任务平均"] + tasks, rows)
    c, small, n30, n60 = (res[x[0]] for x in configs)
    print_observations([
        f"1200 条全对 {100 * c:.1f}%；只有 150 条全对的 {100 * small:.1f}%；1200 条里 30% 是错的 {100 * n30:.1f}%；60% 是错的 {100 * n60:.1f}%。",
        "错误答案的伤害在于：模型学会了'有时候答一个别的国家的答案'——它在测试国家上也会这样做。这是在**训练模型编造**。"
        "第 5 讲会看到 Gekhman et al.（2024）的发现：用模型不知道的事实做微调，会让它在别的问题上也更爱编造。",
        "LIMA 的结论'少而精'在这里的形式是：几百条正确、覆盖全面的数据，就接近数据量饱和时的效果（02 号脚本演示 3）；"
        "大量带噪声的数据不能替代质量。Zhou et al.（2023）的消融也发现，提高数据的多样性和质量比单纯增加数量有用得多。",
    ])


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


if __name__ == "__main__":
    main()