"""
tiny_gpt.py —— 一个可以在 CPU 上几十秒训练完的小 GPT（Decoder-only Transformer）
=====================================================================

【这个文件要回答的问题】
    第 1 讲要用一个"真的"神经语言模型做三件事：
        1. 看清楚三类语言模型（encoder-only / decoder-only / encoder-decoder）的区别在哪里——
           关键只是注意力掩码（attention mask）；
        2. 在城市语言上预训练它，和 n-gram、理论下界比较；
        3. 用它核对"参数量公式"和"训练 FLOPs ≈ 6ND"这两条贯穿全课的估算规则。
    Transformer 内部每个部件为什么这样设计，第 6 讲（Transformer）和第 7 讲（预训练 LLM）会细讲；
    这里先把它当作一个"给定前文、输出下一个 token 分布"的函数来用，代码注释只讲到够用为止。

【结构】（GPT-2 风格，Pre-LN）
    输入 id ──► token 嵌入 + 位置嵌入 ──► [Block] × n_layer ──► LayerNorm ──► 线性层（与 token 嵌入共享权重）──► logits
    Block:  x = x + Attention(LN(x))      ← 不同位置之间交换信息（掩码决定"谁能看谁"）
            x = x + MLP(LN(x))            ← 每个位置各自做非线性变换（d → 4d → d）

【本文件提供】
    GPTConfig、TinyGPT（forward / hidden_states / mask 可选 "causal" 或 "full"）
    count_params、param_formula（参数量的精确公式，04 号脚本逐个核对）
    make_stream、get_batch（把句子"打包"成连续的 token 流，按随机窗口取小批量）
    train_lm、load_or_train（带缓存：同样的配置和数据只训练一次）
    sentence_nll、perplexity、next_token_distribution、generate

【可复现】
    torch.set_num_threads(1) + 固定种子：同一台机器上重复训练，得到的权重逐位相同。

【系统联系】
    这个模型的每一步计算都会在后面的讲次里被单独拿出来优化：嵌入查表与 softmax（第 10 讲的融合内核）、
    注意力（第 21 讲 FlashAttention）、逐 token 生成（第 9 讲推测解码、第 22–24 讲服务系统）、
    训练循环（第 14–18 讲分布式训练）。
"""

import hashlib
import json
import math
import os
import time
from dataclasses import asdict, dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(1)   # 单线程：结果可复现，而且小矩阵上多线程本来也快不了多少


# =====================================================================
# 1. 模型
# =====================================================================

@dataclass
class GPTConfig:
    """
    模型的"形状"。

    【属性】
        vocab_size : 词表大小 V
        ctx_len    : 最长上下文 T（位置嵌入表的行数）
        n_layer    : Block 的层数 L
        n_head     : 注意力头数（d_model 必须能被它整除）
        d_model    : 每个 token 的向量维度 d
        d_ff       : MLP 的中间维度；0 表示用默认的 4·d
        tie_weights: 输出层是否和 token 嵌入共享同一个权重矩阵（GPT-2 的做法，省 V·d 个参数）
    【Python 小知识】@dataclass 自动生成 __init__ 等方法：GPTConfig(vocab_size=80, d_model=64)。
    """
    vocab_size: int
    ctx_len: int = 64
    n_layer: int = 2
    n_head: int = 4
    d_model: int = 64
    d_ff: int = 0
    tie_weights: bool = True

    @property
    def ff(self):
        """MLP 中间维度（d_ff 为 0 时取 4·d_model）。【Python 小知识】@property 让方法像属性一样用：cfg.ff。"""
        return self.d_ff or 4 * self.d_model


