"""
tiny_transformer.py —— 一个在 CPU 上几十秒就能训练好的小 Transformer + 合成语料
=========================================================================
（第 5 讲首次出现；第 17、18、20 讲继续使用；计划第 21 讲和附加讲 A2）

【这个文件要回答的问题】
    讲注意力、KV 缓存、推测解码、LoRA、MoE、量化时，最好有一个**真的训练过、真的会做事**的模型，
    这样"改了实现以后结果还对不对""量化以后准确率掉了多少"才有意义。
    但真模型（GPT-2、Llama）要下载几百 MB、要 GPU。所以我们自己造一个"迷你世界"：

    LikesWorld 语料（每篇文档一行，词级别的 token）：
        <bos> bob likes jam . ann likes tea . eve likes figs . q ann tea . q bob jam . <eos>
        └───────────── 事实 facts ─────────────┘ └──────── 问答 queries ───────┘

    - 事实："某人 likes 某物 ."，一篇 2–6 条，人名互不相同、物品互不相同；
    - 问答："q 某人 某物 ."（q = query，读作"问：某人 → 某物"），问的人一定在事实里出现过，
      答案就是这个人喜欢的东西。紧跟在被问的人名后面的那个词就是答案。
    要答对，模型必须"回头看"：在上文里找到 "ann likes ___"，把空格里的词抄过来。
    这正是**注意力**最擅长的事（第 5 讲演示 2 会把它找到的位置画出来）。

    语料的几个性质，后面几讲会用到：
      * 很多 token 是"套话"（likes . q 的位置完全确定）→ 小模型也能猜对 → 推测解码的接受率高（第 17 讲）；
      * 问谁是随机的 → 有些 token 本质上不可预测 → 接受率不会是 100%；
      * 同一批事实可以被很多请求共享 → 正是前缀复用（第 16 讲 RadixAttention）的场景；
      * 把问法反过来（"q tea ann ."：谁喜欢茶？）就是一个需要新查找回路的任务 → 第 18 讲里 LoRA 很难学会的任务，第 20 讲剪枝实验的第二个"领域"。

【模型】一个 Llama 风格的"仅解码器"（decoder-only）Transformer，默认约 8.6 万个参数：
        词嵌入
        → n_layers 个块：x = x + Attn(RMSNorm(x));  x = x + SwiGLU-MLP(RMSNorm(x))
           （注意力里用 RoPE 旋转位置编码）
        → RMSNorm → 线性层输出每个词的分数（logits）
    注意力的"核心计算"（softmax(QKᵀ/√d)·V）被做成可替换的函数 attn_fn，
    第 5 讲把它换成自己写的 FlashAttention / Flash-Decoding，再检查模型输出是否不变。

【主要接口】
    Config                     模型超参数（frozen dataclass，可以当字典的键）
    make_document(rng, ...)    生成一篇文档（词的列表）
    encode / decode            词 ↔ 编号
    sample_batch(rng, B)       一批训练数据（补齐到 max_len）
    TinyTransformer            模型；forward(tokens, cache=None) 支持 KV 缓存
    KVCache                    每层的 K、V 缓存
    reference_attention        标准注意力（PyTorch 张量运算）
    generate(...)              贪心生成，可选是否用 KV 缓存
    answer_accuracy(...)       问答里答案词的准确率
    load_or_train(cfg)         有缓存就读，没有就训练并缓存到 .cache/ 目录

【运行】python3 tiny_transformer.py    （自测：训练或读取模型，打印准确率；第一次约 1 分钟）

【Python 小知识】
    - `@dataclass(frozen=True)`：自动生成 __init__ 等方法；frozen=True 表示创建后不能改字段，
      这样对象可以做字典的键、可以算哈希。
    - `torch.nn.Module`：PyTorch 里"带参数的层"的基类。在 __init__ 里把子层赋值给 self.xxx，
      PyTorch 会自动登记它们的参数；调用 model(x) 等于调用 model.forward(x)。
    - `torch.no_grad()`：with 块里不记录计算图（不需要求梯度时更快、更省内存）。
"""

