"""
=============================================================================
第 3 讲 · 代码 3：四种注意力掩码，以及"一行文本"在不同目标下变成什么训练样本
=============================================================================

对应讲义：lectures/lecture03_architectures_objectives.md 第 4、5、6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
BERT、GPT、T5、prefix LM 的区别，归根到底只有两件事：**谁能看见谁（注意力掩码）**，以及**要预测什么（训练目标）**。
这个文件不训练任何模型，只把这两件事画出来、数出来：

    演示 1：四种注意力模式（encoder-only 双向、decoder-only 因果、prefix LM、encoder-decoder + 交叉注意力），
            画成 ASCII 网格。
    演示 2：迷你互联网里的一行评论，在 7 种目标下分别变成什么样的训练样本：
            因果 LM（GPT）、masked LM（BERT 的 80/10/10）、span corruption（T5）、prefix LM、
            BART 的文本填充、ELECTRA 的替换 token 检测、UL2 的三种去噪器 + 模式 token。
    演示 3：每种目标"每 1000 个输入 token 提供多少个训练信号"——这是理解样本效率（ELECTRA、第 4 讲）的关键。
    演示 4：为什么代码 04、05 用"词级"token：在字符级模型上做 masked LM，遮住单个字符太容易，
            遮住整个词又会碰上"每个位置独立预测"的问题——拼不出一个完整的词。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_masks_and_objectives.py          # 约 2 秒

需要 PyTorch（只用来调用 arch_tools.py 里真正用于训练的 batch 函数，保证这里展示的和代码 04、05 训练的完全一样）。
"""

import os
import random
import sys

import torch

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

from arch_tools import (IGNORE, MASK, SENTINELS, Corpus, WordTokenizer, causal_mask, clm_batch,   # noqa: E402
                        corrupt_spans, full_mask, mlm_batch, prefix_batch, prefix_mask, span_batch, t5_batch)
from console import pad, print_header                                                          # noqa: E402
from fm_world import FOOD_NEG, FOOD_POS, build_world                                           # noqa: E402


def grid(mask, rows, cols, title):
    """
    把一个布尔掩码画成网格：# = 可以看，. = 看不到。

    【参数】mask：(len(rows), len(cols)) 的布尔张量；rows / cols：行、列的 token 文字
    """
    w = max(len(c) for c in cols) + 1
    print(f"  {title}")
    print("  " + " " * 7 + "".join(pad(c, w, ">") for c in cols))
    for i, r in enumerate(rows):
        print("  " + pad(r, 7, ">") + "".join(pad("#" if mask[i][j] else ".", w, ">") for j in range(len(cols))))
    print()


# =============================================================================
# 演示 1：四种注意力模式
# =============================================================================
def demo_masks():
    print_header("演示 1：四种注意力模式（行 = 当前位置，列 = 被看的位置；# = 可以看）")
    src = ["37", "+", "45", "="]
    tgt = ["82", "."]
    seq = src + tgt
    T = len(seq)
    grid(full_mask(T).tolist(), seq, seq, "(a) encoder-only（BERT）：双向，每个位置看所有位置")
    grid(causal_mask(T).tolist(), seq, seq, "(b) decoder-only（GPT）：因果，只看自己和左边")
    grid(prefix_mask(T, len(src)).tolist(), seq, seq,
         f"(c) prefix LM：前缀 {''.join(src)} 内部双向，之后因果")
    dec_in = ["[S]", "82", "."]
    print("  (d) encoder-decoder（T5、BART）：三块注意力")
    grid(full_mask(len(src)).tolist(), src, src, "    编码器自注意力：双向")
    grid(causal_mask(len(dec_in)).tolist(), dec_in, dec_in, "    解码器自注意力：因果（[S] = 开始写）")
    grid([[True] * len(src) for _ in dec_in], dec_in, src, "    交叉注意力：解码器的每个位置看编码器的所有位置")
    print("""观察 Observations：
  · (a)(b)(c) 是同一个 6×6 的矩阵，只是填法不同。prefix LM 正好介于两者之间：前缀部分像 BERT，后面像 GPT。
    前缀长度为 1 时就是 GPT，前缀等于整段时就是 BERT。
  · (d) 把"读"和"写"分给两套参数：编码器只管读（双向），解码器只管写（因果），中间用交叉注意力连接。
    它和 (c) 的信息流几乎一样——输出能看到全部输入、只能看到自己左边的输出——区别是参数是否共享，
    以及解码器看的是编码器**最后一层**的输出（(c) 里每一层的输出位置看的是同一层的前缀表示）。
  · 掩码决定了模型"能做什么"：(b)(c)(d) 都能一个一个地生成输出；(a) 没有"左边的"概念，
    它的每个位置都看到了整段输入，只能给已有的位置打分或填空。""")