class SelfAttention(nn.Module):
    """
    多头自注意力。每个位置产生一个查询 q、一个键 k、一个值 v；位置 i 的输出是所有"允许看的"位置 j 的
    v_j 的加权平均，权重 ∝ exp(q_i · k_j / √d_head)。

    【掩码 mask】决定"位置 i 能看哪些位置 j"：
        "causal" : 只能看 j ≤ i（自己和过去）——decoder-only 语言模型（GPT）
        "full"   : 能看所有 j（过去和未来）——encoder-only 模型（BERT）
        布尔张量 [T, T]：自定义，True 表示允许看（例如 prefix-LM 或编码器-解码器的组合）
    【边界】这里用最直白的写法（显式算出 T×T 的分数矩阵再 softmax），方便看清楚掩码的作用、
        也方便 FLOP 计数器统计。真实系统会用融合的注意力内核（第 21 讲 FlashAttention）。
    """

    def __init__(self, cfg):
        super().__init__()
        assert cfg.d_model % cfg.n_head == 0, "d_model 必须能被 n_head 整除"
        self.n_head = cfg.n_head
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)   # 一次矩阵乘法同时算出 q、k、v
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)       # 多个头的输出拼起来后再混合一次

    def forward(self, x, mask="causal", cache=None):
        """
        【参数】
            x     : [B, T, d] 的张量（B 个序列，每个 T 个位置）
            mask  : 见类说明
            cache : None，或一个字典（KV 缓存）。给定时，x 只包含"新来的"位置，之前位置的 k、v 从缓存里取，
                    新算出的 k、v 追加进缓存。只有 mask="causal" 时才能这样做（02 号脚本解释为什么）。
        【返回】[B, T, d]
        【Python 小知识】
            - tensor.split(d, dim=2) 沿第 2 维每 d 列切一块；
            - view / transpose 只改变"怎么看"这块内存（形状、步长），不复制数据（前置讲 P1 会细讲步长）。
        """
        B, T, C = x.shape
        H = self.n_head
        q, k, v = self.qkv(x).split(C, dim=2)
        # [B, T, C] → [B, H, T, C/H]：把 d 维拆成 H 个头，每个头各自做注意力
        q = q.view(B, T, H, C // H).transpose(1, 2)
        k = k.view(B, T, H, C // H).transpose(1, 2)
        v = v.view(B, T, H, C // H).transpose(1, 2)
        past = 0
        if cache is not None:
            if mask != "causal":
                raise ValueError("只有因果（causal）注意力才能使用 KV 缓存")
            if "k" in cache:
                past = cache["k"].size(2)
                k = torch.cat([cache["k"], k], dim=2)   # 旧的键 + 新的键
                v = torch.cat([cache["v"], v], dim=2)
            cache["k"], cache["v"] = k, v
        att = (q @ k.transpose(-2, -1)) / math.sqrt(C // H)   # [B, H, T, past+T]：每对 (i, j) 的相关分数
        if isinstance(mask, str):
            if mask == "causal":
                # 新位置 i 的绝对位置是 past+i，它只能看绝对位置 j ≤ past+i 的键
                rows = torch.arange(T, device=x.device).unsqueeze(1) + past
                cols = torch.arange(past + T, device=x.device).unsqueeze(0)
                att = att.masked_fill(cols > rows, float("-inf"))
            elif mask != "full":
                raise ValueError(f"未知的 mask: {mask}")
        else:
            att = att.masked_fill(~mask, float("-inf"))
        att = F.softmax(att, dim=-1)     # 每一行变成概率分布（被禁止的位置 exp(−inf) = 0）
        y = att @ v                       # [B, H, T, C/H]
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)


class MLP(nn.Module):
    """逐位置的前馈网络：d → ff → d，中间用 GELU 激活。"""

    def __init__(self, cfg):
        super().__init__()
        self.fc = nn.Linear(cfg.d_model, cfg.ff)
        self.proj = nn.Linear(cfg.ff, cfg.d_model)

    def forward(self, x):
        return self.proj(F.gelu(self.fc(x)))


class Block(nn.Module):
    """一个 Transformer 层：x + Attn(LN(x))，再 + MLP(LN(x))（Pre-LN；残差让梯度能直接穿过很多层）。"""

    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.attn = SelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = MLP(cfg)

    def forward(self, x, mask="causal", cache=None):
        x = x + self.attn(self.ln1(x), mask, cache)
        x = x + self.mlp(self.ln2(x))
        return x


class TinyGPT(nn.Module):
    """
    Decoder-only Transformer 语言模型（mask="full" 时就变成一个 encoder-only 的双向模型）。

    【用法】
        model = TinyGPT(GPTConfig(vocab_size=80))
        logits, loss = model(x, y)           # x, y: [B, T] 的 id 张量；y 是 x 右移一位
    """

    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)   # 查表：id → d 维向量
        self.pos_emb = nn.Embedding(cfg.ctx_len, cfg.d_model)      # 第 t 个位置 → d 维向量
        # 【Python 小知识】nn.ModuleList 让列表里的子模块被 PyTorch "看见"（参数会被注册、会被 .parameters() 返回）
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)  # d 维向量 → V 个分数（logits）
        if cfg.tie_weights:
            self.head.weight = self.tok_emb.weight   # 共享同一个 [V, d] 矩阵（同一个 Parameter 对象）
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m):
        """GPT-2 风格的初始化：权重 ~ N(0, 0.02²)，偏置为 0。【边界】共享的权重会被初始化两次，结果一样。"""
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def hidden_states(self, idx, mask="causal", caches=None):
        """
        最后一层（经过 ln_f）的隐藏状态 [B, T, d]——每个位置对"前文"（或整句）的总结。

        【参数】caches：None，或长度为 n_layer 的字典列表（每层一个 KV 缓存，见 new_cache）。
            给定时 idx 只是"新来的" token，它们的位置编号从缓存里已有的长度开始。
        【边界】总长度不能超过 ctx_len（位置嵌入表只有 ctx_len 行），否则抛出错误。
        """
        B, T = idx.shape
        past = 0
        if caches is not None and "k" in caches[0]:
            past = caches[0]["k"].size(2)
        if past + T > self.cfg.ctx_len:
            raise ValueError(f"序列长度 {past + T} 超过了 ctx_len={self.cfg.ctx_len}")
        pos = torch.arange(past, past + T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)    # 广播：[B, T, d] + [T, d]
        for i, blk in enumerate(self.blocks):
            x = blk(x, mask, None if caches is None else caches[i])
        return self.ln_f(x)

    def new_cache(self):
        """一份空的 KV 缓存：每层一个空字典。"""
        return [{} for _ in self.blocks]

    def forward(self, idx, targets=None, mask="causal", caches=None):
        """
        【参数】idx：[B, T] 输入 id；targets：[B, T] 目标 id（-100 表示这个位置不计损失），可为 None；
            caches：KV 缓存（见 hidden_states），只在逐 token 生成时使用
        【返回】(logits [B, T, V], loss 或 None)；loss 是所有有效位置的平均交叉熵（nat/token）
        """
        logits = self.head(self.hidden_states(idx, mask, caches))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
        return logits, loss


