"""
=============================================================================
第 9 讲 · 代码 2：在无标签的评论上预训练 —— 迷你 GPT 和迷你 BERT 学到了什么？
=============================================================================

对应讲义：lectures/lecture09_pretraining.md 第 2、3、4、6 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
在 4 万条**没有任何标签**的评论（影评 + 餐厅点评，约 100 万个词）上预训练两个迷你 Transformer：

    迷你 GPT ：解码器 + 因果语言模型（预测下一个词）
    迷你 BERT：编码器 + 掩码语言模型（预测被遮住的词）

两个模型结构相同（3 层、64 维、4 个头，约 17 万参数），唯一的区别是有没有因果掩码、用什么目标训练。
训练时没有人告诉它们哪条评论是正面的、哪个词是褒义的 —— 那它们学到了什么？

    演示 1：预训练迷你 GPT；它生成的评论"前后一致"吗？
    演示 2：预训练迷你 BERT；让它做完形填空（包括带否定的句子）
    演示 3：预训练之后的词向量：superb 和谁最像？只看词向量能分出褒义词和贬义词吗？
    演示 4：上下文相关的表示："not boring" 里的 boring 和 "boring" 里的 boring 一样吗？（ELMo 的动机）

训练好的模型保存到 checkpoints/clm.pt 和 checkpoints/mlm.pt，03、04 号文件和第 10 讲会加载它们。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_pretrain.py           # 完整运行（两个模型各训练 4 轮，约 6~10 分钟）
    python 02_pretrain.py --quick   # 冒烟测试（少量数据、1 轮，不保存）

需要 PyTorch。

【关于结果】编写讲义时这个文件只运行了 --quick 冒烟测试（确认代码能跑通），没有完整运行。
    下面各个演示的"预期会看到"描述的是**应该出现的趋势**，具体数字请以你自己运行的结果为准；
    如果你的结果和预期不符，这本身就是一个值得分析的发现（第 8 讲）。
"""

import os

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

import numpy as np
import torch
import torch.nn.functional as F

import lm_utils as U
from toy_reviews import DOMAINS, NEGATIVE, POSITIVE, VERDICT_WORDS, Vocab, sentiment_score

torch.set_num_threads(1)


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


# =============================================================================
# 演示 1：生成的评论是否前后一致
# =============================================================================
def verdict_agreement(reviews):
    """
    在含有总结句的评论里，统计"总结句的倾向"和"其余部分的倾向"一致的比例。

    【思路】总结句（i loved it / a waste of time ……）和前面的形容词是分开生成的，
        但真实语料里两者大多一致（生成器让 75% 的评价句和整体倾向一致）。
        如果语言模型只学会了"局部的语法"，它生成的总结句和前文应该没有关系（一致率约 50%）；
        如果它学会了"整条评论的情感"，一致率应该接近真实语料。

    【参数】reviews：list[list[str]]
    【返回】(一致率, 参与统计的评论条数)

    【边界】只统计"有总结句、并且其余部分的打分不为 0"的评论。
    """
    agree = n = 0
    for words in reviews:
        v = sum(VERDICT_WORDS.get(w, 0.0) for w in words)
        rest = sentiment_score([w for w in words if w not in VERDICT_WORDS])
        if v != 0 and rest != 0:
            n += 1
            agree += (v > 0) == (rest > 0)
    return agree / max(1, n), n


