"""
=============================================================================
第 2 讲 · 代码 4：测试集污染（contamination）——泄漏的考题怎样抬高分数，又怎样被检测出来
=============================================================================

对应讲义：lectures/lecture02_data.md 第 7 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
基准测试（benchmark）和预训练数据都来自互联网：考题被贴在网上、被转载、被收进爬取，
模型训练时就"看过考卷"。这个文件回答四个问题：

    演示 1：第 1 讲的验证集有多"干净"？——验证集里有多少行在训练集里一字不差地出现过，
            base model 在这些行上的 bit/字符和在真正没见过的行上差多少（train–test overlap）。
    演示 2：泄漏的考题能把分数抬高多少？——一个"跨格式问答"基准（60 道题，第 1 讲 base model 约 0.37）：
            把其中一部分题目连同答案、或只有题目、或换一种问法的题目和答案，混进预训练语料，
            用和 base model 完全相同的配置重新预训练，看分数怎样变化。加法题也泄漏一部分。
    演示 3：怎样检测？——GPT-3 式的 n-gram 重叠检测：考题的任何一个 n-gram 出现在训练数据里，就标为"脏"。
            n 取多大？常见短语怎么处理？换一种问法还查得出来吗？
    演示 4：干净子集（clean subset）——只在"没被标脏"的题上算分，能不能还原真实水平？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 04_contamination.py          # 第一次约 60 秒（预训练一个泄漏了考题的模型，3000 步）；之后约 2 秒

需要 PyTorch。依赖本目录的 lm_lab.py 和第 1 讲的 fm_world.py、tiny_gpt.py、console.py（只读导入）。
第 1 讲的 base model 只读地从 code/lecture01/checkpoints/ 加载（不存在时会按第 1 讲的配置训练并保存——
得到的模型和第 1 讲完全相同）。泄漏模型缓存在 code/lecture02/checkpoints/leaked_model.pt。
"""

import math
import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import pad, print_header  # noqa: E402
from fm_world import ALPHABET, addition_line, build_world, qa_prompt  # noqa: E402
from lm_lab import math_accuracy, train_or_load  # noqa: E402
from tiny_gpt import BASE_TRAIN, CharTokenizer, completion_logprob, generate, load_or_train_base_model  # noqa: E402

COPIES = 3          # 泄漏的问答题，每道在语料里出现几次
MATH_COPIES = 5     # 泄漏的加法题，每道出现几次


def answer(p, kind):
    """真实答案：kind = "city" 时是城市，"job" 时是职业。"""
    return p.city if kind == "city" else p.job


def item_text(p, kind):
    """一道考题的完整文本（题目 + 答案），例如 "q: where does kari live? a: varno."——检测时用它。"""
    return f"{qa_prompt(p.name, kind)} {answer(p, kind)}."


def paraphrase(p, kind):
    """同一道题换一种问法（连同答案），例如 "q: which city is kari in? a: varno."。"""
    if kind == "city":
        return f"q: which city is {p.name} in? a: {p.city}."
    return f"q: what job does {p.name} have? a: {p.job}."


def build_benchmark(world):
    """
    "跨格式问答"基准：那 30 个**没有**问答文本的人，每人问城市、职业两道题，共 60 道。
    第 1 讲里这正是"问答·没见过"那两列：考的是"把叙述句里学到的事实用到问答格式上"的能力。

    【返回】dict：组名 -> list[(Person, kind)]，4 组各 15 道：
        clean      : 不泄漏
        leak_qa    : 题目 + 答案原样泄漏（input-and-output contamination）
        leak_input : 只泄漏题目，没有答案（input contamination）
        leak_para  : 换一种问法的题目 + 答案泄漏（检测器很难发现）
    """
    nonqa = [p for p in world.people if p.name not in world.qa_names]
    items = [(p, k) for p in nonqa for k in ("city", "job")]
    random.Random(5).shuffle(items)
    names = ["clean", "leak_qa", "leak_input", "leak_para"]
    return {n: items[15 * i:15 * (i + 1)] for i, n in enumerate(names)}


GROUP_CN = {"clean": "不泄漏", "leak_qa": "泄漏：题目 + 答案", "leak_input": "泄漏：只有题目",
            "leak_para": "泄漏：换个问法 + 答案"}