def count_params(model):
    """
    模型的参数个数。

    【边界】共享的权重只算一次（model.parameters() 本身会对同一个 Parameter 去重）。
    【Python 小知识】p.numel() 是张量的元素个数。
    """
    return sum(p.numel() for p in model.parameters())


def param_formula(cfg):
    """
    参数量的精确公式（不用建模型）。04 号脚本会把它和 count_params 逐个核对。

    【推导】d = d_model，f = ff，V = vocab_size，T = ctx_len，L = n_layer
        token 嵌入 V·d；位置嵌入 T·d
        每层：LayerNorm 2 个 × (2d)        = 4d         （每个 LayerNorm 有缩放 γ 和平移 β）
              注意力 qkv: d·3d + 3d；proj: d·d + d = 4d² + 4d
              MLP fc: d·f + f；proj: f·d + d   = 2df + f + d
              合计 4d² + 2df + 9d + f；f = 4d 时 = 12d² + 13d
        最后的 LayerNorm 2d；输出层：共享权重时 0，否则 V·d
    【返回】字典：embedding、position、per_layer、layers、final_norm、head、total、non_embedding
        （non_embedding = 总数 − token 嵌入 − 位置嵌入，Kaplan 等人的规模定律用它）
    """
    d, f, V, T, L = cfg.d_model, cfg.ff, cfg.vocab_size, cfg.ctx_len, cfg.n_layer
    per_layer = 4 * d * d + 2 * d * f + 9 * d + f
    out = {"embedding": V * d, "position": T * d, "per_layer": per_layer, "layers": L * per_layer,
           "final_norm": 2 * d, "head": 0 if cfg.tie_weights else V * d}
    out["total"] = out["embedding"] + out["position"] + out["layers"] + out["final_norm"] + out["head"]
    out["non_embedding"] = out["total"] - out["embedding"] - out["position"]
    return out