def demo_clm(vocab, corpus, dev, quick):
    print_header(f"演示 1：预训练迷你 GPT（{len(corpus)} 条无标签评论）")
    model = U.new_model(vocab, causal=True)
    print(f"\n    参数 {U.count_parameters(model):,} 个；每轮约 {sum(len(w) + 1 for w in corpus):,} 个训练信号（词）")
    t0 = time.perf_counter()
    U.pretrain(model, corpus, vocab, objective="clm", epochs=1 if quick else 4, dev=dev)
    print(f"\n  预训练用时 {time.perf_counter() - t0:.0f} 秒；开发集困惑度 {np.exp(U.evaluate_lm(model, dev, vocab, 'clm')):.2f}")
    if not quick:
        U.CKPT_DIR.mkdir(exist_ok=True)
        torch.save(model.state_dict(), U.CKPT_DIR / "clm.pt")
        print("  （已保存到 checkpoints/clm.pt）")

    bos = vocab.stoi[Vocab.BOS]
    gen = U.sample(model, [[bos]] * 500, vocab, seed=0)
    gen_words = [[vocab.itos[i] for i in g] for g in gen]
    print("\n  它生成的评论（温度 1.0）：")
    for words in gen_words[:6]:
        print(f"    {' '.join(words)}")
    rate_gen, n_gen = verdict_agreement(gen_words)
    rate_real, n_real = verdict_agreement(dev)
    rng = random.Random(0)
    shuffled = []
    for words in dev:                       # 对照：把总结句随机换成另一条评论的总结句
        other = rng.choice(dev)
        shuffled.append([w for w in words if w not in VERDICT_WORDS] + [w for w in other if w in VERDICT_WORDS])
    rate_shuf, _ = verdict_agreement(shuffled)
    print(f"\n  总结句和前文的情感一致吗？")
    print(f"    真实评论 {rate_real:.1%}（{n_real} 条）；迷你 GPT 生成的 {rate_gen:.1%}（{n_gen} 条）；"
          f"总结句随机打乱的对照 {rate_shuf:.1%}")
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体以你运行的结果为准）：
    - 生成的评论基本合乎语法、用词得当（影评里谈 acting，餐厅点评里谈 food）。
    - 更重要的是：如果模型学到了整条评论的情感，它生成的总结句和前面的评价应该**大体一致**，
      一致率接近真实评论、明显高于随机打乱的对照（约 50%）——
      为了预测好"下一个词"，模型必须知道"前面整体是夸还是骂"。
      没有任何标签，它自己学到了"情感"这个概念。这正是预训练对下游任务有用的原因。""")
    return model


# =============================================================================
# 演示 2：完形填空
# =============================================================================
@torch.no_grad()
def fill_mask(model, vocab, sentence, k=6):
    """
    让编码器预测句子里 <mask> 位置的词，返回概率最高的 k 个 (词, 概率)。

    【参数】sentence：用空格分开的句子，其中恰好有一个 "<mask>"
    """
    model.eval()
    words = sentence.split()
    ids = vocab.encode(words)
    pos = ids.index(vocab.stoi[Vocab.MASK])
    X = torch.tensor([ids])
    logits, _ = model(X, pad_mask=X != 0)
    p = torch.softmax(logits[0, pos], dim=-1)
    p[:len(Vocab.SPECIALS)] = 0                    # 不考虑特殊符号
    top = torch.topk(p, k)
    return [(vocab.itos[i], v) for v, i in zip(top.values.tolist(), top.indices.tolist())]


def demo_mlm(vocab, corpus, dev, quick):
    print_header("演示 2：预训练迷你 BERT（同样的语料，掩码语言模型）")
    model = U.new_model(vocab, causal=False)
    t0 = time.perf_counter()
    U.pretrain(model, corpus, vocab, objective="mlm", epochs=1 if quick else 4, dev=dev)
    print(f"\n  预训练用时 {time.perf_counter() - t0:.0f} 秒")
    if not quick:
        torch.save(model.state_dict(), U.CKPT_DIR / "mlm.pt")
        print("  （已保存到 checkpoints/mlm.pt）")

    print("\n  完形填空（概率最高的 6 个词）：")
    for s in ["the acting was <mask> . i loved it .",
              "the acting was not <mask> . i loved it .",
              "the acting was <mask> . a waste of time .",
              "the food was <mask> . highly recommended .",
              "the plot was great but the music was <mask> ."]:
        top = fill_mask(model, vocab, s)
        print(f"\n    {s}")
        print("      " + "  ".join(f"{w}({p:.2f})" for w, p in top))
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体的词和概率以你运行的结果为准）：
    - 后面跟着 "i loved it"，排在前面的应该是褒义词；跟着 "a waste of time"，应该是贬义词 ——
      模型利用了**右边**的上下文（双向模型的优势）。
    - "not ___ . i loved it"：如果模型学到了"not + 贬义 = 夸奖"，排在前面的应该是**贬义词**。
      这是一个比较难的测试，小模型不一定完全做到；看看你的结果。
    - 餐厅点评的句子里，应该出现餐厅的形容词（delicious、tasty……）。
    - "great but ___"：but 后面通常是相反的评价。""")
    return model


