"""
=============================================================================
第 9 讲 · 代码 4：零样本分类 —— 不用任何标注数据，让语言模型自己"判断"情感
=============================================================================

对应讲义：lectures/lecture09_pretraining.md 第 8 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
GPT-2 的论文（Radford 等 2019）提出一个想法：语言模型在预训练时已经见过无数"任务"的影子，
不微调、不需要任何标注数据，只要把任务写成"续写"的形式，比较不同续写的概率，就能完成任务。

对情感分类：评论后面接 "i loved it ." 还是 "i hated it ."，哪个概率更高？

    演示 1：零样本分类的准确率 —— 用不同的"答案词"（verbalizer）
    演示 2：校准 —— 减去答案本身的"先验"偏好
    演示 3：和有监督的方法比：零样本相当于多少条标注数据？

为了公平，测试评论里的总结句（i loved it / a waste of time ……）被**删掉**了，
否则模型只要"抄"一下原来的总结句就行了。标签用删掉总结句之后的打分重新计算。

需要 02 号文件保存的 checkpoints/clm.pt（找不到时会先自动预训练，约 3~5 分钟）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 04_zero_shot.py           # 完整运行（约 1 分钟，不含预训练）
    python 04_zero_shot.py --quick   # 冒烟测试

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

import os

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

import numpy as np
import torch

import lm_utils as U
from toy_reviews import VERDICT_WORDS, VERDICTS, Vocab, make_dataset, sentiment_score

base = importlib.import_module("01_baselines_and_metrics")
torch.set_num_threads(1)


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


def strip_verdict(words):
    """
    删掉评论里的总结句（含有 loved、waste 等关键词的那一句）。

    【返回】(删掉之后的词列表, 新的标签)；如果删掉之后打分为 0（无法判断），返回 None
    """
    sents, cur = [], []
    for w in words:
        cur.append(w)
        if w == ".":
            sents.append(cur)
            cur = []
    if cur:
        sents.append(cur)
    kept = [w for s in sents if not any(x in VERDICT_WORDS for x in s) for w in s]
    score = sentiment_score(kept)
    if not kept or score == 0:
        return None
    return kept, int(score > 0)


@torch.no_grad()
def verbalizer_scores(model, vocab, reviews, verbalizer, batch_size=128):
    """
    对每条评论，计算 log P(答案句 | 评论)。

    【参数】
        reviews    : list[list[str]]
        verbalizer : 答案句（词列表），例如 ["i", "loved", "it", "."]
    【返回】np.ndarray (N,)

    【思路】前缀 = [<s>] + 评论，续写 = 答案句，用 lm_utils.sequence_logprob 计算续写的对数概率。
    """
    cont = vocab.encode(verbalizer, add_bos=False, add_eos=False)
    out = []
    for start in range(0, len(reviews), batch_size):
        chunk = reviews[start:start + batch_size]
        prefixes = [vocab.encode(r, add_eos=False) for r in chunk]
        out.append(U.sequence_logprob(model, prefixes, [cont] * len(chunk)).numpy())
    return np.concatenate(out)


def zero_shot(model, vocab, reviews, pos_verbs, neg_verbs, calibrate=False):
    """
    零样本分类：比较"正面答案句"和"负面答案句"的平均对数概率。

    【参数】
        pos_verbs, neg_verbs : 正面 / 负面答案句的列表（每个是词列表），取平均
        calibrate            : True = 每个答案句先减去它在"空评论"后面的对数概率（它本身有多常见）

    【返回】np.ndarray (N,) 的预测（1 = 正面）

    【思路】（讲义第 8.2 节）有些答案句本身就更常见（比如模型可能天生更爱说 "i loved it"），
        不管评论是什么，它的概率都偏高，导致预测偏向一边。
        校准：log P(答案 | 评论) − log P(答案 | 空评论)，衡量的是"评论让这个答案的概率提高了多少"。
        （Zhao 等 2021 的 contextual calibration 用的是类似的想法。）
    """
    def avg(verbs):
        s = np.mean([verbalizer_scores(model, vocab, reviews, v) for v in verbs], axis=0)
        if calibrate:
            s -= np.mean([verbalizer_scores(model, vocab, [[]], v)[0] for v in verbs])
        return s
    return (avg(pos_verbs) > avg(neg_verbs)).astype(int)


def main(quick=False):
    vocab = Vocab()
    model = U.load_or_pretrain("clm", vocab, quick=quick)
    raw = make_dataset(300 if quick else 2000, seed=62)
    data = [x for x in (strip_verdict(w) for w, _ in raw) if x is not None]
    reviews, y = [w for w, _ in data], np.array([l for _, l in data])
    print(f"\n  测试集：{len(data)} 条影评（已删掉总结句），正面占 {y.mean():.1%}")
    print(f"  例子：{' '.join(reviews[0])}  → {'正面' if y[0] else '负面'}")

    # ---------------- 演示 1 ----------------
    print_header("演示 1：零样本分类 —— 评论后面接哪句话的概率更高？")
    pairs = [(["i", "loved", "it", "."], ["i", "hated", "it", "."]),
             (["highly", "recommended", "."], ["avoid", "it", "."]),
             (["a", "must", "see", "."], ["a", "waste", "of", "time", "."])]
    print(f"\n    {'正面答案句 vs 负面答案句':<44s}{'准确率':>8s}{'预测为正面的比例':>14s}")
    for p, n in pairs:
        pred = zero_shot(model, vocab, reviews, [p], [n])
        print(f"    {' '.join(p) + '  vs  ' + ' '.join(n):<40s}{(pred == y).mean():>10.1%}{pred.mean():>14.1%}")
    pred_all = zero_shot(model, vocab, reviews, VERDICTS[1], VERDICTS[-1])
    print(f"    {'5 句正面 vs 5 句负面（取平均）':<36s}{(pred_all == y).mean():>10.1%}{pred_all.mean():>14.1%}")
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体数字以你运行的结果为准）：
    - 没有用任何一条标注数据，准确率也应该明显高于 50%：预训练时模型学到了"前面夸 → 后面说 i loved it"。
    - 不同的答案句，准确率可能差别不小 —— 零样本 / 提示的结果对"怎么问"很敏感（第 11、14 讲）。
    - 注意"预测为正面的比例"：如果明显偏离测试集的真实比例，说明某个答案句本身就更"常见"，模型有偏好。""")

    # ---------------- 演示 2 ----------------
    print_header("演示 2：校准 —— 减去答案句本身的先验")
    for name, (pv, nv) in [("loved vs hated", ([pairs[0][0]], [pairs[0][1]])), ("5 句 vs 5 句", (VERDICTS[1], VERDICTS[-1]))]:
        raw_pred = zero_shot(model, vocab, reviews, pv, nv)
        cal_pred = zero_shot(model, vocab, reviews, pv, nv, calibrate=True)
        print(f"\n    {name}：不校准 {(raw_pred == y).mean():.1%}（预测为正面 {raw_pred.mean():.0%}）"
              f" → 校准后 {(cal_pred == y).mean():.1%}（预测为正面 {cal_pred.mean():.0%}）")
    print("""
  预期会看到：如果不校准时预测明显偏向某一边，校准之后预测的比例应该更接近真实比例，准确率也可能提高；
    如果本来就不怎么偏，校准的作用就不大。（这个技巧在真实的大模型上效果很明显，见 Zhao 等 2021。）""")

    # ---------------- 演示 3 ----------------
    print_header("演示 3：零样本相当于多少条标注数据？")
    pool = make_dataset(2000, seed=60)
    stripped_pool = [x for x in (strip_verdict(w) for w, _ in pool) if x is not None]
    print(f"\n    （逻辑回归也在删掉了总结句的评论上训练和测试）")
    print(f"    零样本（迷你 GPT，5 句 vs 5 句）：{(pred_all == y).mean():.1%}")
    for n in ([20, 50] if quick else [10, 20, 50, 100, 200, 500]):
        accs = []
        for seed in range(3):
            order = np.random.default_rng(seed).permutation(len(stripped_pool))
            train = [stripped_pool[i] for i in order[:n]]
            if len({l for _, l in train}) < 2:
                continue
            predict = base.train_lr(train, n_max=2)
            accs.append((predict(reviews) == y).mean())
        print(f"    逻辑回归，{n:>4d} 条标注数据（3 次的平均）：{np.mean(accs):.1%}")
    print("""
  预期会看到：零样本的准确率大致相当于用几十到几百条标注数据训练的逻辑回归（具体以你的结果为准）。
    在我们这个只有 17 万参数、只见过评论的迷你模型上，这已经说明了思路是对的。
    真实的大模型（GPT-3 及以后）规模大得多、见过的文本也丰富得多，零样本和少样本的能力强得多；
    而且出现了"上下文学习"：在提示里放几个例子，模型就能照着做（讲义第 8.3 节）。
    我们的迷你模型只见过一条条孤立的评论，从没见过"例子 + 例子 + 问题"这样的文本，不会有这种能力。""")


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