# =====================================================================
# 2. 数据：把句子打包成 token 流
# =====================================================================

def make_stream(sentences, vocab):
    """
    把句子首尾相接成一条长长的 token 流：<bos> s1 <bos> s2 <bos> s3 ...

    【为什么要"打包"（packing）而不是"每句一行、短的补齐"（padding）？】
        句子长短不一，补齐会让很多位置算了也白算（第 03 号脚本会算出浪费的比例）。打包后每个位置都是
        真实 token，算力不浪费。代价是一个训练窗口里会跨越句子边界——在这门语言里句子之间相互独立，
        模型会学会"看到 <bos> 就重新开始"。
    【参数】sentences：token 列表的列表；vocab：city_corpus.Vocab
    【返回】一维 LongTensor
    """
    ids = []
    for s in sentences:
        ids.append(vocab.bos_id)
        ids.extend(vocab.encode(s))
    return torch.tensor(ids, dtype=torch.long)


def get_batch(stream, batch_size, ctx_len, gen):
    """
    从 token 流里随机取 batch_size 个长度为 ctx_len 的窗口，构成一个小批量。

    【思路】输入 x = stream[i : i+T]，目标 y = stream[i+1 : i+T+1]（右移一位）：
        位置 t 的任务是"看着 x_0..x_t，预测 x_{t+1}"。一个窗口同时提供 T 个训练样本——
        这正是 decoder-only 模型训练效率高的原因之一（02 号脚本会对比）。
    【参数】stream：一维 LongTensor；gen：torch.Generator（控制随机性）
    【返回】(x, y)：两个 [batch_size, ctx_len] 的 LongTensor
    【边界】stream 长度必须大于 ctx_len + 1。
    """
    starts = torch.randint(0, len(stream) - ctx_len - 1, (batch_size,), generator=gen)
    x = torch.stack([stream[i:i + ctx_len] for i in starts.tolist()])
    y = torch.stack([stream[i + 1:i + ctx_len + 1] for i in starts.tolist()])
    return x, y


# =====================================================================
# 3. 训练
# =====================================================================

def lr_at(step, steps, lr, warmup):
    """
    学习率调度：先线性预热 warmup 步，再余弦衰减到 lr 的 10%。

    【为什么】一开始参数是随机的，梯度方向很乱，步子太大容易发散；后期步子变小，便于收敛到更低的损失。
    """
    if step < warmup:
        return lr * (step + 1) / warmup
    progress = (step - warmup) / max(1, steps - warmup)
    return lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress)))


