"""
=============================================================================
第 2 讲 · 代码 03：标签偏差与校准 —— "先校准，再使用"
=============================================================================

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

【这个脚本要回答的问题】
    Zhao et al.（2021，"Calibrate Before Use"）把少样本提示的不稳定归结为三种偏差：
    多数标签偏差（示例里哪个标签多就偏向哪个）、近因偏差（偏向最后几个示例的标签）、常见词偏差（偏向常见的标签词）。
    他们的修正方法叫**上下文校准**：把一个"无内容"的输入（如 "N/A"）放进同一个提示，看模型偏向哪边，再把这个偏向除掉。
    1. 零样本时，两个迷你 GPT 各自偏向哪个标签？（演示 1）
    2. 用"无内容"输入直接测量：示例的标签多数、标签的顺序，各自怎样推动预测？（演示 2）
    3. "无内容"输入真的没有内容吗？换一个会怎样？（演示 3）
    4. 上下文校准、批量校准（Zhou et al. 2024）在真实的迷你模型上效果如何？（演示 4）
    5. 在 SimICL（"大模型"玩具模型）上复现 Zhao et al. 的图：示例越多、校准前后的均值和波动。（演示 5）

【运行】python3 03_bias_calibration.py     （约 25 秒；第一次运行还要训练 bursty 模型约 45 秒）
【依赖】PyTorch；import 同目录的 icl_lab.py，第 1 讲的 console.py。
"""

import os
import random
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 icl_lab import (CONTENT_FREE, TEMPLATES, VERBALIZERS, SimICL, accuracy, calibrate, demo_pool,  # noqa: E402
                     get_base_model, get_bursty_model, label_probs, mean_probs, mean_positive, predictions,
                     render, render_english, sample_demos, test_set)

N_SETS = 10


def zero_shot_bias(models, tok, test):
    """演示 1：零样本时模型偏向哪边。"""
    print_header("演示 1：零样本的标签偏向——模型把多少条评论判成好评？（测试集里好评占 46.0%）")
    rows = []
    res = {}
    for tname, vname in [("rating", "positive/negative"), ("sentiment", "positive/negative"), ("qa", "positive/negative"),
                         ("rating", "great/awful"), ("sentiment", "great/awful")]:
        t, verb = TEMPLATES[tname], VERBALIZERS[vname]
        row = [f"{tname} · {vname}"]
        for name, m in models.items():
            probs = label_probs(m, tok, [render(t, verb, [], x) for x, _, _ in test], [verb["positive"], verb["negative"]])
            pred_pos = sum(1 for p in predictions(probs) if p == "positive") / len(test)
            res[(name, tname, vname)] = (pred_pos, accuracy(probs, test))
            row += [f"{100 * pred_pos:.1f}%", f"{100 * accuracy(probs, test):.1f}%"]
        rows.append(row)
    print_table(["设计", "base：判成好评", "base：准确率", "bursty：判成好评", "bursty：准确率"], rows)
    bs, us = res[("base", "sentiment", "positive/negative")], res[("bursty", "sentiment", "positive/negative")]
    bg, ug = res[("base", "sentiment", "great/awful")], res[("bursty", "sentiment", "great/awful")]
    br = res[("base", "rating", "great/awful")]
    print_observations([
        "rating · positive/negative 下两个模型都不偏（判成好评的比例正好是真实的 46%）：熟悉的格式里，模型靠输入本身做判断。",
        f"换成陌生的 sentiment 字段：base 只把 {100 * bs[0]:.1f}% 判成好评（准确率 {100 * bs[1]:.1f}%），"
        f"bursty 却几乎不偏（{100 * us[0]:.1f}%，准确率 {100 * us[1]:.1f}%）——bursty 模型把 '<unk> :' 也当成了评分字段。",
        f"sentiment · great/awful 对两个模型都太陌生：判成好评的只有 {100 * bg[0]:.1f}% 和 {100 * ug[0]:.1f}%，"
        f"准确率都在 55% 左右（= 全猜差评）。base 在 rating · great/awful 下也明显偏向 awful（{100 * br[0]:.1f}%）。",
        "这种偏向是**整体的**：不管输入是什么，都往同一个方向推。这正是校准能修正的那种偏差（演示 4）。",
    ])
    return res


