"""
=============================================================================
第 10 讲 · 代码 1：指令微调（SFT）—— 教一个只会续写的模型听指令
=============================================================================

对应讲义：lectures/lecture10_post_training.md 第 2 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 9 讲预训练的迷你 GPT 只会"续写评论"。我们想让它变成一个能听指令的"助手"：

    <user> write a positive review <assistant>                      → 写一条正面的评论
    <user> write a negative review <assistant>                      → 写一条负面的评论
    <user> is this review positive or negative ? <评论> <assistant>  → 回答 positive 或 negative

做法：**监督微调**（Supervised Fine-Tuning，SFT）—— 准备一批"指令 + 示范回答"，在上面继续训练语言模型，
只在回答的部分计算损失。

    演示 1：对话模板和损失掩码长什么样
    演示 2：三个模型比一比 —— 只预训练、没有预训练直接 SFT、预训练 + SFT

训练好的模型保存到 checkpoints/sft.pt，02、04 号文件会加载它。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 01_sft.py           # 完整运行（约 3~6 分钟；如果第 9 讲的 clm.pt 不存在，还要先预训练）
    python 01_sft.py --quick   # 冒烟测试

【关于结果】编写讲义时这个文件只运行了 --quick 冒烟测试，没有完整运行。
    "预期会看到"描述的是应该出现的趋势，具体数字以你运行的结果为准。
"""

import os

os.environ.setdefault("OMP_NUM_THREADS", "1")
import copy
import sys
import time

import torch

import post_utils as P                 # 先导入它：它会把第 7、8、9 讲的文件夹加进搜索路径
import lm_utils as U  # noqa: E402     第 9 讲的工具模块
from toy_reviews import Vocab, make_dataset  # noqa: E402

torch.set_num_threads(1)


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


def evaluate(model, vocab, cls_test, n_gen):
    """
    评估一个模型听不听指令。

    【返回】dict：
        pos  : "写正面评论"时，生成的评论确实是正面的比例
        neg  : "写负面评论"时，生成的评论确实是负面的比例
        cls  : 判断褒贬的准确率
        stop : 生成的回复里，在 60 个词以内自己结束（生成了 </s>）的比例
        example : 一条"写正面评论"的回复
    """
    pos = P.generate(model, vocab, P.INSTRUCTIONS["pos"], n_gen, seed=1)
    neg = P.generate(model, vocab, P.INSTRUCTIONS["neg"], n_gen, seed=2)
    stop = sum(len(r) < 60 for r in pos + neg) / (2 * n_gen)
    return {"pos": P.follow_rate(pos, True), "neg": P.follow_rate(neg, False),
            "cls": P.classify_accuracy(model, vocab, cls_test), "stop": stop, "example": pos[0]}


def main(quick=False):
    vocab = Vocab()

    # ---------------- 演示 1 ----------------
    print_header("演示 1：对话模板和损失掩码")
    data = P.sft_data(300 if quick else 3000)
    for ins, resp in [data[0], next(d for d in data if d[0][0] == "is")]:
        ids, mask = P.format_example(vocab, ins, resp)
        print("\n    " + " ".join(vocab.itos[i] for i in ids))
        print("    " + " ".join(("^" if m else "·") * len(vocab.itos[i]) for i, m in zip(ids, mask)))
    print("""
    （^ 标出的位置计算损失：回复的每个词和最后的 </s>；· 是提示部分，不计算损失。）

  观察：
    - <user>、<assistant> 是预训练时从没出现过的特殊符号，它们的词向量还是随机初始化的值 ——
      模型要在 SFT 阶段学会"<assistant> 之后该我说话了""说完要输出 </s>"。
    - 只在回复上计算损失：我们要教的是"怎样回答"，而不是"怎样写出用户的问题"。""")

    # ---------------- 演示 2 ----------------
    print_header(f"演示 2：只预训练 vs 没有预训练直接 SFT vs 预训练 + SFT（SFT 数据 {len(data)} 条）")
    n_gen = 50 if quick else 200
    cls_test = make_dataset(100 if quick else 500, seed=77)
    base = U.load_or_pretrain("clm", vocab, quick=quick)
    results = {"只预训练（没有 SFT）": evaluate(base, vocab, cls_test, n_gen)}

    t0 = time.perf_counter()
    print("\n  没有预训练，直接 SFT（随机初始化）：")
    scratch = P.sft_train(U.new_model(vocab, causal=True, seed=0), data, vocab, epochs=1 if quick else 3)
    results["没有预训练，直接 SFT"] = evaluate(scratch, vocab, cls_test, n_gen)
    print("\n  预训练 + SFT：")
    sft = P.sft_train(copy.deepcopy(base), data, vocab, epochs=1 if quick else 3)
    results["预训练 + SFT"] = evaluate(sft, vocab, cls_test, n_gen)
    print(f"\n  （两次 SFT 共用时 {time.perf_counter() - t0:.0f} 秒）")
    if not quick:
        P.CKPT_DIR.mkdir(exist_ok=True)
        torch.save(sft.state_dict(), P.CKPT_DIR / "sft.pt")
        print("  （预训练 + SFT 的模型已保存到 checkpoints/sft.pt）")

    print(f"\n    {'模型':<20s}{'写正面':>8s}{'写负面':>8s}{'判断褒贬':>10s}{'会自己停下':>10s}")
    for name, r in results.items():
        print(f"    {name:<18s}{r['pos']:>10.1%}{r['neg']:>10.1%}{r['cls']:>10.1%}{r['stop']:>12.1%}")
    print("\n  \"write a positive review\" 的一条回复：")
    for name, r in results.items():
        print(f"    {name}：{' '.join(r['example'][:30])}{' ……' if len(r['example']) > 30 else ''}")
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体数字以你运行的结果为准）：
    - 只预训练的模型不认识 <user>、<assistant>，基本是在随便续写：写正面 / 负面的比例都在 50% 附近，
      判断褒贬也接近瞎猜，而且常常停不下来。它"知道"很多关于评论的东西，只是不知道我们在要求它做什么。
    - 预训练 + SFT：几千条示范之后，三个任务都应该做得很好 —— 写出符合要求的评论、正确判断褒贬、说完就停。
    - 没有预训练、直接 SFT：同样的数据，效果应该明显更差（尤其是判断褒贬：它没见过足够多的评论，
      不知道 superb、tedious 是什么意思）。
    这正是"表面对齐假说"（LIMA，Zhou 等 2023）的意思：模型的知识和能力主要来自预训练，
    SFT 主要是教它**格式和行为方式**，所以少量高质量的示范就够了。""")


if __name__ == "__main__":
    main(quick="--quick" in sys.argv)