# =============================================================================
# 演示 2：一行文本 → 7 种训练样本
# =============================================================================
def show_ids(tok, ids):
    """把编号列表显示成用空格隔开的 token（换行显示成 ⏎）。"""
    return " ".join(tok.vocab[i].replace("\n", "⏎") for i in ids)


def demo_examples(world, tok):
    print_header("演示 2：同一行评论，在 7 种目标下变成什么训练样本")
    line = "review: the soup was tasty and the host was kind. rating: good."
    ids = tok.encode(line)
    n = len(ids)
    print(f"原文（{n} 个词级 token）：{show_ids(tok, ids)}\n")
    rng = random.Random(3)

    print("(1) 因果 LM（GPT）：每个位置预测下一个 token，全部 {} 个位置都有损失".format(n - 1))
    for t in [3, 6, 12]:
        print(f"      看到 [{show_ids(tok, ids[:t + 1])}] → 预测 {tok.vocab[ids[t + 1]]!r}")

    print("\n(2) masked LM（BERT）：选 15% 的位置；被选中的 80% 换成 [M]、10% 换成随机词、10% 不变；只在选中处算损失")
    sel = [3, 6, 13]                        # 为了展示，固定选这 3 个位置（约 15%），并让它们分别落在三种处理里
    x = list(ids)
    x[sel[0]] = tok.mask_id
    x[sel[1]] = tok.stoi["pilot"]           # 随机替换（10%）
    # sel[2] 保持不变（10%）
    print(f"      输入：{show_ids(tok, x)}")
    print("      要预测：" + "，".join(f"位置 {t} → {tok.vocab[ids[t]]!r}" for t in sel)
          + "（第 2 个被换成了随机词 'pilot'，第 3 个保持原样，模型不知道哪些位置被动过）")

    print("\n(3) span corruption（T5）：挖掉若干片段，每个片段换成一个哨兵；目标 = 哨兵 + 被挖掉的内容")
    selected = [i in (3, 4, 10) for i in range(n)]     # 挖掉 "soup was" 和 "kind"
    c, tgt = corrupt_spans(ids, selected, tok)
    print(f"      编码器输入：{show_ids(tok, c)}")
    print(f"      解码器目标：{show_ids(tok, tgt)}")

    P = 9
    print(f"\n(4) prefix LM：前 {P} 个 token 是前缀（双向可见），之后像 GPT 一样逐个预测")
    print(f"      前缀：{show_ids(tok, ids[:P])}")
    print(f"      目标：{show_ids(tok, ids[P:])}")

    print("\n(5) BART 的文本填充（text infilling）：片段长度服从 Poisson(3)，每个片段只换成**一个** [M]；"
          "解码器要写出**完整的原文**")
    bart_in = ids[:4] + [tok.mask_id] + ids[7:10] + [tok.mask_id] + ids[12:]
    print(f"      编码器输入：{show_ids(tok, bart_in)}")
    print(f"      解码器目标：{show_ids(tok, ids)}（模型还得自己判断每个 [M] 里缺了几个词）")

    print("\n(6) ELECTRA：一个小生成器把 15% 的位置换成\"看起来合理\"的词；判别器对**每个**位置判断：原词还是被换过？")
    e = list(ids)
    picked = [5, 8, 14]                        # 生成器在这 3 个位置采样：
    e[5] = tok.stoi["bland"]                   #   tasty → bland：语法上完全合理，但和后面的 good 矛盾
    e[8] = tok.stoi["staff"]                   #   host → staff：合理，意思也没变，但仍然算"被换过"
    e[14] = tok.stoi["good"]                   #   good → good：生成器碰巧采到了原词，标签算"原词"
    print(f"      输入：{show_ids(tok, e)}")
    print("      生成器动过的位置：" + "，".join(f"{tok.vocab[ids[i]]} → {tok.vocab[e[i]]}" for i in picked))
    labels = ["换" if e[i] != ids[i] else "原" for i in range(n)]
    print(f"      判别器的标签（每个位置一个）：{' '.join(labels)}")
    print(f"      （{n} 个位置全部有标签，其中 {labels.count('换')} 个是\"换\"。要发现 bland 被换过，"
          "判别器得懂\"评分由形容词决定\"）")

    print("\n(7) UL2 的混合去噪器（mixture of denoisers）：同一个模型轮流做三类任务，输入开头放一个模式 token")
    two = tok.encode("\n" + world.train_lines[2] + "\n" + world.train_lines[3])
    m = len(two)
    r_sel = [i in (5, 6, 15) for i in range(m)]
    c_r, t_r = corrupt_spans(two, r_sel, tok)
    print(f"      [R] 常规去噪（短片段、15%）  输入：[R] {show_ids(tok, c_r)}")
    print(f"                                  目标：{show_ids(tok, t_r)}")
    half = m // 2
    print(f"      [S] 顺序去噪（= prefix LM）  输入：[S] {show_ids(tok, two[:half])}")
    print(f"                                  目标：{show_ids(tok, two[half:])}")
    x_sel = [3 <= i < 9 or 12 <= i < 19 for i in range(m)]
    c_x, t_x = corrupt_spans(two, x_sel, tok)
    print(f"      [X] 极端去噪（长片段、约 50%）输入：[X] {show_ids(tok, c_x)}")
    print(f"                                  目标：{show_ids(tok, t_x)}")
    print("""
观察 Observations：
  · 同一行文本，GPT 从中得到 {} 道"预测下一个词"的题；BERT 只得到被选中的那 2–3 道；T5 得到几道"写出片段"的题；
    ELECTRA 得到 {} 道"这个词被换过吗"的判断题——每个位置都有标签。
  · BERT 的 80/10/10：如果被选中的词总是变成 [M]，模型就只需要在 [M] 处认真"思考"，其他位置可以偷懒；
    而下游使用时输入里根本没有 [M]。随机替换和保持原样逼着模型对**每个**位置都保持一个好的表示。
  · T5 和 BART 都是"去噪"，但格式不同：T5 的目标只有被挖掉的片段（短，省计算）；BART 要重写全文，
    而且一个 [M] 可能代表 0 个、1 个或好几个词（它还用了句子打乱等别的噪声）。
  · UL2 的 [R] 和 T5 的 span corruption 相同；[S] 就是 prefix LM；[X] 挖得又长又多，要求模型写出大段内容，
    更接近生成。论文把 causal LM 也看成一种极端情形：把整段都当作要"恢复"的片段。
    模式 token 在下游使用时也加在输入前面，告诉模型"现在做的是哪一类任务"（mode switching）。""".format(n - 1, n))