# =============================================================================
# 演示 3：词向量
# =============================================================================
def nearest(emb, vocab, word, k=6):
    """在词向量空间里找和 word 余弦相似度最高的 k 个词（不含特殊符号和它自己）。"""
    E = F.normalize(emb, dim=1)
    sims = E @ E[vocab.stoi[word]]
    sims[:len(Vocab.SPECIALS)] = -1
    sims[vocab.stoi[word]] = -1
    top = torch.topk(sims, k)
    return [(vocab.itos[i], s) for s, i in zip(top.values.tolist(), top.indices.tolist())]


def polarity_accuracy(emb, vocab):
    """
    只看词向量，能不能分出褒义词和贬义词？

    【思路】对每个形容词，留出它自己，用其余褒义词的平均向量和其余贬义词的平均向量当作两个"中心"，
        看它离哪个中心更近（余弦相似度）。全部 48 个形容词的正确率就是"留一法"的准确率。

    【返回】float，0~1（瞎猜是 0.5）
    """
    E = F.normalize(emb, dim=1)
    words = [(w, 1) for w in POSITIVE] + [(w, 0) for w in NEGATIVE]
    right = 0
    for w, y in words:
        pos = torch.stack([E[vocab.stoi[x]] for x in POSITIVE if x != w]).mean(0)
        neg = torch.stack([E[vocab.stoi[x]] for x in NEGATIVE if x != w]).mean(0)
        e = E[vocab.stoi[w]]
        right += int((e @ pos > e @ neg).item()) == y
    return right / len(words)


