"""
=============================================================================
第 1 讲 · 代码 05：在笔记本上预训练一个迷你 GPT，和 n-gram 比一比
=============================================================================

对应讲义：lectures/lecture01_history_of_lms.md §7–§8（预训练范式；GPT-2 / GPT-3）

【这个脚本要回答的问题】
    1. "生成式预训练"到底是在做什么？一个 26 万参数的 GPT 在 23 万词的语料上训练一分钟，损失怎样变化？（演示 1）
    2. 它比 n-gram 好在哪里？按文档类型看困惑度。（演示 2）
    3. 它学到的是"字符串"还是"事实"？知识能不能跨格式使用？能不能反过来问？（演示 3）
    4. 需要看远处的任务（not、but）上，它和 n-gram 的差距有多大？（演示 4）
    5. 同一个模型，换一个采样温度，写出来的东西有什么不同？（演示 5）

【怎么读输出】
    困惑度越低越好；准确率附 95% 置信区间。"7-gram" 是演示 2 里测试困惑度最好的 n-gram。

【运行】python3 05_tiny_gpt.py           （第一次约 70 秒：其中约 50 秒在训练迷你 GPT 并存进 checkpoints/；
                                           之后读缓存，约 20 秒）
        python3 05_tiny_gpt.py --quick   （smoke test：只训练 60 步、不写缓存，几秒钟；数字没有意义）
【依赖】numpy、PyTorch；import 同目录的 tiny_gpt.py、ngram_lm.py、pretrain_world.py、console.py。
"""

import math
import re
import sys

import torch
import torch.nn.functional as F

from console import pct, pct_ci, print_header, print_observations, print_table
from ngram_lm import NGramLM
from pretrain_world import (NEGATIVE, POSITIVE, RELATIONS, build_corpus, corpus_text, fact_probes,
                            question_text, review_probes)
from tiny_gpt import generate, get_model, score_options, top_next

QUICK = "--quick" in sys.argv


# ---------------------------------------------------------------------------
def demo1_training(model, info):
    """演示 1：训练过程。"""
    print_header("演示 1：生成式预训练 = 在一大堆文本上反复练习'预测下一个词'")
    hist = info["history"]
    rows = [[h["step"], f"{h['train_loss']:.3f}", f"{h['eval_loss']:.3f}", f"{math.exp(h['eval_loss']):.2f}"] for h in hist]
    src = "读取缓存" if info["cached"] else "刚刚训练"
    print(f"迷你 GPT（{src}）：{model.n_params():,} 个参数，{model.cfg.n_layer} 层，每层 {model.cfg.n_head} 个头，"
          f"向量维度 {model.cfg.d_model}，上下文窗口 {model.cfg.block_size} 个 token；训练用时 {info['train_seconds']:.0f} 秒")
    print_table(["训练步数", "训练损失", "验证损失", "验证困惑度"], rows)
    print()
    print_table(["模型", "参数", "层数", "向量维度", "上下文", "训练数据"], [
        ["本讲的迷你 GPT", f"{model.n_params() / 1e3:.0f} K", model.cfg.n_layer, model.cfg.d_model,
         model.cfg.block_size, "约 23 万词（玩具语料）"],
        ["GPT-1（2018）", "117 M", 12, 768, 512, "BooksCorpus（约 7000 本书）"],
        ["GPT-2（2019）", "1.5 B", 48, 1600, 1024, "WebText（约 40 GB 网页文本）"],
        ["GPT-3（2020）", "175 B", 96, 12288, 2048, "约 3000 亿 token"],
    ])
    first, last = hist[0], hist[-1]
    print_observations([
        f"损失从随机猜的 ln(290) ≈ {math.log(290):.2f} 一路降到 {last['eval_loss']:.3f}（困惑度 {math.exp(last['eval_loss']):.2f}）："
        "训练只有一个目标——预测下一个词。没有人告诉它什么是首都、什么是褒贬。",
        "验证损失（在没见过的句子上）和训练损失几乎一样，说明模型没有死记训练句子。"
        "（第一版语料只有现在的 1/3，验证损失在第 1000 步之后就开始上升——数据太少时模型会背答案。）",
        "结构上，它和 GPT-3 是同一个东西：一叠因果 Transformer 层 + 预测下一个词。区别全在规模——参数差了约 70 万倍。",
    ])