# =============================================================================
# 演示 3：每 1000 个输入 token 提供多少训练信号
# =============================================================================
def demo_signal(world, tok):
    print_header("演示 3：同样的输入长度，每种目标提供多少个\"要预测的 token\"")
    corpus = Corpus(world.train_lines, tok)
    g = torch.Generator().manual_seed(0)
    T, B, reps = 32, 32, 50
    rows = []
    for name, fn in [("因果 LM（GPT）", lambda: clm_batch(corpus, B, T, g)),
                     ("prefix LM", lambda: prefix_batch(corpus, B, T, g)),
                     ("masked LM（BERT，15%）", lambda: mlm_batch(corpus, B, T, g)),
                     ("span corruption（单塔，15%）", lambda: span_batch(corpus, B, T, g))]:
        n_pos = n_loss = n_sent = 0
        for _ in range(reps):
            x, m, y = fn()
            n_pos += y.numel()
            n_loss += int((y != IGNORE).sum())
            n_sent += int(((y >= tok.sentinel_ids[0]) & (y != IGNORE)).sum())
        rows.append((name, n_pos, n_loss, n_sent))
    n_pos = n_loss = n_sent = 0
    for _ in range(reps):
        src, vis, tin, y = t5_batch(corpus, B, g)
        n_pos += int(vis.sum()) + int((y != IGNORE).sum())     # 编码器的真实 token + 解码器的 token
        n_loss += int((y != IGNORE).sum())
        n_sent += int(((y >= tok.sentinel_ids[0]) & (y != IGNORE)).sum())
    rows.append(("span corruption（T5 编解码）", n_pos, n_loss, n_sent))
    print(pad("目标", 30) + pad("每 1000 个位置的损失项", 24, ">") + pad("其中是哨兵", 12, ">")
          + pad("真正的\"内容\"预测", 18, ">"))
    for name, npos, nloss, nsent in rows:
        print(pad(name, 30) + pad(f"{1000 * nloss / npos:.0f}", 24, ">") + pad(f"{1000 * nsent / npos:.0f}", 12, ">")
              + pad(f"{1000 * (nloss - nsent) / npos:.0f}", 18, ">"))
    print(pad("ELECTRA 判别器", 30) + pad("1000", 24, ">") + pad("0", 12, ">") + pad("1000（二分类）", 18, ">"))
    print("""
观察 Observations：
  · 同样处理 1000 个位置（计算量大致相同），GPT 得到 1000 个"预测下一个词"的训练信号；prefix LM 约一半；
    BERT 只有约 140 个；T5 的目标里一大半是哨兵（哨兵按 <0>、<1>、<2> 的顺序出现，几乎不用猜），
    真正"写出被挖掉的词"的信号只有 80–110 个。
  · 这是 ELECTRA 的出发点：BERT 每处理一个序列，只从 15% 的位置学习。ELECTRA 让判别器对 100% 的位置做判断，
    同样的算力得到多得多的训练信号（虽然每个信号只是"换没换过"的二分类，信息量比"猜出是哪个词"少）。
  · 它也是 decoder-only 在大规模预训练里很有吸引力的原因之一：**每个 token 都是一道题**，没有浪费。
    代码 04 会看到：在相同的算力下，训练信号多的目标学得快得多。""")