import hashlib
import json
import math
import os
import time
from dataclasses import asdict, dataclass
from typing import Callable, List, Optional, Sequence

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

# ---------------------------------------------------------------------------
# 1. 语料：LikesWorld
# ---------------------------------------------------------------------------
NAMES = ("ann", "bob", "cat", "dan", "eve", "fay", "gus", "hal")
ITEMS = ("tea", "jam", "figs", "rice", "soup", "cake", "nuts", "plums")
SPECIALS = ("<pad>", "<bos>", "<eos>")
WORDS = ("likes", ".", "q")
VOCAB = SPECIALS + WORDS + NAMES + ITEMS          # 3 + 3 + 8 + 8 = 22 个词
STOI = {w: i for i, w in enumerate(VOCAB)}         # string → index
PAD, BOS, EOS = STOI["<pad>"], STOI["<bos>"], STOI["<eos>"]
Q_TOKEN = STOI["q"]                                # "q 某人 某物 ."：q 后面第 2 个词是答案

MIN_FACTS, MAX_FACTS = 2, 6
MIN_QUERIES, MAX_QUERIES = 1, 3
# 最长的文档：<bos> + 6 条事实 × 4 词 + 3 个问答 × 4 词 + <eos> = 38 个词
MAX_DOC_LEN = 2 + MAX_FACTS * 4 + MAX_QUERIES * 4


def make_document(rng: np.random.Generator, n_facts: Optional[int] = None,
                  n_queries: Optional[int] = None, reverse: bool = False) -> List[str]:
    """生成一篇 LikesWorld 文档。

    【思路】先随机选 n_facts 个不同的人和 n_facts 个不同的物品配对成事实，再随机抽问题。
    【参数】
        rng       : numpy 的随机数生成器（np.random.default_rng(seed)），由调用者固定种子 → 可复现
        n_facts   : 事实条数；None 表示在 [2, 6] 里随机
        n_queries : 问答个数；None 表示在 [1, 3] 里随机（同一个人可能被问两次）
        reverse   : False 问"某人喜欢什么"（q ann tea .），True 反过来问"谁喜欢某物"（q tea ann .）。
                    基础模型只在 reverse=False 上训练；reverse=True 留给第 18、20 讲当"新任务"。
    【返回】词的列表，第一个是 "<bos>"，最后一个是 "<eos>"。
    【边界】n_facts 不能超过 8（人名、物品各只有 8 个）；越界时报 ValueError。
    """
    if n_facts is None:
        n_facts = int(rng.integers(MIN_FACTS, MAX_FACTS + 1))
    if n_queries is None:
        n_queries = int(rng.integers(MIN_QUERIES, MAX_QUERIES + 1))
    if not 1 <= n_facts <= min(len(NAMES), len(ITEMS)):
        raise ValueError(f"n_facts 必须在 1..{min(len(NAMES), len(ITEMS))} 之间，收到 {n_facts}")
    # 【Python 小知识】rng.choice(8, size=k, replace=False) 从 0..7 里不放回地抽 k 个
    names = rng.choice(len(NAMES), size=n_facts, replace=False)
    items = rng.choice(len(ITEMS), size=n_facts, replace=False)
    facts = list(zip(names, items))
    words = ["<bos>"]
    for n, it in facts:
        words += [NAMES[n], "likes", ITEMS[it], "."]
    for _ in range(n_queries):
        n, it = facts[int(rng.integers(n_facts))]
        if reverse:
            words += ["q", ITEMS[it], NAMES[n], "."]
        else:
            words += ["q", NAMES[n], ITEMS[it], "."]
    words.append("<eos>")
    return words


def encode(words: Sequence[str]) -> List[int]:
    """词 → 编号。遇到词表里没有的词会抛 KeyError（这个世界的词表是封闭的）。"""
    return [STOI[w] for w in words]


def decode(ids: Sequence[int]) -> List[str]:
    """编号 → 词。"""
    return [VOCAB[int(i)] for i in ids]