def demo_embeddings(vocab, clm, mlm):
    print_header("演示 3：预训练之后的词向量")
    random_model = U.new_model(vocab, causal=True, seed=123)
    models = [("迷你 GPT", clm), ("迷你 BERT", mlm), ("随机初始化", random_model)]
    for word in ["superb", "boring", "delicious", "rude"]:
        print(f"\n    和 {word} 最像的词：")
        for name, m in models[:2]:
            print(f"      {name:<8s}" + "  ".join(f"{w}({s:.2f})" for w, s in nearest(m.tok_emb.weight.detach(), vocab, word)))
    print(f"\n    只看词向量分褒贬（48 个形容词，留一法，最近中心）：")
    for name, m in models:
        print(f"      {name:<8s}{polarity_accuracy(m.tok_emb.weight.detach(), vocab):.1%}")
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体以你运行的结果为准）：
    - 预训练之后，褒义词的近邻大多是褒义词、贬义词的近邻大多是贬义词 —— 和第 1、2 讲的 word2vec 一样，
      "出现在相似上下文里的词，向量相似"（分布假设）。
    - 只看词向量，褒贬分类的准确率应该明显高于 50%；随机初始化的词向量只能瞎猜。
    - 这就是预训练帮助下游任务的第一个原因：标注数据里只出现过 great 是褒义词，
      模型也能"猜到" superb、stunning 也是褒义词，因为它们的向量很像。""")


# =============================================================================
# 演示 4：上下文相关的表示
# =============================================================================
def probe_phrases():
    """
    探针用的短句："the {方面} was {形容词} ." 和 "the {方面} was not {形容词} ."
    【返回】list[(词列表, 形容词的下标, 短句的情感 0/1, 形容词)]
    """
    out = []
    aspects = DOMAINS["movie"]["aspects"][:4] + DOMAINS["restaurant"]["aspects"][:4]
    for adj, pol in [(w, 1) for w in POSITIVE] + [(w, 0) for w in NEGATIVE]:
        for asp in aspects:
            out.append((["the", asp, "was", adj, "."], 3, pol, adj))
            out.append((["the", asp, "was", "not", adj, "."], 4, 1 - pol, adj))
    return out


@torch.no_grad()
def features(model, vocab, items, kind):
    """
    取每个短句里形容词的表示。

    【参数】kind："static"（输入词向量，和上下文无关）或 "context"（最后一层的隐藏状态，看过了整个上下文）
    【返回】(N, D) 的张量
    """
    model.eval()
    if kind == "static":
        return torch.stack([model.tok_emb.weight[vocab.stoi[w[i]]] for w, i, _, _ in items])
    X = U.pad_batch([vocab.encode(w) for w, _, _, _ in items])
    _, H = model(X, pad_mask=X != 0)
    return torch.stack([H[b, i + 1] for b, (_, i, _, _) in enumerate(items)])     # +1：前面有 <s>


def train_probe(Xtr, ytr, Xte, yte, steps=500):
    """训练一个逻辑回归探针，返回测试准确率。"""
    torch.manual_seed(0)
    mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6
    Xtr, Xte = (Xtr - mu) / sd, (Xte - mu) / sd
    w = torch.zeros(Xtr.shape[1], requires_grad=True)
    b = torch.zeros(1, requires_grad=True)
    opt = torch.optim.Adam([w, b], lr=0.05)
    for _ in range(steps):
        loss = F.binary_cross_entropy_with_logits(Xtr @ w + b, ytr) + 1e-3 * (w ** 2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return (((Xte @ w + b) > 0).float() == yte).float().mean().item()


def demo_contextual(vocab, clm, mlm):
    print_header("演示 4：上下文相关的表示 —— 同一个词，在不同的句子里")
    items = probe_phrases()
    rng = random.Random(0)
    adjs = POSITIVE + NEGATIVE
    test_adjs = set(rng.sample(POSITIVE, 12) + rng.sample(NEGATIVE, 12))      # 一半形容词只用来测试
    tr = [i for i, it in enumerate(items) if it[3] not in test_adjs]
    te = [i for i, it in enumerate(items) if it[3] in test_adjs]
    y = torch.tensor([float(it[2]) for it in items])
    print(f"""
  探针任务：短句 "the acting was boring ." / "the acting was not boring ." 是褒还是贬？
  特征：只用**形容词那个位置**的向量。用 {len(set(adjs) - test_adjs)} 个形容词的短句训练逻辑回归，
  在另外 {len(test_adjs)} 个**没见过的**形容词上测试（瞎猜 50%；一半的短句有 not）。
""")
    for name, m in [("迷你 GPT", clm), ("迷你 BERT", mlm)]:
        for kind, label in [("static", "输入词向量（静态，和上下文无关）"), ("context", "最后一层的隐藏状态（上下文相关）")]:
            X = features(m, vocab, items, kind)
            acc = train_probe(X[tr], y[tr], X[te], y[te])
            print(f"    {name:<8s}{label:<24s}{acc:>8.1%}")
    print("""
  预期会看到（编写讲义时只做了冒烟测试，具体以你运行的结果为准）：
    - 静态词向量（word2vec、GloVe 和预训练模型的输入层都是这样）：boring 永远是同一个向量，
      不管前面有没有 not。它最多只能判断"这个词本身"的褒贬，一半带 not 的短句必然判错，准确率约 50%。
    - 上下文相关的表示：boring 的向量已经"读过"了前面的 not，所以 "not boring" 里的 boring
      和 "boring" 里的 boring 是不同的向量 —— 线性探针在没见过的形容词上应该也能比较好地判断整个短句的褒贬，
      准确率明显高于静态词向量。
    - 这正是 ELMo（2018）的出发点：一个词的表示应该取决于它所在的句子。
      预训练的语言模型免费提供了这样的表示（讲义第 2 节）。""")


def main(quick=False):
    vocab = Vocab()
    corpus = U.pretrain_corpus(2000 if quick else 40000)
    dev = U.pretrain_corpus(2000, seed=101)
    clm = demo_clm(vocab, corpus, dev, quick)
    mlm = demo_mlm(vocab, corpus, dev, quick)
    demo_embeddings(vocab, clm, mlm)
    demo_contextual(vocab, clm, mlm)


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