# ---------------------------------------------------------------------------
def stream_with_kinds(docs, tok):
    """把文档拼成 token 流，同时记下每个 token 属于哪种文档。"""
    ids, kinds = [], []
    for d in docs:
        t = tok.encode("<doc> " + d.text)
        ids += t
        kinds += [d.kind] * len(t)
    return ids, kinds


@torch.no_grad()
def gpt_token_nll(model, ids, batch=64):
    """
    迷你 GPT 在一条 token 流上每个位置的负对数概率（第 0 个位置没有前文，记为 None）。

    【思路】滑动窗口：窗口长 block_size，每次往前挪半个窗口；除了第一个窗口，只给后半段的位置打分，
        这样每个被打分的位置都至少有半个窗口（32 个 token）的前文。
    【返回】长度为 len(ids) 的列表
    """
    T = model.cfg.block_size
    S = T // 2
    nll = [None] * len(ids)
    starts = list(range(0, max(1, len(ids) - T), S))
    x_all = torch.tensor(ids, dtype=torch.long)
    for b0 in range(0, len(starts), batch):
        chunk = starts[b0:b0 + batch]
        xs = torch.stack([x_all[s:s + T] for s in chunk])
        ys = torch.stack([x_all[s + 1:s + T + 1] for s in chunk])
        logits = model(xs)[0]
        lp = -F.log_softmax(logits, dim=-1).gather(-1, ys.unsqueeze(-1)).squeeze(-1)
        for row, s in enumerate(chunk):
            lo = 0 if s == 0 else S
            for j in range(lo, T):
                nll[s + j + 1] = lp[row, j].item()
    return nll


def ngram_token_nll(lm, ids):
    """n-gram 在每个位置的负对数概率。"""
    out = [None]
    for i in range(1, len(ids)):
        out.append(-math.log(lm.prob(ids[i], ids[max(0, i - lm.n + 1):i])))
    return out


def demo2_perplexity(model, tok, world, ngrams):
    """演示 2：按文档类型的测试困惑度。"""
    print_header("演示 2：测试困惑度，按文档类型分开看")
    test_docs = build_corpus(world, seed=1, scale=0.2)
    ids, kinds = stream_with_kinds(test_docs, tok)
    nlls = {f"{n}-gram": ngram_token_nll(lm, ids) for n, lm in ngrams.items()}
    nlls["迷你 GPT"] = gpt_token_nll(model, ids)
    kinds_order = ["article", "faq", "quiz", "reviews", "story", "全部"]
    rows = []
    table = {}
    for name, nl in nlls.items():
        r = [name]
        for k in kinds_order:
            vals = [v for v, kd, g in zip(nl, kinds, nlls["迷你 GPT"]) if v is not None and g is not None
                    and (k == "全部" or kd == k)]
            ppl = math.exp(math.fsum(vals) / len(vals))
            table[(name, k)] = ppl
            r.append(f"{ppl:.3f}")
        rows.append(r)
    print(f"测试语料：{len(test_docs)} 篇新生成的文档、{len(ids):,} 个 token（两种模型在同一批位置上打分）")
    print_table(["模型"] + kinds_order, rows)
    best_ng = min(ngrams, key=lambda n: table[(f"{n}-gram", "全部")])
    gain = {k: table[(f"{best_ng}-gram", k)] / table[("迷你 GPT", k)] for k in kinds_order}
    biggest = max(kinds_order[:-1], key=lambda k: gain[k])
    print_observations([
        f"迷你 GPT 的总体困惑度（{table[('迷你 GPT', '全部')]:.3f}）低于所有 n-gram（最好的 {best_ng}-gram 是 "
        f"{table[(f'{best_ng}-gram', '全部')]:.3f}）。它的'上下文'是整个 64 个 token 的窗口，而且不要求这个上下文原样出现过。",
        f"差距最大的是 {biggest}（{best_ng}-gram 的困惑度是 GPT 的 {gain[biggest]:.2f} 倍）：这类文档里要预测的词依赖远处的信息"
        "（同一篇文章前面提过的国名、故事开头的人物……）。",
        "这份语料的随机性（每个模板随机选词）有一个下限，谁都不可能低于它；所以两者的差距看起来不大，"
        "但在'需要远处信息的那几个词'上差别是决定性的——演示 3、4 专门看这些词。",
    ])
    return best_ng