def sample_batch(rng: np.random.Generator, batch_size: int, max_len: int,
                 reverse_frac: float = 0.0) -> torch.Tensor:
    """一批训练数据：形状 (batch_size, max_len) 的整数张量，短的文档在右边用 <pad> 补齐。

    【参数】reverse_frac：有多大比例的文档用反向问法（基础训练为 0）。
    【边界】文档比 max_len 长时报错（默认 max_len = 48 ≥ 最长文档 38）。
    """
    out = np.full((batch_size, max_len), PAD, dtype=np.int64)
    for b in range(batch_size):
        ids = encode(make_document(rng, reverse=bool(rng.random() < reverse_frac)))
        if len(ids) > max_len:
            raise ValueError(f"文档长度 {len(ids)} 超过 max_len={max_len}")
        out[b, :len(ids)] = ids
    return torch.from_numpy(out)


# ---------------------------------------------------------------------------
# 2. 模型
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class Config:
    """模型超参数。默认值（3000 步，约 1 分钟）能把答案准确率练到 100%。"""
    vocab_size: int = len(VOCAB)
    d_model: int = 64       # 隐藏维度
    n_heads: int = 4        # 注意力头数；每个头的维度 d_head = d_model / n_heads = 16
    n_layers: int = 2       # 层数。要"先找到问的是谁，再去事实里找这个人"，至少需要 2 层
    d_ff: int = 128         # SwiGLU 中间层宽度
    max_len: int = 48       # 最长序列
    rope_base: float = 10000.0   # RoPE 的频率基数（Llama 的默认值）

    @property
    def d_head(self) -> int:
        return self.d_model // self.n_heads


class KVCache:
    """每一层一对 (K, V) 张量，形状 (batch, heads, 已缓存的长度, d_head)。

    【为什么需要】生成第 t 个词时，前面 t-1 个词的 K、V 和上一步算的完全一样，存下来就不用重算
    （第 5 讲 §8）。代价是显存：每个词每层要存 2 × d_model 个数。
    """

    def __init__(self, n_layers: int):
        self.k: List[Optional[torch.Tensor]] = [None] * n_layers
        self.v: List[Optional[torch.Tensor]] = [None] * n_layers

    @property
    def length(self) -> int:
        """已经缓存了多少个位置（所有层一样长）。"""
        return 0 if self.k[0] is None else int(self.k[0].shape[-2])

    def append(self, layer: int, k: torch.Tensor, v: torch.Tensor):
        """把新算出的 k、v 拼到第 layer 层的缓存后面，返回拼好的完整 K、V。"""
        if self.k[layer] is None:
            self.k[layer], self.v[layer] = k, v
        else:
            self.k[layer] = torch.cat([self.k[layer], k], dim=-2)
            self.v[layer] = torch.cat([self.v[layer], v], dim=-2)
        return self.k[layer], self.v[layer]


def causal_mask(q_start: int, tq: int, tk: int, device=None) -> torch.Tensor:
    """因果掩码：形状 (tq, tk) 的布尔张量，True 表示"这个 key 在 query 的未来，要屏蔽"。

    【思路】第 i 个 query 的位置是 q_start + i，第 j 个 key 的位置是 j；key 位置 > query 位置就屏蔽。
            q_start = 0 是普通的训练 / 预填充（prefill）；q_start = 已缓存长度是解码（decode）。
    """
    qpos = torch.arange(q_start, q_start + tq, device=device)[:, None]
    kpos = torch.arange(tk, device=device)[None, :]
    return kpos > qpos


def attention_probs(q: torch.Tensor, k: torch.Tensor, q_start: int) -> torch.Tensor:
    """注意力权重 P = softmax(QKᵀ/√d_head + 因果掩码)，形状 (B, H, Tq, Tk)。"""
    s = (q @ k.transpose(-1, -2)) / math.sqrt(q.shape[-1])
    s = s.masked_fill(causal_mask(q_start, q.shape[-2], k.shape[-2], q.device), float("-inf"))
    return torch.softmax(s, dim=-1)