def leaked_lines(bench, math_leak):
    """
    要混进预训练语料的"泄漏"文本（就像考题被贴到了网上、又被爬了进来）。

    【返回】list[str]：问答题按组的方式各泄漏 COPIES 份；加法题 math_leak 每道泄漏 MATH_COPIES 份
    """
    out = []
    for p, k in bench["leak_qa"]:
        out += [item_text(p, k)] * COPIES
    for p, k in bench["leak_input"]:
        out += [qa_prompt(p.name, k)] * COPIES
    for p, k in bench["leak_para"]:
        out += [paraphrase(p, k)] * COPIES
    for a, b in math_leak:
        out += [addition_line(a, b)] * MATH_COPIES
    return out


def qa_accuracy(model, tok, items):
    """问答题的 exact-match 准确率（贪心生成到句号）。"""
    ok = 0
    for p, k in items:
        out = generate(model, tok, "\n" + qa_prompt(p.name, k), max_new_chars=10, temperature=0, stop=".\n")
        ok += out.strip(" .\n") == answer(p, k)
    return ok / len(items)


# =============================================================================
# 演示 1：第 1 讲的验证集有多干净
# =============================================================================
def demo_overlap(world, tok, base):
    """
    演示 1：第 1 讲验证集里有多少行在训练集里一字不差地出现过，base model 在重叠 / 没见过的行上的 bit/字符。

    【参数】world：第 1 讲默认的世界；tok：CharTokenizer；base：第 1 讲的 base model（只读）
    【返回】无（只打印）
    """
    print_header("演示 1：第 1 讲的验证集有多\"干净\"？—— 验证行在训练集里一字不差出现过的比例")
    train_set = set(world.train_lines)
    print(pad("体裁", 8) + pad("验证行", 8, ">") + pad("在训练集里", 12, ">") + pad("bpc：重叠的行", 16, ">")
          + pad("bpc：没见过的行", 18, ">"))
    tot = {True: [0.0, 0], False: [0.0, 0]}
    for genre in ["bio", "qa", "math", "review", "story"]:
        lines = [l for l in world.val_lines if world.line_genres[l] == genre]
        res = {True: [0.0, 0], False: [0.0, 0]}
        for l in lines:
            bits = -completion_logprob(base, tok, "\n", l + "\n") / math.log(2)
            res[l in train_set][0] += bits
            res[l in train_set][1] += len(l) + 1
            tot[l in train_set][0] += bits
            tot[l in train_set][1] += len(l) + 1
        frac = sum(l in train_set for l in lines) / len(lines)
        fmt = lambda r: f"{r[0] / r[1]:.3f}" if r[1] else "—"
        print(pad(genre, 8) + pad(len(lines), 8, ">") + pad(f"{frac:.0%}", 12, ">") + pad(fmt(res[True]), 16, ">")
              + pad(fmt(res[False]), 18, ">"))
    n_all = len(world.val_lines)
    n_in = sum(l in train_set for l in world.val_lines)
    print(f"\n合计：{n_in} / {n_all} 行（{n_in / n_all:.0%}）在训练集里。按字符平均：重叠的行 "
          f"{tot[True][0] / tot[True][1]:.3f} bit/字符，没见过的行 {tot[False][0] / tot[False][1]:.3f} bit/字符。")
    print("""
观察 Observations：
  · 第 1 讲的验证集有 60%（223 / 369 行）在训练集里一字不差地出现过。人物、问答、故事的验证行 100% 重叠——
    它们都是模板句，种类有限（每个人 5 个模板、故事只有 64 种句子），训练集里早就有；
    只有加法（19%）和评论（34%）有真正"没见过"的行。
  · 合计的数字差别看起来很大（重叠的行 0.290 bit/字符，没见过的行 0.478），但这主要是**体裁构成**造成的：
    没见过的行全是加法和评论。在同一个体裁里比，差别很小（加法 1.854 vs 1.884，评论 0.208 vs 0.214）——
    这个 11 万参数的模型几乎没有"背下"具体的行，它学的是模板和规律。
  · 大模型会背得多得多。Lee et al. (2021)：C4 验证集里 4.6% 的样本在训练集里有近似副本（RealNews 14.35%）；
    Transformer-XL 在 LM1B 验证集上，"训练集里有副本"的样本困惑度 10.11，"没有副本"的 23.58。
    验证集要先对训练集去重，上面的数字才衡量"泛化"而不是"记忆"。""")