# ---------------------------------------------------------------------------
def answer_ngram(lm, tok, prompt):
    """n-gram 的贪心下一词。"""
    d = lm.dist(tok.encode(prompt))
    return tok.itos[max(range(len(d)), key=lambda i: (d[i], -i))]


def demo3_knowledge(model, tok, world, ngrams, best_ng):
    """演示 3：知识能否跨格式、跨方向使用。"""
    print_header("演示 3：学到的是'字符串'还是'事实'？——跨格式与反转")
    lm = ngrams[best_ng]
    probes = fact_probes(world)
    gpt_ans = lambda p: top_next(model, tok, p, 1)[0][0]
    ng_ans = lambda p: answer_ngram(lm, tok, p)
    rows = []
    res = {}
    for label, seen in [("问答见过的国家（16 国 × 4 关系）", True), ("只在百科出现的国家（8 国 × 4 关系）", False)]:
        sub = [p for p in probes if p.seen == seen]
        r = [label]
        for name, f in [(f"{best_ng}-gram", ng_ans), ("迷你 GPT", gpt_ans)]:
            ok = sum(f(f"<doc> faq : q : {question_text(p.country, p.relation)} a :") == p.answer for p in sub)
            res[(name, seen)] = (ok, len(sub))
            r.append(pct_ci(ok, len(sub)))
        rows.append(r)
    print("(a) 用问答格式提问：'<doc> faq : q : <问题> a :'")
    print_table(["国家", f"{best_ng}-gram", "迷你 GPT"], rows)

    countries = world.countries
    names = {c.name for c in countries}
    checks = [
        ("首都 · 正着问", "the capital of {c} is", lambda c: c.capital),
        ("首都 · 反着问", "{capital} is the capital of", lambda c: c.name),
        ("领导人 · 正着问", "the leader of {c} is", lambda c: c.leader),
        ("领导人 · 反着问", "{leader} is the leader of", lambda c: c.name),
    ]
    rows = []
    rev = {}
    for label, tpl, gold in checks:
        r = [label, repr(tpl)]
        for name, f in [(f"{best_ng}-gram", ng_ans), ("迷你 GPT", gpt_ans)]:
            preds = [f("<doc> article : " + tpl.format(c=c.name, capital=c.capital, leader=c.leader)) for c in countries]
            ok = sum(p == gold(c) for p, c in zip(preds, countries))
            is_country = sum(p in names for p in preds)
            rev[(name, label)] = ok
            rev[(name, label, "country")] = is_country
            r.append(f"{ok}/24" + (f"（{is_country}/24 次答了某个国名）" if "反着" in label else ""))
        rows.append(r)
    print("\n(b) 用百科的句式提问（全部 24 国）。训练语料里，首都两个方向的说法都有（'the capital of X is C' 和 "
          "'C is the capital of X'），\n    领导人**只有**一个方向（'the leader of X is L'，L 从来不出现在 X 前面）。")
    print_table(["问法", "提示（前面加 '<doc> article :'）", f"{best_ng}-gram", "迷你 GPT"], rows)
    g_held = res[("迷你 GPT", False)]
    print_observations([
        f"{best_ng}-gram 在问答见过的国家上很好，在只在百科出现过的国家上是 {pct(*res[(f'{best_ng}-gram', False)])}："
        "它从没见过 'capital of <这个国家> ? a :' 这串字符。",
        f"迷你 GPT 在这些国家上答对了 {pct(*g_held)}：首都是在百科的句子里学到的，却能用问答的格式取出来。"
        "知识被存成了'关于这个国家的事实'，而不是'某个句子'——这正是提示能起作用的前提：同一个知识可以用很多种问法问出来。",
        f"但**反着问**就不行了：领导人反着问只对了 {rev[('迷你 GPT', '领导人 · 反着问')]}/24（随机猜 ≈ 1/24），"
        f"虽然它 {rev[('迷你 GPT', '领导人 · 反着问', 'country')]}/24 次答的是一个国名（知道'该答一个国家'），却答不对是哪一个。"
        f"首都反着问能答对（{rev[('迷你 GPT', '首都 · 反着问')]}/24），只因为训练语料里本来就有那个方向的句子。",
        "这就是 Berglund et al.（2023）报告的**反转诅咒**（reversal curse）：在 'A 是 B' 上训练的模型，不会自动学会 'B 是 A'。"
        "他们报告 GPT-4 回答'某位名人的母亲是谁'的正确率是 79%，反过来问'这位母亲的儿子是谁'只有 33%。"
        "提问的方向要和知识在训练数据里出现的方向一致。",
    ])