def content_free_probe(models, tok, pool):
    """演示 2：无内容输入上的 P(positive)，随示例标签的多少和顺序怎样变。"""
    print_header("演示 2：用无内容输入（<unk>）测偏差——P(positive) 随示例标签怎样变（10 组示例的平均）")
    patterns = ["", "P", "N", "PP", "NN", "PPPP", "NNNN", "PPNN", "NNPP", "PPPN", "NNNP"]
    tname = "sentiment"
    t, verb = TEMPLATES[tname], VERBALIZERS["positive/negative"]
    rows, vals = [], {}
    for pat in patterns:
        row = [pat or "（无示例）", f"{pat.count('P')}:{pat.count('N')}", pat[-1] if pat else "–"]
        for name, m in models.items():
            ps = []
            for s in range(N_SETS):
                demos = sample_demos(pool, random.Random(s), pattern=pat) if pat else []
                ps.append(label_probs(m, tok, [render(t, verb, demos, CONTENT_FREE[0])], ["positive", "negative"])[0][0])
            vals[(name, pat)] = sum(ps) / len(ps)
            row.append(f"{vals[(name, pat)]:.2f}")
        rows.append(row)
    print_table(["示例的标签顺序", "好 : 差", "最后一个", "base", "bursty"], rows)

    def d(name, a, b):
        return vals[(name, a)] - vals[(name, b)]
    print_observations([
        f"多数 vs 近因：PPPN（多数是好评、最后一个是差评）和 NNNP 比，bursty 的 P(positive) 是 {vals[('bursty', 'PPPN')]:.2f} vs "
        f"{vals[('bursty', 'NNNP')]:.2f}，PPNN vs NNPP 是 {vals[('bursty', 'PPNN')]:.2f} vs {vals[('bursty', 'NNPP')]:.2f}："
        "决定方向的是**最后一个标签**，不是多数。bursty 模型有很强的近因偏差。",
        f"标签全是好评（PPPP）和全是差评（NNNN）：bursty {vals[('bursty', 'PPPP')]:.2f} vs {vals[('bursty', 'NNNN')]:.2f}，"
        f"差别很大；base {vals[('base', 'PPPP')]:.2f} vs {vals[('base', 'NNNN')]:.2f}，几乎一样——base 模型对无内容输入的判断"
        f"主要取决于'有没有至少两个示例'（0、1 个示例时 {min(vals[('base', q)] for q in ('', 'P', 'N')):.2f}–"
        f"{max(vals[('base', q)] for q in ('', 'P', 'N')):.2f}，2 个以上 {min(vals[('base', q)] for q in ('PP', 'NN', 'PPPP', 'NNNN')):.2f}–"
        f"{max(vals[('base', q)] for q in ('PPPP', 'NNNN', 'PPNN', 'NNPP', 'PPPN', 'NNNP')):.2f}），而不是示例写了什么标签。",
        f"base 模型对标签顺序的反应弱得多（PPNN − NNPP = {d('base', 'PPNN', 'NNPP'):+.2f}，bursty 是 {d('bursty', 'PPNN', 'NNPP'):+.2f}）："
        "它的训练语料里同一页的评分互相独立，'前面的评分'对它来说本来就不是线索。",
        "这说明近因偏差、多数标签偏差不是 Transformer 天生的毛病，而是模型从预训练数据里学来的'合理推断'：在一份同页评分"
        "相关的语料里，上一条评分确实是下一条评分的好线索。用在少样本提示里，它就成了偏差（讲义 §5.2）。",
    ])
    return vals


def which_content_free(models, tok, test):
    """演示 3：换一个'无内容'输入，测出来的偏差和校准结果都会变。"""
    print_header("演示 3：'无内容'输入真的中性吗？（sentiment · positive/negative，零样本）")
    t, verb = TEMPLATES["sentiment"], VERBALIZERS["positive/negative"]
    candidates = ["<unk>", "the", "the food was", "the soup was <unk>"]
    rows = []
    out = {}
    for name, m in models.items():
        probs = label_probs(m, tok, [render(t, verb, [], x) for x, _, _ in test], ["positive", "negative"])
        for cf in candidates:
            ref = label_probs(m, tok, [render(t, verb, [], cf)], ["positive", "negative"])[0]
            out[(name, cf)] = (ref[0], accuracy(calibrate(probs, ref), test))
            rows.append([name, repr(cf), f"{ref[0]:.3f}", f"{100 * accuracy(probs, test):.1f}%", f"{100 * out[(name, cf)][1]:.1f}%"])
        bc = accuracy(calibrate(probs, mean_probs(probs)), test)
        out[(name, "batch")] = bc
        rows.append([name, "（批量校准：用整批测试输入的平均）", f"{mean_probs(probs)[0]:.3f}", f"{100 * accuracy(probs, test):.1f}%",
                     f"{100 * bc:.1f}%"])
    print_table(["模型", "无内容输入", "它的 P(positive)", "校准前", "校准后"], rows)
    accs = [v[1] for k, v in out.items() if k[0] == "base" and k[1] != "batch"]
    print_observations([
        f"同一个模型、同一个提示，只换'无内容'输入，校准后的准确率就在 {100 * min(accs):.1f}%–{100 * max(accs):.1f}% 之间（base）。"
        "'the food was' 这种半句话，模型会把它补成一条评论来判断，它根本不是'无内容'的。",
        "Zhao et al. 用 'N/A'、空字符串、'[MASK]' 三种无内容输入取平均来减轻这个问题；后来的工作（Fei et al. 2023 的领域上下文校准、"
        "Zhou et al. 2024 的批量校准）指出，这些输入本身也带偏向，于是改用随机的领域内文字，或直接用测试输入的平均。",
        f"批量校准的前提是'测试集的两类大致平衡'：它把平均预测拉回 50:50。bursty 模型本来就几乎不偏，批量校准反而让它从 "
        f"89.0% 掉到 {100 * out[('bursty', 'batch')]:.1f}%；测试集严重不平衡（例如 90% 是差评）时，这一步会出更大的错。",
    ])


