"""
=============================================================================
第 3 讲 · 代码 01：指令微调前后 —— 同一组提示，base 模型和 SFT 模型怎样回答
=============================================================================

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

【这个脚本要回答的问题】
    1. 用对话格式提问（"<user> what is the capital of volel ? <end> <assistant>"），base 模型会怎样？
       指令微调之后呢？和"为 base 模型精心伪装的原生格式"比，谁更好？（演示 1–2）
    2. 指令微调有没有让模型忘掉预训练学到的东西（"对齐税"）？混入预训练语料（replay）有没有用？（演示 3）
    3. 对话模板是一种"文档格式"：换掉、去掉模板里的标记，指令微调过的模型还会回答吗？（演示 4）

【怎么读输出】
    测试国家是 8 个"只在预训练的百科词条里出现过"的国家：指令微调数据里**一个**都没有。
    所以 SFT 模型在它们上答对，靠的是预训练学到的知识。每种任务 8 个国家 × 4 种问法 = 32 道题（情感任务 60 条评论 × 4）。

【运行】python3 01_base_vs_sft.py     （约 20 秒；第一次运行要训练 2 个 SFT 模型，每个约 6 秒，之后读缓存）
【依赖】PyTorch；import 同目录的 instruct_lab.py，第 1 讲的 tiny_gpt.py、console.py。
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import pct, print_header, print_observations, print_table  # noqa: E402
from instruct_lab import (EVAL_COUNTRIES, EVAL_REVIEWS, TASKS, TRAIN_TEMPLATES, answer, base_model,  # noqa: E402
                          evaluate, extend_model, extended_tokenizer, get_sft_model, gold_answer, native_scores,
                          question, render_chat)
from tiny_gpt import generate, score_options  # noqa: E402

ALL_TASKS = list(TASKS)


def native_prompt(task, c=None, review=None):
    """
    为 base 模型"精心伪装"的提示：把每种任务写成预训练语料里最像的那种文档的开头（第 1 讲 §8.3 的技巧）。

    【返回】(提示, 判对方式)：判对方式是 "first_word"（续写的第一个词就是答案）或 "rating"（比较 positive/negative）
    【边界】region 和 revcap 在预训练里从没被"问"过，只能借百科的句式："X lies in the" → 方位；"K is the capital of" → 国名。
    """
    if task == "sentiment":
        return f"<doc> reviews : review : {review[0]} . rating :", "rating"
    if task == "region":
        return f"<doc> article : {c.name} lies in the", "first_word"
    if task == "revcap":
        return f"<doc> article : {c.capital} is the capital of", "first_word"
    rel = task
    return f"<doc> faq : q : {question(task, 0, c)} a :", "first_word"


def native_accuracy(model, tok, task):
    """base 模型在原生格式下的准确率（测试国家 / 测试评论）。"""
    from tiny_gpt import top_next
    ok = n = 0
    items = [(None, r) for r in EVAL_REVIEWS] if task == "sentiment" else [(c, None) for c in EVAL_COUNTRIES]
    for c, rv in items:
        p, how = native_prompt(task, c, rv)
        if how == "rating":
            ok += score_options(model, tok, p, ["positive", "negative"])[2] == rv[1]
        else:
            ok += top_next(model, tok, p, 1)[0][0] == gold_answer(task, c, rv)
        n += 1
    return ok / n


def demo1(base_ext, base, btok, sft, tok):
    """演示 1：三种方式问同一组问题。"""
    print_header("演示 1：同一组问题——base 模型用对话格式、base 模型用'伪装'的原生格式、SFT 模型用对话格式")
    rows = []
    res = {}
    for task in ALL_TASKS:
        a0, _ = evaluate(base_ext, tok, task)
        a1 = native_accuracy(base, btok, task)
        a2, _ = evaluate(sft, tok, task)
        res[task] = (a0, a1, a2)
        n = len(EVAL_REVIEWS) * 4 if task == "sentiment" else len(EVAL_COUNTRIES) * 4
        rows.append([task, TASKS[task][0], pct(round(a0 * n), n), f"{100 * a1:.1f}%", pct(round(a2 * n), n)])
    print_table(["任务", "问法之一", "base · 对话格式", "base · 原生格式（只用 1 种）", "SFT · 对话格式（4 种问法）"], rows)
    print_observations([
        "base 模型完全不会用对话格式回答（全部 0%）：<user>、<assistant> 这些标记它从没见过（它们的嵌入只是旧词嵌入的平均），"
        "指令里的 name / tell / which 也是新词。",
        f"同一个 base 模型，换成它熟悉的原生格式，首都 / 语言 / 出口 / 领导人都能答出来（{100 * res['capital'][1]:.0f}% / "
        f"{100 * res['language'][1]:.0f}% / {100 * res['export'][1]:.0f}% / {100 * res['leader'][1]:.0f}%）："
        "**知识一直都在**，只是要用对的'文档开头'才取得出来。",
        f"指令微调后，同一个模型用四种不同的问法、统一的对话格式，就能取出这些知识（SFT 列）。测试国家一个都没出现在微调数据里——"
        "指令微调教会的是'怎样回应一个请求'，知识来自预训练。",
        f"region 和 revcap 在预训练里从没被'问'过，但百科里有现成的句式（'X lies in the …'、'K is the capital of …'），"
        f"借这些句式 base 模型能答对 {100 * res['region'][1]:.0f}% 和 {100 * res['revcap'][1]:.0f}%；SFT 模型用问句的形式答到 "
        f"{100 * res['region'][2]:.0f}% 和 {100 * res['revcap'][2]:.0f}%，比精心伪装的原生格式还低一点。"
        "指令微调换来的是一个**统一的、不需要伪装的接口**，不是更多的知识。",
    ])
    return res


def demo2(base_ext, base, btok, sft, tok):
    """演示 2：看几个具体的回答。"""
    print_header("演示 2：几个具体的回答（测试国家 " + EVAL_COUNTRIES[0].name + "）")
    c = EVAL_COUNTRIES[0]
    rows = []
    prompts = [("capital", 0), ("capital", 2), ("leader", 3), ("region", 0), ("revcap", 1)]
    for task, tpl in prompts:
        q = question(task, tpl, c)
        p = render_chat(q)
        b = generate(base_ext, tok, p, max_new_tokens=8, stop=("<end>", "<doc>", "."))
        s = answer(sft, tok, p)
        rows.append([q, b, s, gold_answer(task, c)])
    raw = f"<doc> {question('capital', 0, c)}"
    rows.append([f"（base 不用对话模板）{raw}", generate(base, btok, raw, max_new_tokens=10, stop=("<doc>",)),
                 answer(sft, tok, render_chat(question("capital", 0, c))), c.capital])
    print_table(["问题", "base 模型续写", "SFT 模型回答（对话模板）", "标准答案"], rows, aligns=["<", "<", "<", "<"])
    print_observations([
        "base 模型拿到对话格式时只会'胡乱续写'：这些标记对它来说只是几个从没见过的词。'tell me the capital of volel .' "
        "被续写成了 'is ruled by kiro .'——它把 volel 后面接成了百科里的句子。",
        "最后一行是第 1 讲 §8.3 的老问题：光秃秃一个问题，base 模型接着出了更多的题；SFT 模型直接回答。",
    ])


def demo3(base, btok, tok):
    """演示 3：对齐税——指令微调有没有让模型忘掉东西？"""
    print_header("演示 3：指令微调的代价——用第 1 讲的原生格式测'还记得多少'（24 个国家 × 4 种关系；300 条评论）")
    rows = []
    res = {}
    configs = [("base（指令微调之前）", None),
               ("SFT · 混入预训练语料（默认配方）", {}),
               ("SFT · 不混入预训练语料", {"mix": 0.0})]
    for name, over in configs:
        if over is None:
            sc = native_scores(base, btok)
            inst = "–"
        else:
            m, t, _, _ = get_sft_model(**over)
            sc = native_scores(m, t)
            inst = f"{100 * sum(evaluate(m, t, task)[0] for task in ALL_TASKS) / len(ALL_TASKS):.1f}%"
        res[name] = sc
        rows.append([name, f"{100 * sc['faq']:.1f}%", f"{100 * sc['rating']:.1f}%", f"{sc['ppl']:.3f}", inst])
    print_table(["模型", "原生问答格式", "原生评论格式", "预训练验证集困惑度", "对话格式（7 种任务平均）"], rows)
    b, m, n = (res[c[0]] for c in configs)
    print_observations([
        f"不混入预训练语料时，困惑度从 {b['ppl']:.3f} 升到 {n['ppl']:.3f}，原生问答格式的准确率从 {100 * b['faq']:.1f}% 变成 "
        f"{100 * n['faq']:.1f}%：模型在适应新格式的同时，'续写文档'的本事变差了。这就是**对齐税**（alignment tax）的一种形式。",
        f"每一步混入一半批量的预训练语料（replay）后，困惑度是 {m['ppl']:.3f}、原生问答 {100 * m['faq']:.1f}%：代价小了很多。"
        "InstructGPT 的 PPO-ptx 用的是同一个思路——在对齐的目标里加上预训练的语言模型损失（讲义 §4.3）。",
        "对话格式上，两种配方的平均准确率见最后一列：replay 不但减少了遗忘，还让新格式下的泛化更好（03 号脚本会细看）。",
    ])


def demo4(sft, tok):
    """演示 4：对话模板是一种文档格式。"""
    print_header("演示 4：模板改一改，SFT 模型还认得吗？（首都任务，测试国家 × 4 种问法）")
    variants = [
        ("训练时的模板", lambda q: render_chat(q)),
        ("去掉开头的 <doc>", lambda q: render_chat(q)[len("<doc> "):]),
        ("去掉结尾的 <assistant>", lambda q: render_chat(q)[:-len(" <assistant>")]),
        ("<user> 换成第 1 讲的 q :", lambda q: render_chat(q).replace("<user>", "q :")),
        ("<assistant> 换成 a :", lambda q: render_chat(q).replace("<assistant>", "a :")),
        ("完全不用对话标记（原生问答格式）", lambda q: f"<doc> faq : q : {q} a :"),
        ("加一个训练时没见过的系统提示", lambda q: render_chat(q, system="answer in a full sentence")),
    ]
    rows = []
    res = {}
    for name, make in variants:
        first = clean = n = 0
        ex = None
        for c in EVAL_COUNTRIES:
            for tpl in TRAIN_TEMPLATES:
                out = answer(sft, tok, make(question("capital", tpl, c)))
                first += (out.split()[:1] == [c.capital])
                clean += out == c.capital
                n += 1
                ex = ex or out
        res[name] = (first / n, clean / n)
        rows.append([name, pct(first, n), pct(clean, n), ex])
    print_table(["模板变体", "第一个词是答案", "只写答案就停", "第一个例子的回答"], rows, aligns=["<", ">", ">", "<"])
    t = res["训练时的模板"]
    nd, na = res["去掉开头的 <doc>"], res["去掉结尾的 <assistant>"]
    qu, qa = res["<user> 换成第 1 讲的 q :"], res["<assistant> 换成 a :"]
    nat, sysp = res["完全不用对话标记（原生问答格式）"], res["加一个训练时没见过的系统提示"]
    print_observations([
        f"训练时的模板：{100 * t[0]:.1f}% 的回答第一个词就是答案，{100 * t[1]:.1f}% 只写答案就停。",
        f"去掉结尾的 <assistant>：第一个词是答案的只有 {100 * na[0]:.1f}%——模型不知道'轮到它说话了'。"
        f"去掉开头的 <doc>：答案仍然对（{100 * nd[0]:.1f}%），但只有 {100 * nd[1]:.1f}% 能停下来，其余的接着写起了问答页。",
        f"把 <user> 换成 q :、<assistant> 换成 a :：第一个词对的分别是 {100 * qu[0]:.1f}% 和 {100 * qa[0]:.1f}%，"
        f"能停下来的 {100 * qu[1]:.1f}% 和 {100 * qa[1]:.1f}%。完全用原生问答格式：答案 {100 * nat[0]:.1f}% 对，"
        f"但{'没有一次' if nat[1] == 0 else f'只有 {100 * nat[1]:.1f}%'}停下来——它回到了 base 模型'续写问答页'的行为。",
        "模板决定了模型'以为自己在写什么文档'：对话模板 = 只答这一句；问答页 = 答完接着出题。调用 chat 模型时，模板一定要和训练时"
        "完全一致——这正是 API 替你做的事（P2 §3.3）；自己部署开源模型时，用错 chat 模板是最常见、也最难发现的错误之一"
        "（输出'差一点'，而不是报错）。",
        f"加一个训练时从没见过的系统提示（'answer in a full sentence'）：第一个词是答案的 {100 * sysp[0]:.1f}%，"
        f"只写答案就停的只有 {100 * sysp[1]:.1f}%——模型被这段陌生的前缀弄糊涂了，开始写半句百科。04 号脚本专门训练模型听系统提示。",
    ])


def main():
    base, btok = base_model()
    tok = extended_tokenizer()
    base_ext = extend_model(base, tok)
    sft, tok, _, _ = get_sft_model()
    demo1(base_ext, base, btok, sft, tok)
    demo2(base_ext, base, btok, sft, tok)
    demo3(base, btok, tok)
    demo4(sft, tok)


if __name__ == "__main__":
    main()