# ---------------------------------------------------------------------------
def demo4_reviews(model, tok, ngrams):
    """演示 4：情感分类（需要看远处的 not、but）。"""
    print_header("演示 4：零样本情感分类——'review : … . rating :' 之后，positive 和 negative 谁的概率大")
    probes = review_probes()
    rows = []
    res = {}
    for name in [f"{n}-gram" for n in ngrams] + ["迷你 GPT"]:
        cnt = {}
        for text, label, kind in probes:
            prompt = f"<doc> reviews : review : {text} . rating :"
            if name == "迷你 GPT":
                pred = score_options(model, tok, prompt, ["positive", "negative"])[2]
            else:
                lm = ngrams[int(name.split("-")[0])]
                ctx = tok.encode(prompt)
                pred = "positive" if lm.prob(tok.stoi["positive"], ctx) >= lm.prob(tok.stoi["negative"], ctx) else "negative"
            ok, n = cnt.get(kind, (0, 0))
            cnt[kind] = (ok + (pred == label), n + 1)
        res[name] = cnt
        rows.append([name] + [pct_ci(*cnt[k]) for k in ["simple", "but", "not"]])
    print(f"{len(probes)} 条新生成的评论；按类型分开报告（括号里是 95% 置信区间）")
    print_table(["模型", "简单句", "but 句", "not 句"], rows)
    print_observations([
        "3-gram 的窗口只有 'rating :' 两个词，什么都看不到，只能按标签的整体频率猜（≈ 50%）。",
        f"5-gram 看得见最后一个形容词，所以简单句和 but 句（以 but 后面为准，决定性的词正好在最后）都对；"
        f"可 not 句只有 {pct(*res['5-gram']['not'])}——'not' 在它的窗口外面，它看到 'cold'，就说 negative。",
        "7-gram 的窗口刚好够到 'not'，三类都对：在这种模板化的玩具语言里，'把窗口开大一点'就够了。"
        "真实语言里决定性的词可以在几十个词之外，n-gram 永远开不了那么大的窗口（演示 2 的稀疏性）。",
        "迷你 GPT 三类全对，而且它的'提示'只是模仿训练语料里评论页面的格式——没有任何指令。"
        "GPT-2 论文用 'TL;DR:' 诱导出摘要能力，是同一个思路：找一个在预训练语料里自然地引出该任务的格式。",
    ])