def calibration_table(models, tok, test, pool):
    """演示 4：三种做法（不校准 / 上下文校准 / 批量校准）× 多种设计 × 示例个数。"""
    print_header("演示 4：校准在真实的迷你模型上——10 组示例的平均 [最低, 最高]")
    designs = [("sentiment", "positive/negative"), ("sentiment", "great/awful"), ("rating", "great/awful"), ("qa", "positive/negative")]
    rows, tally = [], {"cc": [0, 0, 0], "bc": [0, 0, 0]}
    best = {}
    for name, m in models.items():
        for tname, vname in designs:
            t, verb = TEMPLATES[tname], VERBALIZERS[vname]
            words = [verb["positive"], verb["negative"]]
            for k in [0, 2, 4]:
                sets = [[]] if k == 0 else [sample_demos(pool, random.Random(s), k) for s in range(N_SETS)]
                acc = {"raw": [], "cc": [], "bc": []}
                for demos in sets:
                    probs = label_probs(m, tok, [render(t, verb, demos, x) for x, _, _ in test], words)
                    ref = label_probs(m, tok, [render(t, verb, demos, CONTENT_FREE[0])], words)[0]
                    acc["raw"].append(accuracy(probs, test))
                    acc["cc"].append(accuracy(calibrate(probs, ref), test))
                    acc["bc"].append(accuracy(calibrate(probs, mean_probs(probs)), test))
                means = {key: sum(v) / len(v) for key, v in acc.items()}
                for key in ("cc", "bc"):
                    diff = means[key] - means["raw"]
                    tally[key][0 if diff > 0.005 else (2 if diff < -0.005 else 1)] += 1
                best[(name, tname, vname, k)] = means
                rows.append([name, f"{tname} · {vname}", k] + [
                    f"{100 * means[key]:.1f}%" + ("" if k == 0 else f" [{100 * min(acc[key]):.1f}, {100 * max(acc[key]):.1f}]")
                    for key in ("raw", "cc", "bc")])
    print_table(["模型", "设计", "示例数", "不校准", "上下文校准（<unk>）", "批量校准"], rows)
    print(f"\n  24 种（模型 × 设计 × 示例数）里，比不校准高 / 差不多（±0.5 个百分点）/ 低："
          f"上下文校准 {tally['cc'][0]} / {tally['cc'][1]} / {tally['cc'][2]}，批量校准 {tally['bc'][0]} / {tally['bc'][1]} / {tally['bc'][2]}")
    cc, bc = tally["cc"], tally["bc"]
    rg = best[("base", "rating", "great/awful", 4)]
    sg0 = best[("base", "sentiment", "great/awful", 0)]
    print_observations([
        f"校准救得最多的是 base 模型的 rating · great/awful（4 个示例）：{100 * rg['raw']:.1f}% → 上下文校准 {100 * rg['cc']:.1f}%、"
        f"批量校准 {100 * rg['bc']:.1f}%。模型其实'看懂了'评论，只是整体偏向 awful；除掉这个偏向就对了。",
        f"也有帮倒忙的：base 的 sentiment · great/awful 零样本，上下文校准把 {100 * sg0['raw']:.1f}% 变成 {100 * sg0['cc']:.1f}%"
        "——<unk> 上的概率很极端，除法把它放大成了反方向的偏差（演示 3）。",
        "上下文校准在 10 组示例之间的波动常常比不校准还大（例如 base 的 rating · great/awful，4 个示例：[70.0, 72.0] → "
        "[61.0, 94.0]）：它的参考值来自**一个**无内容输入，本身就有噪声。批量校准用 300 条测试输入的平均做参考，稳定得多。",
        f"总账：上下文校准在 24 种情形里有 {cc[0]} 种变好、{cc[2]} 种变差；批量校准 {bc[0]} 种变好、{bc[2]} 种变差。"
        "校准不是免费的午餐，**要在开发集上验证它对你的模型和提示是否有用**。",
        "Zhao et al. 报告上下文校准让 GPT-3 和 GPT-2 的平均准确率最多提高 30 个百分点、并减小了不同提示之间的波动；"
        "我们的迷你模型上，效果有好有坏——校准修正的是'整体的偏向'，修正不了'没看懂输入'。",
    ])