def train_lm(cfg, train_stream, steps=1500, batch_size=32, lr=3e-3, warmup=100, seed=0,
             log_every=0, eval_fn=None):
    """
    预训练：最小化下一个 token 的交叉熵（原课第 1 讲的 Training Objective）。

    【思路】标准训练循环的五步：取一个小批量 → 前向算损失 → 反向算梯度 → 裁剪梯度 → 优化器更新。
    【参数】
        cfg          : GPTConfig
        train_stream : make_stream 的输出
        steps        : 优化步数；batch_size、lr、warmup：见上
        seed         : 同时决定初始化和取批的随机性
        log_every    : >0 时每隔这么多步调用一次 eval_fn(model) 并记录
        eval_fn      : 函数 model → 字典（例如验证集困惑度），可为 None
    【返回】(model, history)：history 是字典，含 "train_loss"（每步的损失）、"evals"（[(step, eval_fn 结果)]）、
        "seconds"（训练耗时，⏱ 与机器有关）
    【边界】AdamW 的 weight_decay=0.01；梯度范数裁剪到 1.0。训练结束时模型处于 eval 模式（本模型没有 dropout，
        所以 train / eval 模式其实没有区别，这里只是养成习惯）。
    """
    torch.manual_seed(seed)
    model = TinyGPT(cfg)
    gen = torch.Generator().manual_seed(seed + 1)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.01)
    history = {"train_loss": [], "evals": []}
    t0 = time.time()
    for step in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, steps, lr, warmup)
        x, y = get_batch(train_stream, batch_size, cfg.ctx_len, gen)
        _, loss = model(x, y)
        opt.zero_grad(set_to_none=True)     # 清空上一步的梯度（PyTorch 默认会累加梯度）
        loss.backward()                      # 反向传播：算出每个参数的梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()                           # 按梯度更新参数
        history["train_loss"].append(loss.item())
        if log_every and eval_fn is not None and ((step + 1) % log_every == 0 or step == 0):
            model.eval()
            history["evals"].append((step + 1, eval_fn(model)))
            model.train()
    history["seconds"] = time.time() - t0
    model.eval()
    return model, history


def _fingerprint(cfg, train_sentences, hparams):
    """
    给"配置 + 训练数据 + 超参数"算一个指纹（SHA-256 的前 12 位），用作缓存文件名。

    【为什么】改了任何一样，指纹就会变，load_or_train 就会重新训练，不会误用过期的缓存。
    """
    h = hashlib.sha256()
    h.update(json.dumps({"cfg": asdict(cfg), "hp": hparams}, sort_keys=True).encode())
    for s in train_sentences:
        h.update((" ".join(s) + "\n").encode())
    return h.hexdigest()[:12]


def load_or_train(cfg, train_sentences, vocab, tag="city_gpt", cache_dir=None, retrain=False,
                  eval_fn=None, log_every=0, **hparams):
    """
    如果同样的配置、数据和超参数训练过，就直接读缓存；否则训练并保存。

    【参数】
        cfg, train_sentences, vocab : 见 train_lm / make_stream
        tag       : 缓存文件名前缀
        cache_dir : 缓存目录；None 表示本文件旁边的 checkpoints/
        retrain   : True 时忽略缓存、重新训练（并覆盖缓存）
        eval_fn, log_every : 训练过程中定期评估（见 train_lm）；它们不影响训练出的权重，所以不计入指纹
        **hparams : 传给 train_lm 的超参数（steps、batch_size、lr、warmup、seed），计入指纹
    【返回】(model, info)：info 含 "path"、"cached"（是否来自缓存）、"history"（训练时的记录，
        缓存里保存的是训练那一次的记录）
    【Python 小知识】**hparams 把多余的关键字参数收集成一个字典；调用时再用 **hparams 展开传下去。
    """
    if cache_dir is None:
        cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints")
    fp = _fingerprint(cfg, train_sentences, hparams)
    path = os.path.join(cache_dir, f"{tag}-{fp}.pt")
    if os.path.exists(path) and not retrain:
        blob = torch.load(path, weights_only=False)
        model = TinyGPT(cfg)
        model.load_state_dict(blob["state_dict"])
        model.eval()
        return model, {"path": path, "cached": True, "history": blob["history"]}
    stream = make_stream(train_sentences, vocab)
    model, history = train_lm(cfg, stream, log_every=log_every, eval_fn=eval_fn, **hparams)
    os.makedirs(cache_dir, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "history": history, "cfg": asdict(cfg)}, path)
    return model, {"path": path, "cached": False, "history": history}


# =====================================================================
# 4. 评估与生成
# =====================================================================