def reference_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_start: int) -> torch.Tensor:
    """标准注意力：O = softmax(QKᵀ/√d + mask)·V。

    【参数】q: (B, H, Tq, d)；k, v: (B, H, Tk, d)；q_start: 第一个 query 的绝对位置。
    【返回】(B, H, Tq, d)。
    【边界】要求 Tk ≥ q_start + Tq（每个 query 至少能看到它自己），否则某一行全被屏蔽，softmax 得到 NaN。
    """
    return attention_probs(q, k, q_start) @ v


# 注意力核心函数的类型：(q, k, v, q_start) -> o
AttnFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int], torch.Tensor]


class RMSNorm(nn.Module):
    """RMSNorm：x / sqrt(mean(x²) + eps) × 可学习的缩放。比 LayerNorm 少减均值一步（Llama 用它）。"""

    def __init__(self, d: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


def apply_rope(x: torch.Tensor, start: int, base: float) -> torch.Tensor:
    """旋转位置编码（Rotary Position Embedding, RoPE），Llama 用的位置编码。

    【直觉】把每个头的向量看成 d/2 个二维小向量，第 i 个小向量按"位置 × 频率_i"的角度旋转。
           两个向量的点积只取决于它们的**角度差**，也就是**位置差**——
           所以 q·k 天然只依赖"隔了几个词"，"看前面第 2 个词"这种规律很容易学（换成可学习的
           绝对位置嵌入，这个模型 4000 步也学不会，见 README 的"计划 vs 实际"）。
    【参数】x: (B, H, T, d)，d 必须是偶数；start: 第一个位置的编号（解码时 = 已缓存长度）。
    【返回】同形状的张量。
    【边界】K 要在**放进缓存之前**旋转（按它自己的位置），之后就不用再动它。
    """
    d = x.shape[-1]
    if d % 2:
        raise ValueError("RoPE 要求 d_head 是偶数")
    pos = torch.arange(start, start + x.shape[-2], dtype=torch.float32, device=x.device)
    freqs = base ** (-torch.arange(0, d, 2, dtype=torch.float32, device=x.device) / d)   # (d/2,)
    ang = pos[:, None] * freqs[None, :]                                                    # (T, d/2)
    cos, sin = torch.cos(ang).to(x.dtype), torch.sin(ang).to(x.dtype)
    x1, x2 = x[..., 0::2], x[..., 1::2]               # 偶数维、奇数维两两配对
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * cos - x2 * sin
    out[..., 1::2] = x1 * sin + x2 * cos
    return out


class Attention(nn.Module):
    """多头自注意力。QKV 投影 → 拆成多个头 → RoPE →（拼上 KV 缓存）→ attn_fn → 拼回 → 输出投影。"""

    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.wqkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.wo = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        self.record = False                      # True 时把注意力权重存到 self.last_probs（演示用）
        self.last_probs: Optional[torch.Tensor] = None

    def forward(self, x, layer: int, cache: Optional[KVCache], attn_fn: AttnFn, q_start: int):
        """q_start：这批新 token 里第一个的绝对位置（由模型统一传入）。

        【边界】不能在这里用 cache.length 现算：KVCache.length 看的是第 0 层的缓存，
               第 1 层执行时第 0 层已经把新 token 拼进去了，会多算 T（见 README"计划 vs 实际"）。
        """
        B, T, _ = x.shape
        H, dh = self.cfg.n_heads, self.cfg.d_head
        # (B, T, 3·d) → 三个 (B, T, d) → (B, H, T, dh)
        q, k, v = self.wqkv(x).split(self.cfg.d_model, dim=-1)
        q, k, v = (t.view(B, T, H, dh).transpose(1, 2) for t in (q, k, v))
        q = apply_rope(q, q_start, self.cfg.rope_base)
        k = apply_rope(k, q_start, self.cfg.rope_base)
        if cache is not None:
            k, v = cache.append(layer, k, v)     # k、v 现在包含了全部历史位置
        o = attn_fn(q, k, v, q_start)
        if self.record:
            self.last_probs = attention_probs(q, k, q_start).detach()
        o = o.transpose(1, 2).reshape(B, T, self.cfg.d_model)
        return self.wo(o)


class SwiGLU(nn.Module):
    """Llama 的前馈层：W2 · (SiLU(W1 x) ⊙ W3 x)。"""

    def __init__(self, cfg: Config):
        super().__init__()
        self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class Block(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.d_model)
        self.mlp = SwiGLU(cfg)

    def forward(self, x, layer, cache, attn_fn, start):
        x = x + self.attn(self.norm1(x), layer, cache, attn_fn, start)
        return x + self.mlp(self.norm2(x))


class TinyTransformer(nn.Module):
    """Llama 风格的仅解码器 Transformer（RoPE、RMSNorm、SwiGLU、没有偏置项）。"""

    def __init__(self, cfg: Config):
        super().__init__()
        if cfg.d_model % cfg.n_heads != 0:
            raise ValueError("d_model 必须能被 n_heads 整除")
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.attn_fn: AttnFn = reference_attention      # 可以换成别的实现

    def forward(self, tokens: torch.Tensor, cache: Optional[KVCache] = None) -> torch.Tensor:
        """【参数】tokens: (B, T) 整数；cache: None（不缓存，从位置 0 开始）或 KVCache（接着缓存往后算）。
        【返回】logits: (B, T, vocab_size)。
        【边界】起始位置 + T 超过 max_len 时报错（RoPE 本身没有长度上限，但模型只在 max_len 以内训练过）。
        """
        B, T = tokens.shape
        start = cache.length if cache is not None else 0
        if start + T > self.cfg.max_len:
            raise ValueError(f"序列长度 {start + T} 超过 max_len={self.cfg.max_len}")
        x = self.tok_emb(tokens)
        for i, blk in enumerate(self.blocks):
            x = blk(x, i, cache, self.attn_fn, start)
        return self.head(self.norm(x))

    def set_record(self, on: bool):
        """打开 / 关闭所有层的注意力权重记录。"""
        for blk in self.blocks:
            blk.attn.record = on


def count_params(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


# ---------------------------------------------------------------------------
# 3. 训练、评估、生成
# ---------------------------------------------------------------------------
def lm_loss(model: TinyTransformer, batch: torch.Tensor) -> torch.Tensor:
    """语言模型损失：用前 T-1 个词预测后 T-1 个词，<pad> 位置不计入。"""
    logits = model(batch[:, :-1])
    target = batch[:, 1:]
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), ignore_index=PAD)


