"""
=============================================================================
第 3 讲 · 代码 03：怎样微调 —— 全量还是只调一点、要不要 replay、损失算在哪里、问法要多样吗
=============================================================================

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

【这个脚本要回答的问题】
    同样 1200 条指令数据，"怎么训练"会不会比"训练什么"更要紧？
    1. 四种训练配方：全部参数 + 混入预训练语料（默认）、全部参数不混入、全部参数不混入但学习率小 3 倍、
       **只训练新加词的嵌入**（其余参数冻结）。在训练用过的国家、没用过的国家上分别测。（演示 1）
    2. 损失只算在回答上，还是连提示一起算？（演示 2）
    3. 每种任务只用 1 种问法训练，还是 4 种？在**训练时没见过的问法**上测。（演示 3）

【怎么读输出】
    "SFT 国家"是微调数据里出现过的 12 个国家；"近邻国家"是 4 个预训练问答页里见过、但不在微调数据里的；
    "测试国家"是 8 个只在预训练百科里出现过的。三者的差距说明模型是在"学会回答"还是在"背微调数据"。
    国家任务的准确率是 6 种任务（首都、语言、出口、领导人、方位、反查首都）× 4 种问法的平均。

【运行】python3 03_how_to_finetune.py     （第一次约 50 秒：训练 6 个 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 (DEV_COUNTRIES, EVAL_COUNTRIES, SFT_COUNTRIES, TRAIN_TEMPLATES, UNSEEN_TEMPLATES,  # noqa: E402
                          evaluate, get_sft_model, native_scores)

COUNTRY_TASKS = ["capital", "language", "export", "leader", "region", "revcap"]


def country_acc(model, tok, countries, templates=TRAIN_TEMPLATES):
    """6 种国家任务在一组国家上的平均准确率。"""
    return sum(evaluate(model, tok, t, countries=countries, templates=templates)[0] for t in COUNTRY_TASKS) / len(COUNTRY_TASKS)


def demo1():
    """演示 1：四种训练配方。"""
    print_header("演示 1：同样的数据，四种训练方法（7 种任务、1200 条样本、400 步）")
    recipes = [
        ("全部参数 · 混入预训练语料（默认）", {}),
        ("全部参数 · 不混入", {"mix": 0.0}),
        ("全部参数 · 不混入 · 学习率 3e-4", {"mix": 0.0, "lr": 3e-4}),
        ("只训练新词的嵌入（28 行）· 不混入", {"mix": 0.0, "train": "new_embeddings", "lr": 3e-3}),
    ]
    rows, res = [], {}
    for name, over in recipes:
        m, tok, _, _ = get_sft_model(**over)
        a = (country_acc(m, tok, SFT_COUNTRIES), country_acc(m, tok, DEV_COUNTRIES), country_acc(m, tok, EVAL_COUNTRIES))
        sent = evaluate(m, tok, "sentiment")[0]
        nat = native_scores(m, tok)
        res[name] = (a, sent, nat)
        rows.append([name, f"{100 * a[0]:.1f}%", f"{100 * a[1]:.1f}%", f"{100 * a[2]:.1f}%", f"{100 * sent:.1f}%",
                     f"{100 * nat['faq']:.1f}%", f"{nat['ppl']:.2f}"])
    print_table(["配方", "SFT 国家", "近邻国家", "测试国家", "情感", "原生问答格式", "预训练困惑度"], rows)
    d, nm, lr, emb = (res[r[0]] for r in recipes)
    print_observations([
        f"全部参数、不混入预训练语料：SFT 国家 {100 * nm[0][0]:.1f}%，测试国家只有 {100 * nm[0][2]:.1f}%——模型把微调数据里的"
        f"事实**背**了下来，却没学会'去预训练的知识里找答案'。同时困惑度涨到 {nm[2]['ppl']:.1f}，忘掉了很多东西。",
        f"把学习率调小 3 倍（{100 * lr[0][2]:.1f}%）或混入预训练语料（{100 * d[0][2]:.1f}%），测试国家都好得多：更新得越'温和'、"
        "越不偏离预训练的模型，越能保留并调用原有的知识。",
        f"最极端的'温和'：只训练 28 个新词的嵌入，原有的约 26 万个参数一个都不动。测试国家 {100 * emb[0][2]:.1f}%，"
        f"原生问答格式 {100 * emb[2]['faq']:.1f}%（和 base 模型一样）。<user>、<assistant> 这几个标记学到的，相当于一段"
        "'软提示'（soft prompt）：它们把对话格式翻译成模型原本就懂的那种文档。",
        f"它的困惑度（{emb[2]['ppl']:.2f}）仍比 base 的 1.89 高，原因不是旧知识被改了，而是新词的嵌入同时就是输出层的一行"
        "（权重共享）：模型开始在普通文本里也给 <end> 之类的新词分一点概率，真实下一个词的概率就少了一点。",
        "这和真实世界的经验一致：Lester et al.（2021）的提示微调（prompt tuning）只训练几十个'虚拟 token'的嵌入，"
        "在大模型上接近全量微调；LoRA 等参数高效方法只改很少的参数；全量微调时，小学习率、少步数、混入预训练数据都是常用的防遗忘手段。"
        "迷你模型的结论是：**微调数据越少、越窄，越要防止模型'背数据'**。",
    ])


def demo2():
    """演示 2：损失算在哪里。"""
    print_header("演示 2：损失只算在回答上（默认），还是连提示一起算？")
    rows, res = [], {}
    for name, over in [("只算回答（默认）", {}), ("提示和回答一起算", {"mask_prompt": False})]:
        m, tok, _, _ = get_sft_model(**over)
        a = country_acc(m, tok, EVAL_COUNTRIES)
        u = country_acc(m, tok, EVAL_COUNTRIES, templates=UNSEEN_TEMPLATES)
        nat = native_scores(m, tok)
        res[name] = (a, u, nat)
        rows.append([name, f"{100 * a:.1f}%", f"{100 * u:.1f}%", f"{100 * nat['faq']:.1f}%", f"{nat['ppl']:.3f}"])
    print_table(["损失", "测试国家（训练过的问法）", "测试国家（没见过的问法）", "原生问答格式", "预训练困惑度"], rows)
    a, b = res["只算回答（默认）"], res["提示和回答一起算"]
    print_observations([
        f"两种做法在测试国家上分别是 {100 * a[0]:.1f}% 和 {100 * b[0]:.1f}%，没见过的问法上 {100 * a[1]:.1f}% 和 {100 * b[1]:.1f}%。",
        "连提示一起算损失，等于同时训练模型'写出用户的问题'。它多学了 4 种问法的写法；可我们要的只是'回答'。"
        "多数指令微调（InstructGPT 的 SFT、Alpaca 等）只在回答上算损失；也有研究发现提示损失在数据很少时有一点正则化的作用"
        "——在我们的迷你模型上，这个选择的影响很小，远小于演示 1 的配方差别。",
    ])


def demo3():
    """演示 3：问法的多样性。"""
    print_header("演示 3：每种任务用 1 种问法还是 4 种问法训练？（同样 1200 条样本）")
    rows, res = [], {}
    for name, tpls in [("只用第 1 种问法", [0]), ("4 种问法（默认）", list(TRAIN_TEMPLATES))]:
        m, tok, _, _ = get_sft_model(templates=tpls)
        cells = []
        for tset, label in [([0], "第 1 种"), ([1, 2, 3], "第 2–4 种"), (list(UNSEEN_TEMPLATES), "没见过的第 5–6 种")]:
            cells.append(country_acc(m, tok, EVAL_COUNTRIES, templates=tset))
        sent = [evaluate(m, tok, "sentiment", templates=t)[0] for t in ([0], list(UNSEEN_TEMPLATES))]
        res[name] = (cells, sent)
        rows.append([name] + [f"{100 * x:.1f}%" for x in cells] + [f"{100 * sent[0]:.1f}% / {100 * sent[1]:.1f}%"])
    print_table(["训练用的问法", "测试：第 1 种问法", "第 2–4 种", "没见过的第 5–6 种", "情感：第 1 种 / 没见过的"], rows)
    one, four = res["只用第 1 种问法"], res["4 种问法（默认）"]
    print_observations([
        f"只用一种问法训练：同一种问法 {100 * one[0][0]:.1f}%，换成别的问法就掉到 {100 * one[0][1]:.1f}%（第 2–4 种）和 "
        f"{100 * one[0][2]:.1f}%（没见过的第 5–6 种）。模型学会的是'这个句式 → 这种答案'，而不是'理解请求'。",
        f"4 种问法训练：第 2–4 种 {100 * four[0][1]:.1f}%，**连训练时没见过的第 5–6 种**也有 {100 * four[0][2]:.1f}%。",
        "FLAN 给每个数据集写 10 个模板、T0 用众包收集了上千个提示，都是为了这个：问法的多样性让模型学到'请求'这个抽象的概念。"
        "Longpre et al.（2023，FLAN 2022）进一步发现，混入零样本、少样本、思维链等多种格式的模板，能同时改善各种提示方式下的表现。",
        "对提示工程师的启示：一个指令微调过的模型对**训练分布里常见的问法**最可靠。用户的问法越奇特，越可能落在训练分布之外。",
    ])


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


if __name__ == "__main__":
    main()