{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验 03：训练一个最小 GPT\n",
    "\n",
    "**运行条件：** CPU 约 2–5 分钟；Apple Silicon 或 CUDA 更快；内存约 1 GB。\n",
    "\n",
    "这个模型使用字符级词表和两层 decoder block。模型很小，但训练目标、因果注意力、残差连接和自回归生成与大型 decoder-only 模型一致。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(11)\n",
    "device = 'mps' if torch.backends.mps.is_available() else ('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "text = ('语言模型根据历史上下文预测下一个符号。注意力让每个位置读取此前的信息。' * 80)\n",
    "chars = sorted(set(text))\n",
    "stoi = {char: i for i, char in enumerate(chars)}\n",
    "itos = {i: char for char, i in stoi.items()}\n",
    "tokens = torch.tensor([stoi[c] for c in text], dtype=torch.long)\n",
    "print('device:', device, 'tokens:', len(tokens), 'vocab:', len(chars))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, width=64, heads=4):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(width)\n",
    "        self.attn = nn.MultiheadAttention(width, heads, batch_first=True)\n",
    "        self.norm2 = nn.LayerNorm(width)\n",
    "        self.mlp = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(), nn.Linear(4*width, width))\n",
    "    def forward(self, x):\n",
    "        length = x.shape[1]\n",
    "        mask = torch.triu(torch.ones(length, length, device=x.device, dtype=torch.bool), diagonal=1)\n",
    "        normalized = self.norm1(x)\n",
    "        attended, _ = self.attn(normalized, normalized, normalized, attn_mask=mask, need_weights=False)\n",
    "        x = x + attended\n",
    "        return x + self.mlp(self.norm2(x))\n",
    "\n",
    "class TinyGPT(nn.Module):\n",
    "    def __init__(self, vocab, context=48, width=64, layers=2):\n",
    "        super().__init__()\n",
    "        self.context = context\n",
    "        self.token = nn.Embedding(vocab, width)\n",
    "        self.position = nn.Embedding(context, width)\n",
    "        self.blocks = nn.ModuleList([Block(width) for _ in range(layers)])\n",
    "        self.norm = nn.LayerNorm(width)\n",
    "        self.head = nn.Linear(width, vocab, bias=False)\n",
    "        self.head.weight = self.token.weight\n",
    "        self.apply(self._initialize)\n",
    "    @staticmethod\n",
    "    def _initialize(module):\n",
    "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
    "            nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if isinstance(module, nn.Linear) and module.bias is not None: nn.init.zeros_(module.bias)\n",
    "    def forward(self, ids, targets=None):\n",
    "        positions = torch.arange(ids.shape[1], device=ids.device)\n",
    "        x = self.token(ids) + self.position(positions)\n",
    "        for block in self.blocks: x = block(x)\n",
    "        logits = self.head(self.norm(x))\n",
    "        loss = None if targets is None else F.cross_entropy(logits.flatten(0,1), targets.flatten())\n",
    "        return logits, loss\n",
    "\n",
    "model = TinyGPT(len(chars)).to(device)\n",
    "print('parameters:', sum(p.numel() for p in model.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练\n",
    "\n",
    "每个输入序列的标签是向右平移一个位置的同一段文本。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)\n",
    "losses = []\n",
    "context, batch_size = model.context, 32\n",
    "for step in range(220):\n",
    "    starts = torch.randint(0, len(tokens)-context-1, (batch_size,))\n",
    "    x = torch.stack([tokens[s:s+context] for s in starts]).to(device)\n",
    "    y = torch.stack([tokens[s+1:s+context+1] for s in starts]).to(device)\n",
    "    _, loss = model(x, y)\n",
    "    optimizer.zero_grad(set_to_none=True); loss.backward(); optimizer.step()\n",
    "    losses.append(loss.item())\n",
    "    if step % 55 == 0: print(f'step={step:03d} loss={loss.item():.4f}')\n",
    "plt.plot(losses); plt.xlabel('step'); plt.ylabel('cross entropy'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 自回归生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def generate(prompt, new_tokens=50, temperature=0.8):\n",
    "    model.eval()\n",
    "    ids = torch.tensor([[stoi[c] for c in prompt if c in stoi]], device=device)\n",
    "    for _ in range(new_tokens):\n",
    "        context_ids = ids[:, -model.context:]\n",
    "        logits, _ = model(context_ids)\n",
    "        probs = torch.softmax(logits[:, -1] / temperature, dim=-1)\n",
    "        ids = torch.cat((ids, torch.multinomial(probs, 1)), dim=1)\n",
    "    return ''.join(itos[i] for i in ids[0].tolist())\n",
    "print(generate('语言模型'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 练习\n",
    "\n",
    "1. 将层数从 2 改为 1 和 4，比较参数量与收敛速度。\n",
    "2. 删除位置向量，观察生成质量。\n",
    "3. 增大训练语料的多样性，解释为什么训练损失可能上升而泛化能力反而改善。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}