# =============================================================================
# 演示 4：为什么不在字符级模型上做 masked LM
# =============================================================================
def demo_char_level(world):
    print_header("演示 4：字符级 masked LM 的两难 —— 遮单个字符太容易，遮整个词又拼不出词")
    words = sorted({w for line in world.train_lines for w in line.replace(".", " ").replace(":", " ")
                    .replace("?", " ").replace(",", " ").replace("+", " ").replace("=", " ").split()})
    rng = random.Random(0)
    total = unique = 0
    for line in world.train_lines[:2000]:
        for w in line.replace(".", " ").replace(":", " ").replace("?", " ").replace(",", " ").split():
            if "+" in w or "=" in w:
                continue
            for i in range(len(w)):
                if rng.random() < 0.15:                     # 这个字符被遮住
                    total += 1
                    pattern = w[:i] + "_" + w[i + 1:]
                    cands = [v for v in words if len(v) == len(w) and v[:i] == w[:i] and v[i + 1:] == w[i + 1:]]
                    unique += len(cands) == 1
    print(f"(a) 每个字符独立地以 15% 的概率被遮住：在 {total} 个被遮住的字符里，只看**同一个词里剩下的字母**、"
          f"查一下词表就能唯一确定答案的占 {unique / total:.0%}。")
    print("    例如 \"va_no\" 只能是 varno，\"r_ting\" 只能是 rating——模型学会拼写就够了，用不着理解上下文。\n")

    pos5 = [w for w in FOOD_POS if len(w) == 5]
    neg5 = [w for w in FOOD_NEG if len(w) == 5]
    print("(b) 改成遮住整个词：\"review: the soup was _____ and the host was kind. rating: good.\"")
    print(f"    右边的上下文说明食物是好的，答案是 5 个字母的正面形容词之一：{pos5}")
    guess = ""
    for i in range(5):
        counts = {}
        for w in pos5:
            counts[w[i]] = counts.get(w[i], 0) + 1
        best = max(sorted(counts), key=lambda c: counts[c])      # 并列时取字母序最小的
        guess += best
        print(f"    第 {i + 1} 个位置：各字母的概率 " + "  ".join(f"{c}:{counts[c] / len(pos5):.1f}" for c in sorted(counts))
              + f"   → 选 {best}")
    print(f"    即使模型把每个位置的分布都估计得**完全正确**，逐位置选最可能的字母也拼出 \"{guess}\"——"
          "不是任何一个食物形容词，也不在迷你互联网的词表里。")
    print(f"    （负面的 5 个字母形容词同理：{neg5}）")
    print("""
观察 Observations：
  · masked LM 对被遮住的多个位置是**分别独立**预测的（给定其余输入）。当一个"意思单位"（一个词）被拆成好几个 token 时，
    每个位置只知道"这里大概是 t、f、s、c、g 中的一个"，组合起来就乱了。XLNet 等工作批评的正是这种独立性假设。
  · 真实的 BERT 用 WordPiece：常见词是一个 token，一个 [MASK] 就是一个词，用一次 softmax 在 3 万个词里挑；
    长词被拆成几块时也会遇到同样的问题，所以后来有了 whole word masking 和 SpanBERT（整段遮住）。
  · T5 的 span corruption 用解码器**一个一个地**写出被挖掉的片段，后写的 token 能看到先写的，所以能写出完整的词和短语。
  · 本讲最初的版本就是在字符级上做 masked LM 的：训练 4000 步后，"填形容词"的正负判断还停在瞎猜的水平
    （见代码 README"关于实验结果"）。所以代码 04、05 改用词级 token：一个 token = 一个词，和 BERT、T5 的情形一致。""")


if __name__ == "__main__":
    world = build_world()
    tok = WordTokenizer(world.train_lines)
    demo_masks()
    demo_examples(world, tok)
    demo_signal(world, tok)
    demo_char_level(world)