# =============================================================================
# 演示 2：泄漏的考题能把分数抬高多少
# =============================================================================
def demo_leak(world, tok, base):
    """
    演示 2：把考题按三种方式混进预训练语料，用和 base model 相同的配置重新预训练（有缓存就读），比较分数。

    【参数】world、tok、base：同演示 1
    【返回】(bench, lines, leaked, math_leak, math_clean)：基准分组、泄漏后的训练行、泄漏模型、两组加法题
    """
    print_header("演示 2：把一部分考题混进预训练语料，用和 base model 完全相同的配置重新预训练")
    bench = build_benchmark(world)
    tp = list(world.test_pairs)
    random.Random(6).shuffle(tp)
    math_leak, math_clean = tp[:50], tp[100:300]
    extra = leaked_lines(bench, math_leak)
    lines = list(world.train_lines) + extra
    random.Random(7).shuffle(lines)
    print(f"第 1 讲的训练语料 {len(world.train_lines):,} 行 + 泄漏的 {len(extra)} 行（占 {len(extra) / len(lines):.1%}）。"
          f"模型配置、步数（{BASE_TRAIN['steps']}）、随机种子都和 base model 相同。")
    leaked, info = train_or_load("leaked_model", tok, lines, BASE_TRAIN["steps"], model_seed=BASE_TRAIN["model_seed"],
                                 data_seed=BASE_TRAIN["data_seed"], lr=BASE_TRAIN["lr"],
                                 batch_size=BASE_TRAIN["batch_size"])
    print("\n泄漏的样子（每种各一行）：")
    for g, (p, k) in [(g, bench[g][0]) for g in ("leak_qa", "leak_input", "leak_para")]:
        shown = {"leak_qa": item_text(p, k), "leak_input": qa_prompt(p.name, k), "leak_para": paraphrase(p, k)}[g]
        print(f"  {pad(GROUP_CN[g], 22)} {shown}")
    print("\n" + pad("考题组", 26) + pad("题数", 6, ">") + pad("base model", 12, ">") + pad("泄漏后的模型", 14, ">"))
    acc = {}
    for g, items in bench.items():
        acc[g] = (qa_accuracy(base, tok, items), qa_accuracy(leaked, tok, items))
        print(pad(GROUP_CN[g], 26) + pad(len(items), 6, ">") + pad(f"{acc[g][0]:.2f}", 12, ">") + pad(f"{acc[g][1]:.2f}", 14, ">"))
    all_items = [it for items in bench.values() for it in items]
    overall = (qa_accuracy(base, tok, all_items), qa_accuracy(leaked, tok, all_items))
    print(pad("整个基准（60 道）", 26) + pad(60, 6, ">") + pad(f"{overall[0]:.2f}", 12, ">") + pad(f"{overall[1]:.2f}", 14, ">"))
    m = {"leak": (math_accuracy(base, tok, math_leak), math_accuracy(leaked, tok, math_leak)),
         "clean": (math_accuracy(base, tok, math_clean), math_accuracy(leaked, tok, math_clean))}
    print(pad(f"加法：泄漏的题（各 {MATH_COPIES} 份）", 26) + pad(50, 6, ">") + pad(f"{m['leak'][0]:.2f}", 12, ">")
          + pad(f"{m['leak'][1]:.2f}", 14, ">"))
    print(pad("加法：没泄漏的题", 26) + pad(200, 6, ">") + pad(f"{m['clean'][0]:.2f}", 12, ">") + pad(f"{m['clean'][1]:.2f}", 14, ">"))
    print("（每组只有 15 道题：一道题就是 0.07。不泄漏的那一组两个模型的差别，可以当作\"换了训练语料带来的随机波动\"的参照。）")
    print("""
观察 Observations：
  · 题目和答案原样泄漏的 15 道题：0.27 → 0.73。整个基准从 0.37 升到 0.47——但模型"把叙述句里的知识
    用到问答上"的能力并没有变（不泄漏的那组 0.33 → 0.33）。分数上涨全部来自背下了考题的答案。
  · 只泄漏题目、没有答案：0.40 → 0.40，没有变化。Dodge et al. (2021) 把污染分成两种：输入和答案都在训练数据里
    （input-and-output contamination，例如 C4 里 XSum 15.49%、TIFU-short 24.88%），以及只有输入在里面
    （input contamination，例如 GLUE 的 QNLI 句子 53.6%）。对考"答案"的基准，前一种危险得多。
  · 换个问法泄漏答案：0.47 → 0.40，没有帮上忙（15 道里少对一道，在噪声范围内）。这个结果并不稳定，
    换一种语料顺序可能不一样（代码 README 的"关于实验结果"有说明）。小模型对"换了说法的同一个事实"
    泛化得很弱（第 1 讲 §4.2）；大模型却可能从换了说法、甚至换了语言的泄漏里学到答案——而这种泄漏最难检测（演示 3）。
  · 加法：每道泄漏 5 份的题 0.12 → 0.18，没泄漏的题 0.07 → 0.04。这个小模型几乎背不住加法题
    （第 1 讲：见过的题也只有 0.12），所以泄漏没把分数抬起来多少。GPT-3 论文里也有类似的情况：
    LAMBADA 有真实的污染，但干净子集的分数和全集相差不到 0.5%。污染会不会抬高分数，取决于模型背不背得下来——
    模型越大，背得越多。""")
    return bench, lines, leaked, math_leak, math_clean