def train(cfg: Config = Config(), steps: int = 3000, batch_size: int = 64, lr: float = 3e-3,
          seed: int = 0, reverse_frac: float = 0.0, verbose: bool = True) -> TinyTransformer:
    """从头训练一个模型。

    【思路】AdamW + 线性预热 100 步 + 余弦退火；每一步现场生成一批新文档（数据无限，不会过拟合）。
    【返回】训练好的模型（eval 模式）。
    【边界】为了逐位可复现，函数内部固定了 torch 的种子并设置单线程；调用后全局线程数仍是 1。
    """
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    rng = np.random.default_rng(seed)
    model = TinyTransformer(cfg)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01, betas=(0.9, 0.98))

    def lr_at(step):
        if step < 100:
            return lr * (step + 1) / 100
        frac = (step - 100) / max(1, steps - 100)
        return lr * 0.5 * (1 + math.cos(math.pi * frac))

    t0 = time.time()
    for step in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        loss = lm_loss(model, sample_batch(rng, batch_size, cfg.max_len, reverse_frac))
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if verbose and (step % 500 == 0 or step == steps - 1):
            print(f"    step {step:5d}  loss {loss.item():.4f}  ({time.time() - t0:.0f}s)")
    model.eval()
    return model


def _cache_path(cfg: Config, steps: int, seed: int, reverse_frac: float) -> str:
    """缓存文件名里带上配置的哈希：换了任何超参数都会重新训练，不会读到旧模型。"""
    key = json.dumps({"cfg": asdict(cfg), "steps": steps, "seed": seed, "rev": reverse_frac,
                      "vocab": VOCAB, "version": 2}, sort_keys=True)
    h = hashlib.sha1(key.encode()).hexdigest()[:10]
    here = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(here, ".cache", f"tiny_{h}.pt")


