"""
=============================================================================
第 9 讲 · 代码 1：三种预训练目标的数据长什么样
=============================================================================

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

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
预训练的核心思想：从**没有标签**的文本里自己造出"题目"和"答案"（自监督学习）。
三种主流架构各有各的造题方式：

    解码器（GPT）       因果语言模型：给前面的词，预测下一个词
    编码器（BERT）      掩码语言模型：遮住一些词，根据左右两边的上下文猜出来
    编码器—解码器（T5） span corruption：删掉几段词，让模型把删掉的内容写出来

    演示 1：同一条评论，三种目标分别造出什么样的输入和目标
    演示 2：BERT 的 80/10/10 规则 —— 在 1 万条评论上统计，验证实现是否正确
    演示 3：T5 的 span corruption —— 统计被删掉的比例、每段的长度、输入输出的长度
    演示 4：每条评论能提供多少个"训练信号"？

所有函数都在 lm_utils.py 里，这里只做演示和检查。只用 Python 标准库和 numpy（会 import PyTorch，但不训练）。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 01_pretraining_objectives.py

约 2 秒。
"""

import random
from collections import Counter

import numpy as np

from lm_utils import mask_tokens, pretrain_corpus, span_corrupt
from toy_reviews import Vocab


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


def show(tokens, width=11):
    """把一列词元排成等宽的一行，方便上下对齐比较。"""
    return "".join(f"{t:>{width}s}" for t in tokens)


def main():
    vocab = Vocab()
    words = "the acting was not boring . i loved it .".split()

    # ---------------- 演示 1 ----------------
    print_header("演示 1：同一条评论，三种预训练目标")
    ids = vocab.encode(words)
    toks = [vocab.itos[i] for i in ids]
    print("\n  (a) 因果语言模型（GPT）：每个位置预测下一个词")
    print("      输入 " + show(toks[:-1], 8))
    print("      目标 " + show(toks[1:], 8))
    print("""      → 每个位置都有一个目标；模型只能看左边（因果掩码），所以它是一个"生成"模型。""")

    rng = random.Random(3)
    inp, tgt = mask_tokens(ids, vocab, rng)
    print("\n  (b) 掩码语言模型（BERT）：遮住约 15% 的词，根据两边的上下文预测")
    print("      输入 " + show([vocab.itos[i] for i in inp], 8))
    print("      目标 " + show([vocab.itos[t] if t else "·" for t in tgt], 8))
    print("""      → "·" 的位置不计算损失。模型能同时看左右两边（双向），但只有被选中的位置提供训练信号。""")

    inp, tgt = span_corrupt(words, random.Random(1), noise_density=0.3)
    print("\n  (c) span corruption（T5）：删掉几段连续的词，用哨兵 <x0>、<x1>…… 代替，让解码器把它们写出来")
    print(f"      输入：{' '.join(inp)}")
    print(f"      目标：{' '.join(tgt)}")
    print("""      → （为了让例子更明显，这里删掉了 30% 的词；T5 默认是 15%。）
        编码器读带哨兵的输入（双向），解码器逐词生成目标（因果）。""")

    # ---------------- 演示 2 ----------------
    print_header("演示 2：BERT 的 80/10/10 规则（在 1 万条评论上统计）")
    corpus = pretrain_corpus(10000, seed=7)
    rng = random.Random(0)
    n_tokens = n_picked = 0
    kinds = Counter()
    for w in corpus:
        ids = vocab.encode(w)
        inp, tgt = mask_tokens(ids, vocab, rng)
        n_tokens += len(w)
        for i, t in enumerate(tgt):
            if t:
                n_picked += 1
                kinds["换成 <mask>" if inp[i] == vocab.stoi[Vocab.MASK] else
                      ("保持原词" if inp[i] == t else "换成随机词")] += 1
    print(f"\n    被选中作为预测目标的比例：{n_picked / n_tokens:.1%}（目标 15%）")
    for k in ["换成 <mask>", "换成随机词", "保持原词"]:
        print(f"      其中{k}：{kinds[k] / n_picked:.1%}")
    print("""
  观察：比例和 BERT 论文一致（"换成随机词"时偶尔会抽到原词本身，所以"保持原词"略多于 10%）。
    为什么不把被选中的词全部换成 <mask>？微调和使用时输入里永远没有 <mask>。
    如果被预测的位置总是 <mask>，模型就可以只在看到 <mask> 时才认真"思考"；
    混入随机词和原词之后，模型不知道哪些位置会被考，只好对**每个**词都建立好的上下文表示。""")

    # ---------------- 演示 3 ----------------
    print_header("演示 3：T5 的 span corruption（在 1 万条评论上统计）")
    rng = random.Random(0)
    noise = total = 0
    span_lens, in_lens, out_lens, orig_lens = [], [], [], []
    for w in corpus:
        inp, tgt = span_corrupt(w, rng)
        n_spans = sum(t in Vocab.SENTINELS for t in inp)
        noise += len(tgt) - n_spans
        total += len(w)
        span_lens.append((len(tgt) - n_spans) / n_spans)
        in_lens.append(len(inp))
        out_lens.append(len(tgt))
        orig_lens.append(len(w))
    print(f"\n    被删掉的词占 {noise / total:.1%}（目标 15%），平均每段 {np.mean(span_lens):.2f} 个词（目标 3）")
    print(f"    平均长度：原文 {np.mean(orig_lens):.1f} 个词 → 输入 {np.mean(in_lens):.1f} 个 + 目标 {np.mean(out_lens):.1f} 个")
    print("""
  观察：
    - 删掉的是**连续的一段**，模型要生成一整段内容（不只是一个词），学到的东西更接近"写作"。
    - 输入比原文短（被删掉的部分换成了一个哨兵），目标也很短：编码器和解码器处理的总长度
      和原文差不多，而 BERT 式的模型输入和原文一样长，却只有 15% 的位置在"学习"。
    - T5 的论文（Raffel 等 2020）系统地比较了很多种目标，发现这种做法效果好而且省计算。""")

    # ---------------- 演示 4 ----------------
    print_header("演示 4：每条评论能提供多少个训练信号？")
    avg = np.mean(orig_lens)
    print(f"\n    一条评论平均 {avg:.1f} 个词")
    print(f"      因果语言模型：每个位置都预测下一个词（含 </s>）→ 约 {avg + 1:.1f} 个训练信号")
    print(f"      掩码语言模型：只预测被选中的 15%           → 约 {0.15 * avg:.1f} 个训练信号")
    print(f"      span corruption：预测被删掉的词和哨兵      → 约 {np.mean(out_lens):.1f} 个训练信号")
    print("""
  观察：同样一条文本，因果语言模型得到的训练信号最多。BERT 每条数据只"学习"15% 的位置，
    这是它数据效率的一个弱点（ELECTRA 的改进：让模型判断**每个**词是不是被替换过，所有位置都提供信号）。
    但"训练信号多"不等于"适合所有任务"：双向的编码器在理解类任务上曾经明显更好 —— 02、03 号文件会比较。""")


if __name__ == "__main__":
    main()