# =============================================================================
# 演示 3：n-gram 重叠检测
# =============================================================================
def ngram_index(lines, n):
    """
    训练语料里每个字符 n-gram 出现在多少行里。

    【参数】lines：训练行；n：n-gram 的长度（字符数）
    【返回】dict：n-gram -> 包含它的行数（同一行里出现多次只算一次）
    【思路】GPT-3 按"词"切 13-gram；我们的模型是字符级的、考题很短，所以按字符切。
    """
    idx = {}
    for line in lines:
        for g in {line[i:i + n] for i in range(len(line) - n + 1)}:
            idx[g] = idx.get(g, 0) + 1
    return idx


def is_dirty(text, idx, n, line_set, max_lines=None):
    """
    一道考题是否被标为"脏"：它的任何一个 n-gram 出现在训练语料里。

    【参数】text：考题文本；idx：ngram_index 的结果；n：n-gram 长度；line_set：训练行的集合
            max_lines：GPT-3 的做法——出现在超过这么多行里的 n-gram 是"常见短语"，忽略它（None 表示不忽略）
    【返回】True / False
    【边界】考题比 n 还短时，检查整道题是否作为一整行出现过（GPT-3：比 13-gram 短的例子整体比较）。
    """
    if len(text) < n:
        return text in line_set
    for i in range(len(text) - n + 1):
        c = idx.get(text[i:i + n], 0)
        if c > 0 and (max_lines is None or c <= max_lines):
            return True
    return False


def demo_detection(bench, lines, math_leak, math_clean):
    """
    演示 3：不同 n、是否忽略常见 n-gram 时，每组考题被标为"脏"的比例。

    【参数】bench：基准分组；lines：泄漏后的训练行；math_leak、math_clean：两组加法题
    【返回】无（只打印）
    """
    print_header("演示 3：n-gram 重叠检测 —— 考题的任何一个 n-gram 出现在训练数据里，就标为\"脏\"")
    line_set = set(lines)
    groups = dict(bench)
    texts = {g: [item_text(p, k) for p, k in items] for g, items in groups.items()}
    texts["math_leak"] = [addition_line(a, b) for a, b in math_leak]
    texts["math_clean"] = [addition_line(a, b) for a, b in math_clean]
    cn = dict(GROUP_CN, math_leak="加法：泄漏的题", math_clean="加法：没泄漏的题")
    configs = [(8, None), (13, None), (13, 10), (20, 10), (30, 10)]
    print("每格是\"被标为脏\"的比例。n = 字符 n-gram 的长度；\"忽略常见\" = 出现在 10 行以上的 n-gram 不算（GPT-3 的做法）。\n")
    header = pad("考题组", 24) + "".join(pad(f"n={n}" + ("，忽略常见" if m else ""), 16, ">") for n, m in configs)
    print(header)
    indexes = {n: ngram_index(lines, n) for n in sorted({n for n, _ in configs})}
    for g, ts in texts.items():
        row = pad(cn[g], 24)
        for n, m in configs:
            rate = sum(is_dirty(t, indexes[n], n, line_set, m) for t in ts) / len(ts)
            row += pad(f"{rate:.0%}", 16, ">")
        print(row)
    p, k = bench["clean"][0]
    t = item_text(p, k)
    hits = sorted({t[i:i + 13] for i in range(len(t) - 12) if indexes[13].get(t[i:i + 13], 0) > 0},
                  key=lambda g: (-indexes[13][g], g))
    print(f"\n一道没泄漏的题：{t!r}")
    print("  它在训练数据里出现过的 13-gram（括号里是出现的行数）：" + "，".join(f"{h!r}（{indexes[13][h]}）" for h in hits))
    print("""
观察 Observations：
  · n 太小（8 或 13 个字符）：几乎每道题都被标为脏，包括完全没泄漏的题。13 个字符只是两三个词，
    "q: what does "、" do? a: tailo" 这样的片段在别人的问答里到处都是——共享的是**格式和常见答案**，不是考题本身。
  · GPT-3 的补救是忽略常见 n-gram（出现在 10 个以上训练文档里的），但 n = 13 时仍有 93% 的干净题被标脏：
    ": what does t"（9 行）这样的片段不算"常见"，却和考题毫无关系。
  · n = 20、忽略常见：干净题 0%，原样泄漏和只泄漏题目 100%——正好分开。但换了问法的泄漏是 0%，检测器看不见它；
    n = 30 连"只泄漏题目"也漏掉了（题目本身不到 30 个字符）。
  · 加法题比 n 还短，就整题比较（GPT-3 对短样本的做法）：泄漏的 100% 被找到，没泄漏的 0%——短而唯一的考题最容易查。
  · n 的选择是假阳性和假阴性之间的权衡。GPT-3 用的是 13 个**词**（按基准里样本长度的第 5 百分位往下调，最少 8 个词）；
    我们的模型是字符级的、考题只有 7 个词左右，所以按字符切，合适的 n 也就不同。""")