def load_or_train(cfg: Config = Config(), steps: int = 3000, seed: int = 0, reverse_frac: float = 0.0,
                  verbose: bool = True) -> TinyTransformer:
    """有缓存就读取，没有就训练并保存到 code/lecture05/.cache/。

    【边界】缓存只是为了省时间：删掉 .cache 目录后重新训练，得到的参数逐位相同（同一台机器上）。
            不同机器 / 不同 PyTorch 版本上训练出的参数可能有微小差别，准确率等结论不受影响。
    """
    path = _cache_path(cfg, steps, seed, reverse_frac)
    torch.set_num_threads(1)
    model = TinyTransformer(cfg)
    if os.path.exists(path):
        model.load_state_dict(torch.load(path, weights_only=True))
        model.eval()
        if verbose:
            print(f"    （读取缓存的模型 {os.path.relpath(path)}）")
        return model
    if verbose:
        print(f"    （没有缓存，开始训练：{steps} 步，约 1 分钟）")
    model = train(cfg, steps=steps, seed=seed, reverse_frac=reverse_frac, verbose=verbose)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save(model.state_dict(), path)
    return model


@torch.no_grad()
def answer_positions(ids: Sequence[int]) -> List[int]:
    """返回"被问的词"所在的位置 i（ids[i-1] 是 q）：模型在位置 i 的预测应该等于答案 ids[i+1]。"""
    return [i for i in range(1, len(ids) - 1) if ids[i - 1] == Q_TOKEN]


@torch.no_grad()
def answer_accuracy(model: TinyTransformer, docs: Sequence[Sequence[int]]) -> float:
    """在一组文档上，统计答案词被模型以最大概率预测对的比例（teacher forcing：输入是正确的上文）。"""
    correct = total = 0
    for ids in docs:
        pred = model(torch.tensor([ids]))[0].argmax(-1)
        for i in answer_positions(ids):
            total += 1
            correct += int(pred[i].item() == ids[i + 1])
    return correct / max(1, total)


def test_docs(n: int, seed: int = 12345, reverse: bool = False) -> List[List[int]]:
    """固定种子的测试文档（和训练用的随机流分开）。"""
    rng = np.random.default_rng(seed)
    return [encode(make_document(rng, reverse=reverse)) for _ in range(n)]


@torch.no_grad()
def generate(model: TinyTransformer, prompt: Sequence[int], max_new: int, use_cache: bool = True,
             stop_at_eos: bool = True, return_logits: bool = False):
    """贪心生成（每一步取概率最大的词）。

    【思路】
        use_cache=True ：先把整个 prompt 喂进去（预填充 prefill），之后每步只喂 1 个新词（解码 decode），
                         K、V 从缓存里拿；
        use_cache=False：每一步都把"到目前为止的整个序列"重新喂一遍（最朴素的做法）。
        两种方式在数学上完全等价（第 5 讲 05 会验证）。
    【返回】新生成的词编号列表；return_logits=True 时同时返回每一步的 logits（列表，每个形状 (vocab,)）。
    【边界】总长度达到 max_len 时提前停止。
    """
    out, logits_list = [], []
    seq = list(prompt)
    cache = KVCache(model.cfg.n_layers) if use_cache else None
    logits = model(torch.tensor([seq]), cache)[0, -1]
    for _ in range(max_new):
        nxt = int(logits.argmax())
        out.append(nxt)
        logits_list.append(logits)
        seq.append(nxt)
        if (stop_at_eos and nxt == EOS) or len(seq) >= model.cfg.max_len:
            break
        if use_cache:
            logits = model(torch.tensor([[nxt]]), cache)[0, -1]
        else:
            logits = model(torch.tensor([seq]))[0, -1]
    return (out, logits_list) if return_logits else out


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    print("一篇示例文档：", " ".join(make_document(rng)))
    m = load_or_train()
    print(f"参数量 {count_params(m):,}；答案准确率 {answer_accuracy(m, test_docs(300)):.3f}")