"""
=============================================================================
第 9 讲 · 代码 4：DPO 的变体 —— IPO、SimPO、KTO 各改了哪个假设
=============================================================================

对应讲义：lectures/lecture09_dpo.md 第 6 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
DPO 之后出现了大量变体。它们大多是对 DPO 的某个假设"动手术"：

    IPO  ：不假设 BT（"概率比越大越好"），改成"概率比等于一个有限的目标"   → 确定性偏好下不发散（代码 03 演示 2）
    SimPO：去掉参考模型，用"按长度平均的对数概率"当奖励，再要求一个间隔 γ → 省掉参考模型、减少长度偏差的初衷
    KTO  ：不需要成对的数据，只需要"这个回答好 / 不好"的标签               → 数据更容易收集
    （ORPO 把偏好损失和 SFT 损失合在一起、也不用参考模型；讲义里讨论，这里不实现。）

    演示 1：同样的 2000 对数据，四种损失训练 5 / 20 / 60 遍后的 KL、真实质量、长度、华丽。
    演示 2：没有参考模型意味着什么？—— SimPO 训练越久，离 SFT 越远。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_variants.py        # 约 25 秒（CPU）

需要 numpy 和 PyTorch。依赖本目录的 dpo_tools.py。
"""

import numpy as np

from dpo_tools import LOSSES, all_logprobs, exact_stats, make_reference, make_reward_model, pad, print_header, \
    train_preference

SETTINGS = [("dpo", 1.0, 0.0, "β = 1"), ("ipo", 0.1, 0.0, "τ = 0.1"), ("kto", 1.0, 0.0, "β = 1"),
            ("simpo", 2.0, 0.5, "β = 2, γ = 0.5")]


def demo_variants(ref, ref_lp, data):
    print_header("演示 1：四种损失，同样的数据")
    chosen, rejected = data
    print("同样的 2000 对偏好（众包标注，有噪声）。每格是“KL / 真实质量”，括号里是平均长度、平均华丽个数：\n")
    checkpoints = [5, 20, 60]
    print(f"  {pad('损失（参数）', 24)}" + "".join(pad(f"{c} 遍", 26, "right") for c in checkpoints))
    final = {}
    for kind, beta, gamma, desc in SETTINGS:
        rec = {}

        def cb(ep, pol, rec=rec):
            if ep in checkpoints:
                s = exact_stats(all_logprobs(pol), ref_lp)
                rec[ep] = (f"{s['kl']:.2f} / {s['gold']:+.2f} ({s['length']:.1f}, {s['fancy']:.1f})", s)
        train_preference(kind, ref, chosen, rejected, beta=beta, gamma=gamma, epochs=checkpoints[-1], callback=cb)
        final[kind] = rec[checkpoints[-1]][1]
        print(f"  {pad(LOSSES[kind] + '（' + desc + '）', 24)}" + "".join(pad(rec[c][0], 26, "right") for c in checkpoints))
    s0 = exact_stats(ref_lp, ref_lp)
    print(f"\n  参照：SFT 策略 KL 0 / {s0['gold']:+.2f}（{s0['length']:.1f}, {s0['fancy']:.1f}）")
    print("""
观察：
  · DPO、IPO、KTO 的结果相近：都离 SFT 一段距离，真实质量都提高了，都学到了一些“华丽”（标注偏差）。
    参数选得合适时，“哪种损失”远没有“什么数据”重要 —— 它们学的都是同一批偏好。
  · KTO 只用了“好 / 不好”的单个标签（把被选中的当好、被拒绝的当坏，丢掉了“谁和谁比”的配对信息），
    效果并不差。在“点赞 / 点踩”这类数据更容易得到的场景里，这是一个实际的优势（Ethayarajh 等 2024）。
  · SimPO 没有参考模型：它离 SFT 的 KL 最大，而且随训练一直增长；回答更长、华丽更多。
    它的长度归一化本意是减少长度偏差，但这里标注员本身偏爱长回答，而没有参考模型“拉住”，
    策略反而把偏差放大了（演示 2）。
  · 【边界】每种方法都只试了一组参数，结论只适用于这个玩具世界。真实世界的比较（例如 SimPO 论文
    在 AlpacaEval 上的结果）依赖于数据、模型和调参，不同论文的结论并不一致。""")
    return final


def demo_no_reference(ref, ref_lp, data):
    print_header("演示 2：没有参考模型意味着什么？")
    chosen, rejected = data
    print("SimPO（β = 2、γ = 0.5）和 DPO（β = 1），训练更久时离 SFT 有多远：\n")
    checkpoints = [10, 40, 80, 120]
    print(f"  {pad('损失', 8)}" + "".join(pad(f"{c} 遍：KL / 长度", 20, "right") for c in checkpoints))
    for kind, beta, gamma in [("dpo", 1.0, 0.0), ("simpo", 2.0, 0.5)]:
        rec = {}

        def cb(ep, pol, rec=rec):
            if ep in checkpoints:
                s = exact_stats(all_logprobs(pol), ref_lp)
                rec[ep] = f"{s['kl']:.2f} / {s['length']:.2f}"
        train_preference(kind, ref, chosen, rejected, beta=beta, gamma=gamma, epochs=checkpoints[-1], callback=cb)
        print(f"  {pad(LOSSES[kind], 8)}" + "".join(pad(rec[c], 20, "right") for c in checkpoints))
    print("""
观察：
  · DPO 的损失里有 log π_ref：“隐式奖励”是**相对于 SFT** 的变化。离 SFT 越远，同样的偏好就越容易满足，
    梯度权重 σ(·) 越小 —— 参考模型像一根橡皮筋把策略拉住。
  · SimPO 的奖励是策略自己的（按长度平均的）对数概率，和 SFT 无关；只要被选中的回答比被拒绝的回答
    概率高出 γ，损失就满意 —— 但这个“满意”可以在离 SFT 任意远的地方实现。
    去掉参考模型省了显存和计算（不用再跑一个参考模型），代价是少了这根橡皮筋，需要靠学习率、训练步数来控制。
  · ORPO（Hong 等 2024）也不用参考模型：它在 SFT 损失上加一个“被选中回答的几率比被拒绝的高”的项，
    靠 SFT 损失把策略留在示范附近。""")


def main():
    ref = make_reference()
    ref_lp = all_logprobs(ref)
    _, _, data = make_reward_model(ref)
    demo_variants(ref, ref_lp, data)
    demo_no_reference(ref, ref_lp, data)


if __name__ == "__main__":
    main()