def sim_curves(test, pool):
    """演示 5：SimICL 上的 Zhao et al. 式曲线。"""
    print_header("演示 5：SimICL（'大模型'玩具模型）——示例数 × 校准，10 组随机示例的平均 / 标准差 / 最低")
    words = ("positive", "negative")
    rng0 = random.Random(0)
    test_sub = rng0.sample(test, 200)
    rows = []
    summary = {}
    for size in ["medium", "large"]:
        sim = SimICL(size)
        for k in [0, 1, 2, 4, 8, 16]:
            acc = {"raw": [], "cc": [], "bc": []}
            for s in range(N_SETS if k else 1):
                rng = random.Random(1000 + s)
                demos = rng.sample(pool.all, k)                          # 随机抽，标签不一定平衡（和 Zhao et al. 一样）
                probs = [sim.label_probs(render_english("review", words, demos, x)) for x, _, _ in test_sub]
                ref = sim.label_probs(render_english("review", words, demos, "N/A"))
                acc["raw"].append(accuracy(probs, test_sub))
                acc["cc"].append(accuracy(calibrate(probs, ref), test_sub))
                acc["bc"].append(accuracy(calibrate(probs, mean_probs(probs)), test_sub))
            cells = []
            for key in ("raw", "cc", "bc"):
                v = acc[key]
                m = sum(v) / len(v)
                sd = (sum((x - m) ** 2 for x in v) / len(v)) ** 0.5
                summary[(size, k, key)] = (m, sd, min(v))
                cells.append(f"{100 * m:.1f}% ± {100 * sd:.1f}（{100 * min(v):.1f}）")
            rows.append([size, k] + cells)
    print_table(["规模", "示例数", "不校准：均值 ± 标准差（最低）", "上下文校准", "批量校准"], rows)
    r0, c0 = summary[("large", 0, "raw")], summary[("large", 0, "cc")]
    r4, c4 = summary[("large", 4, "raw")], summary[("large", 4, "cc")]
    m0, m1 = summary[("medium", 0, "raw")], summary[("medium", 1, "raw")]
    print_observations([
        f"SimICL-large 零样本只有 {100 * r0[0]:.1f}%（模板和常见词两项偏差把它推向一边），上下文校准后 {100 * c0[0]:.1f}%。",
        f"不校准时，4 个示例的 10 组之间标准差 {100 * r4[1]:.1f} 个百分点，最差的一组 {100 * r4[2]:.1f}%："
        "多数标签、近因、顺序三项偏差让'抽到哪几个、什么顺序'变得很要紧——这是 Zhao et al. 在 GPT-3 上看到的样子。"
        f"上下文校准后：均值 {100 * c4[0]:.1f}%、标准差 {100 * c4[1]:.1f}、最低 {100 * c4[2]:.1f}%。",
        f"SimICL-medium 的 1 个示例（{100 * m1[0]:.1f}%）比 0 个（{100 * m0[0]:.1f}%）还差：只有一个示例时，"
        "多数和近因两项偏差全部来自那一个标签，而 medium 从示例里学映射的能力（a_icl）又不够强。",
        "在这个模拟器里，偏差几乎是'加在分数上的常数'（只有易感度 u(x) 一项因输入而异），正是校准能减掉的那种，"
        "所以校准效果很好、但不完美。真实的迷你模型上，偏差的形状复杂得多，校准有好有坏（演示 4）。",
        "**这个演示不能证明任何关于真实模型的事**：结果完全由我们设定的旋钮决定。它的用处是把'偏差是加性的，所以能被减掉'"
        "这个论证变得具体。",
    ])


def main():
    test, pool = test_set(), demo_pool()
    base, tok = get_base_model()
    bursty, _ = get_bursty_model()
    models = {"base": base, "bursty": bursty}
    zero_shot_bias(models, tok, test)
    content_free_probe(models, tok, pool)
    which_content_free(models, tok, test)
    calibration_table(models, tok, test, pool)
    sim_curves(test, pool)


if __name__ == "__main__":
    main()