@torch.no_grad()
def sentence_nll(model, sentences, vocab, batch_size=256, mask="causal"):
    """
    逐句计算负对数似然：每句单独喂给模型（前面只有一个 <bos>），和 city_corpus.true_logprob 的口径一致。

    【思路】输入 [<bos>, w1, ..., w_{n−1}]，目标 [w1, ..., wn]；不同长度的句子补齐到同一长度，
        补齐的位置目标设为 -100（不计损失）。
    【返回】列表：每句的 −log p(句子)（nat）
    【边界】句子长度（含 <bos>）不能超过模型的 ctx_len。
    【Python 小知识】@torch.no_grad() 让函数里的计算不记录计算图：省内存、更快（评估时不需要梯度）。
    """
    out = []
    for i in range(0, len(sentences), batch_size):
        chunk = sentences[i:i + batch_size]
        T = max(len(s) for s in chunk)
        x = torch.full((len(chunk), T), vocab.bos_id, dtype=torch.long)
        y = torch.full((len(chunk), T), -100, dtype=torch.long)
        for r, s in enumerate(chunk):
            ids = vocab.encode(s)
            x[r, 1:len(ids)] = torch.tensor(ids[:-1])
            y[r, :len(ids)] = torch.tensor(ids)
        logits, _ = model(x, mask=mask)
        logp = F.log_softmax(logits, dim=-1)
        valid = y != -100
        picked = logp.gather(-1, y.clamp(min=0).unsqueeze(-1)).squeeze(-1) * valid
        out.extend((-picked.sum(dim=1)).tolist())
    return out


def perplexity(model, sentences, vocab):
    """
    测试集困惑度 = exp(总负对数似然 / 总 token 数)。

    【返回】(交叉熵 nat/token, 困惑度)
    """
    nll = sentence_nll(model, sentences, vocab)
    n_tok = sum(len(s) for s in sentences)
    ce = sum(nll) / n_tok
    return ce, math.exp(ce)


@torch.no_grad()
def next_token_distribution(model, vocab, prefix):
    """
    模型给出的 P(· | <bos> + 前缀)。

    【参数】prefix：token 列表（不含 <bos>）
    【返回】字典 {token: 概率}（整个词表，含 <bos>）
    """
    x = torch.tensor([[vocab.bos_id] + vocab.encode(prefix)])
    logits, _ = model(x)
    probs = F.softmax(logits[0, -1], dim=-1)
    return {vocab.itos[i]: probs[i].item() for i in range(len(vocab))}


@torch.no_grad()
def generate(model, vocab, gen, n=1, max_len=20, temperature=1.0, prefix=()):
    """
    从模型里采样 n 句话（一次并行生成 n 句，每步给每句都追加一个 token）。

    【思路】自回归生成：每一步把"目前为止的全部 token"喂给模型，取最后一个位置的分布，按温度采样一个 token。
        这里为了简单**没有**用 KV 缓存，每步都从头重算整个前缀；02 号脚本会说明 decoder-only 模型为什么
        可以缓存、缓存能省多少计算。
    【参数】gen：torch.Generator；n：句数；max_len：最长 token 数；temperature：温度（>0；越小越保守）；
        prefix：所有句子共同的开头（token 列表，不含 <bos>）
    【返回】n 个 token 列表（完整的句子，包含 prefix）；每句在第一个 "." 处截断，<bos> 会被去掉
    【边界】生成到 ctx_len 时强制停止。
    """
    start = [vocab.bos_id] + vocab.encode(list(prefix))
    x = torch.tensor([start] * n)
    done = torch.zeros(n, dtype=torch.bool)
    dot = vocab.stoi["."]
    for _ in range(max_len):
        if x.size(1) >= model.cfg.ctx_len:
            break
        logits, _ = model(x)
        probs = F.softmax(logits[:, -1] / temperature, dim=-1)
        nxt = torch.multinomial(probs, 1, generator=gen)
        x = torch.cat([x, nxt], dim=1)
        done |= nxt.squeeze(1) == dot
        if done.all():
            break
    outs = []
    for row in x.tolist():
        toks = [vocab.itos[i] for i in row[1:]]
        if "." in toks:
            toks = toks[:toks.index(".") + 1]
        outs.append([t for t in toks if t != "<bos>"])
    return outs