# ---------------------------------------------------------------------------
def judge_review(text, rating):
    """
    检查一条生成的评论是否"合乎这个世界的语法"，以及评分和内容是否一致。

    【返回】(合语法, 评分一致)
    【思路】按生成器的规则反推：分句是 "the <方面> was [not] <形容词>"，用 and 或 but 连接；
        and 连接的分句极性必须相同；but 前后极性必须相反；整句的极性 = 最后一个分句的极性。
    """
    parts = text.split(" but ")
    if len(parts) > 2:
        return False, False
    pol = []
    for part in parts:
        ps = []
        for cl in part.split(" and "):
            m = re.fullmatch(r"the (\w+) was (not )?(\w+)", cl)
            if not m or m.group(3) not in POSITIVE + NEGATIVE:
                return False, False
            p = 1 if m.group(3) in POSITIVE else -1
            ps.append(-p if m.group(2) else p)
        if len(set(ps)) != 1:
            return False, False
        pol.append(ps[0])
    if len(pol) == 2 and pol[0] == pol[1]:
        return False, False
    return True, rating == ("positive" if pol[-1] > 0 else "negative")


def demo5_temperature(model, tok):
    """演示 5：采样温度。"""
    print_header("演示 5：同一个模型、同一个提示，换一个采样温度")
    prompt = "<doc> reviews : review :"
    rows, examples = [], {}
    n = 60
    for temp in [0.0, 0.5, 1.0, 1.5, 2.0]:
        outs = []
        for seed in range(n if temp > 0 else 1):
            g = generate(model, tok, prompt, max_new_tokens=30, temperature=temp, stop=("<doc>",), seed=seed)
            m = re.match(r"(.*?) \. rating : (\w+) \.", g)
            outs.append((m.group(1), m.group(2)) if m else (g, None))
        if temp == 0.0:
            outs = outs * n                           # 贪心解码：每次都一样
        judged = [judge_review(t, r) if r else (False, False) for t, r in outs]
        distinct = len({t for t, _ in outs})
        rows.append([temp if temp > 0 else "0（贪心）", f"{distinct}/{n}", pct(sum(j[0] for j in judged), n),
                     pct(sum(j[1] for j in judged), n)])
        examples[temp] = outs[:2]
    print(f"提示：{prompt!r}；每个温度采样 {n} 次，只看生成的第一条评论")
    print_table(["温度", "不同的评论数", "合乎语法", "评分与内容一致"], rows)
    print()
    for temp, ex in examples.items():
        for t, r in ex[:1 if temp == 0 else 2]:
            if r is None:
                print(f"  T={temp}：review : {t}    ← 没有写出 '. rating : …' 就结束了")
            else:
                print(f"  T={temp}：review : {t} . rating : {r}")
    print_observations([
        "温度 0（贪心）每次都写同一条评论：语言模型给出的是一个**分布**，贪心解码只取了它的众数。",
        "温度升高，写出的评论越来越多样；但到了 1.5、2.0，开始出现不合语法、或者评分和内容对不上的评论——"
        "低概率的词被抽到的机会变大了。",
        "所以温度是'多样性 vs 可靠性'的旋钮：要稳定、可复现的答案（分类、抽取）用低温度；要多样的候选（头脑风暴、"
        "self-consistency 的多次采样，第 4 讲）用较高的温度。",
    ])


def main():
    torch.set_num_threads(1)
    model, tok, info = get_model("base", quick=QUICK)
    world = info["world"]
    if QUICK:
        print("  （--quick：只训练了 60 步，下面的数字没有意义，只用来检查代码能跑通）")
    demo1_training(model, info)
    ids = tok.encode(corpus_text(build_corpus(world, seed=0)))
    ngrams = {n: NGramLM(n).fit(ids, len(tok)) for n in [3, 5, 7]}
    best_ng = demo2_perplexity(model, tok, world, ngrams)
    demo3_knowledge(model, tok, world, ngrams, best_ng)
    demo4_reviews(model, tok, ngrams)
    demo5_temperature(model, tok)


if __name__ == "__main__":
    main()