# =============================================================================
# 演示 4：干净子集
# =============================================================================
def demo_clean_subset(world, tok, base, leaked, bench, lines):
    """
    演示 4：用 n = 20（忽略常见 n-gram）的检测器把基准分成"脏"和"干净"两部分，分别算两个模型的分数。

    【参数】world、tok、base：同演示 1；leaked：泄漏模型；bench：基准分组；lines：泄漏后的训练行
    【返回】无（只打印）
    """
    print_header("演示 4：干净子集 —— 只在\"没被标脏\"的题上算分（检测器：n = 20，忽略常见 n-gram）")
    line_set = set(lines)
    idx = ngram_index(lines, 20)
    all_items = [(g, it) for g, items in bench.items() for it in items]
    dirty = [is_dirty(item_text(*it), idx, 20, line_set, 10) for _, it in all_items]
    subsets = {"全部 60 道": [it for _, it in all_items],
               "被标为脏的": [it for (_, it), d in zip(all_items, dirty) if d],
               "干净子集（没被标脏的）": [it for (_, it), d in zip(all_items, dirty) if not d]}
    print(pad("子集", 26) + pad("题数", 6, ">") + pad("base model", 12, ">") + pad("泄漏后的模型", 14, ">"))
    for name, items in subsets.items():
        print(pad(name, 26) + pad(len(items), 6, ">") + pad(f"{qa_accuracy(base, tok, items):.2f}", 12, ">")
              + pad(f"{qa_accuracy(leaked, tok, items):.2f}", 14, ">"))
    comp = {}
    for (g, _), d in zip(all_items, dirty):
        comp.setdefault("脏" if d else "干净", {}).setdefault(GROUP_CN[g], 0)
        comp["脏" if d else "干净"][GROUP_CN[g]] += 1
    for key in ("脏", "干净"):
        print(f"  标为{key}的题来自：" + "，".join(f"{g} {c}" for g, c in comp.get(key, {}).items()))
    print("""
观察 Observations：
  · 用 n = 20 的检测器，被标脏的 30 道题恰好是"原样泄漏"和"只泄漏题目"两组。泄漏后的模型在脏题上 0.57、
    在干净子集上 0.37。GPT-3 论文的判据是：干净子集的分数如果比全集低 1–2 个百分点以上，就提示分数被污染抬高了——
    这里低了 0.10。
  · 干净子集上的 0.37，和 base model 在全集上的 0.37 一样：这一次，干净子集还原了真实水平。但这有运气成分——
    换了问法的泄漏留在了"干净子集"里，只是它恰好没帮上忙（演示 2）。如果它帮了忙，干净子集本身就是脏的。
  · 脏题里还混着"只泄漏题目"的 15 道，它们并没有被抬分。GPT-3 论文正是这样：四分之一的基准有超过 50% 的样本被标脏，
    但大多数基准的干净子集和全集差别可以忽略；阅读理解被大量标脏，人工检查发现只是原文在训练数据里，问题和答案不在。
    所以"被标脏的比例"不等于"分数被抬高的幅度"，要比较的是干净子集和全集的分数。
  · 整个分析的前提是检测器够好。它查不出换了说法、换了语言、或者只泄漏了答案的情况——
    这是 2023 年之后污染问题变得更难的原因之一（讲义 §11）。""")


if __name__ == "__main__":
    world = build_world()
    tok = CharTokenizer(ALPHABET)
    base, _ = load_or_train_base_model(world, tok)     # 第 1 讲的 base model：默认世界、默认路径，只读
    demo_overlap(world, tok, base)
    bench, lines, leaked, math_leak, math_clean = demo_leak(world, tok, base)
    demo_detection(bench, lines, math_leak, math_clean)
    demo_clean_subset(world, tok, base, leaked, bench, lines)