{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM5-L3: Pretrain a tiny GPT and diagnose the run\n",
    "\n",
    "You will train a very small GPT (about 100,000 parameters) to predict the next **character** of a synthetic text, using the same recipe as large language models:\n",
    "next-token cross-entropy, AdamW, gradient clipping, a training loop with a validation pass, and sampling at different stages. Then you will write a small\n",
    "diagnosis function and use it on a healthy run and on two broken ones.\n",
    "\n",
    "**This is a real training run, on a toy problem.** The text is generated from a few templates, so the model can learn it almost perfectly in a few hundred steps. That makes the\n",
    "learning curve fast and clean; it does not say anything about how well a model of this size would do on natural language.\n",
    "\n",
    "**Hardware:** CPU is enough: about 1 to 2 minutes for the whole notebook. A GPU is not needed. No downloads.\n",
    "\n",
    "Cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`). Theory: Lessons LM5-T1 to LM5-T4."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "import random\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(0)\n",
    "print(\"torch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Given: the data\n",
    "\n",
    "A synthetic corpus of 1,200 short sentences about a robot, and a character-level tokenizer (each distinct character is a token).\n",
    "The stream is split **first** into 90% training and 10% validation (Lesson LM2-T6), so no validation window overlaps training text."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "SUBJ = [\"the robot\", \"the arm\", \"the drone\", \"the cart\"]\n",
    "VERB = [\"picks up\", \"moves to\", \"inspects\", \"avoids\"]\n",
    "COLOR = [\"red\", \"green\", \"blue\"]\n",
    "NOUN = [\"box\", \"bin\", \"cone\", \"tool\"]\n",
    "\n",
    "def make_corpus(n_sentences=1200, seed=0):\n",
    "    r = random.Random(seed)\n",
    "    return \" \".join(f\"{r.choice(SUBJ)} {r.choice(VERB)} the {r.choice(COLOR)} {r.choice(NOUN)}.\" for _ in range(n_sentences))\n",
    "\n",
    "text = make_corpus()\n",
    "chars = sorted(set(text))\n",
    "stoi = {c: i for i, c in enumerate(chars)}\n",
    "itos = {i: c for c, i in stoi.items()}\n",
    "VOCAB = len(chars)\n",
    "data = torch.tensor([stoi[c] for c in text])\n",
    "cut = int(0.9 * len(data))\n",
    "train_data, val_data = data[:cut], data[cut:]\n",
    "print(f\"{len(text):,} characters, vocabulary {VOCAB}, train {len(train_data):,}, validation {len(val_data):,}\")\n",
    "print(text[:70])\n",
    "WORDS = set(text.replace(\".\", \"\").split())\n",
    "print(\"loss of a uniform guess, ln(vocab) =\", round(math.log(VOCAB), 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. Given: a small GPT\n",
    "\n",
    "The same structure as Module 4 (pre-LN blocks, learned positions), written compactly. Width 64, 4 heads, 2 blocks, context 64."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, d, heads):\n",
    "        super().__init__()\n",
    "        self.h = heads\n",
    "        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)\n",
    "        self.qkv = nn.Linear(d, 3 * d, bias=False)\n",
    "        self.proj = nn.Linear(d, d)\n",
    "        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, n, d = x.shape\n",
    "        q, k, v = self.qkv(self.n1(x)).view(B, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)\n",
    "        x = x + self.proj(F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, n, d))\n",
    "        return x + self.ff(self.n2(x))\n",
    "\n",
    "class TinyGPT(nn.Module):\n",
    "    def __init__(self, vocab, context=64, d=64, heads=4, layers=2):\n",
    "        super().__init__()\n",
    "        self.context = context\n",
    "        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)\n",
    "        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(layers)])\n",
    "        self.norm, self.head = nn.LayerNorm(d), nn.Linear(d, vocab, bias=False)\n",
    "\n",
    "    def forward(self, idx):\n",
    "        return self.head(self.norm(self.blocks(self.tok(idx) + self.pos(torch.arange(idx.shape[1])))))\n",
    "\n",
    "model = TinyGPT(VOCAB)\n",
    "print(\"parameters:\", sum(p.numel() for p in model.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 3. TODO: batches\n",
    "\n",
    "`get_batch(data, batch_size, context, generator)` picks `batch_size` random start positions (`torch.randint(0, len(data) - context - 1, (batch_size,), generator=generator)`)\n",
    "and returns the input windows `x` and the targets `y`, both of shape `(batch_size, context)`, where `y` is `x` shifted one position to the right (Lesson LM2-T5)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_batch(data, batch_size, context, generator):\n",
    "    starts = torch.randint(0, len(data) - context - 1, (batch_size,), generator=generator)\n",
    "    x = torch.stack([data[s : s + context] for s in starts])\n",
    "    y = torch.stack([data[s + 1 : s + context + 1] for s in starts])\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "g = torch.Generator().manual_seed(1)\n",
    "x, y = get_batch(train_data, 8, 64, g)\n",
    "assert x.shape == (8, 64) and y.shape == (8, 64), \"both x and y should have shape (batch_size, context)\"\n",
    "assert torch.equal(x[:, 1:], y[:, :-1]), \"the targets must be the inputs shifted by one position\"\n",
    "assert x.max() < VOCAB\n",
    "print(\"self-check 1 passed: batches\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 4. TODO: the loss\n",
    "\n",
    "Cross-entropy between the logits `(B, n, V)` and the targets `(B, n)`. `F.cross_entropy` wants logits of shape `(N, V)` and targets of shape `(N,)`: flatten the first two dimensions."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def loss_fn(model, x, y):\n",
    "    logits = model(x)\n",
    "    return F.cross_entropy(logits.flatten(0, 1), y.flatten())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "### Self-check 2: an untrained model should start near ln(V)"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "init_loss = loss_fn(model, x, y).item()\n",
    "print(f\"initial loss {init_loss:.3f}   ln(vocab) = {math.log(VOCAB):.3f}\")\n",
    "assert abs(init_loss - math.log(VOCAB)) < 0.35 * math.log(VOCAB), \"an untrained model should be near ln(vocab); check the loss and the flattening\"\n",
    "print(\"self-check 2 passed: loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 5. TODO: evaluation\n",
    "\n",
    "`estimate_loss(model, data, n_batches, generator)`: the mean loss over `n_batches` batches of 32 windows, with the model in evaluation mode and gradients disabled\n",
    "(`torch.no_grad()`). Put the model back in training mode before returning."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def estimate_loss(model, data, n_batches, generator):\n",
    "    model.eval()\n",
    "    total = 0.0\n",
    "    with torch.no_grad():\n",
    "        for _ in range(n_batches):\n",
    "            xb, yb = get_batch(data, 32, model.context, generator)\n",
    "            total += loss_fn(model, xb, yb).item()\n",
    "    model.train()\n",
    "    return total / n_batches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "### Self-check 3"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "g = torch.Generator().manual_seed(2)\n",
    "v0 = estimate_loss(model, val_data, 5, g)\n",
    "assert model.training, \"estimate_loss must leave the model in training mode\"\n",
    "assert abs(v0 - init_loss) < 0.5, \"validation loss of the untrained model should also be near ln(vocab)\"\n",
    "print(\"self-check 3 passed: evaluation, validation loss\", round(v0, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 6. TODO: the training loop\n",
    "\n",
    "For each step: draw a batch, compute the loss, `optimizer.zero_grad()`, `loss.backward()`, clip the gradient norm to `clip` with `torch.nn.utils.clip_grad_norm_` (if `clip` is not None), `optimizer.step()`.\n",
    "Every `eval_every` steps (including step 0) append `(step, train_loss, val_loss)` to `history`, using `estimate_loss` on both sets. Use `torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def train(model, train_data, val_data, steps, lr, eval_every=50, clip=1.0, batch_size=32, wd=0.1, seed=1):\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)\n",
    "    history = []\n",
    "    for step in range(steps + 1):\n",
    "        if step % eval_every == 0:\n",
    "            history.append((step, estimate_loss(model, train_data, 10, g), estimate_loss(model, val_data, 10, g)))\n",
    "        if step == steps:\n",
    "            break\n",
    "        xb, yb = get_batch(train_data, batch_size, model.context, g)\n",
    "        loss = loss_fn(model, xb, yb)\n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        if clip is not None:\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        opt.step()\n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "### Self-check 4: train for 400 steps (about 10 seconds on a CPU)"
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def sample(model, prompt, n_chars=120, temperature=1.0, top_k=None, seed=0):\n",
    "    \"\"\"Given: autoregressive sampling with temperature and optional top-k (Lessons LM5-T5, LM5-T6).\"\"\"\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    model.eval()\n",
    "    idx = torch.tensor([[stoi[c] for c in prompt]])\n",
    "    for _ in range(n_chars):\n",
    "        logits = model(idx[:, -model.context:])[:, -1, :] / temperature\n",
    "        if top_k is not None:\n",
    "            kth = torch.topk(logits, top_k).values[:, -1:]\n",
    "            logits = logits.masked_fill(logits < kth, float(\"-inf\"))\n",
    "        nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=g)\n",
    "        idx = torch.cat([idx, nxt], dim=1)\n",
    "    model.train()\n",
    "    return \"\".join(itos[i] for i in idx[0].tolist())\n",
    "\n",
    "def word_quality(s):\n",
    "    words = s.replace(\".\", \" \").split()\n",
    "    return sum(w in WORDS for w in words) / max(1, len(words))\n",
    "\n",
    "print(\"BEFORE training:\", repr(sample(model, \"the robot \", 100)))\n",
    "t0 = time.time()\n",
    "torch.manual_seed(0)\n",
    "model = TinyGPT(VOCAB)\n",
    "history = train(model, train_data, val_data, steps=400, lr=3e-3)\n",
    "print(f\"trained in {time.time() - t0:.1f} s\")\n",
    "for step, tr, va in history[::2]:\n",
    "    print(f\"step {step:4d}  train {tr:.3f}  val {va:.3f}\")\n",
    "assert len(history) == 9 and history[0][0] == 0 and history[-1][0] == 400, \"history should have an entry at steps 0, 50, ..., 400\"\n",
    "assert history[-1][2] < 0.5 * history[0][2], \"validation loss should fall below half of its initial value\"\n",
    "assert all(history[i + 1][1] < history[i][1] + 0.05 for i in range(len(history) - 1)), \"training loss should keep falling (allowing small noise)\"\n",
    "print(\"AFTER training: \", repr(sample(model, \"the robot \", 100)))\n",
    "assert word_quality(sample(model, \"the robot \", 300)) > 0.8, \"after training, most generated words should be words of the corpus\"\n",
    "print(\"self-check 4 passed: the model learned the toy language\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "## 7. Sampling at different temperatures\n",
    "\n",
    "Same trained model, three temperatures. Read the output (no assertion): at a low temperature the text is the most likely continuation and repeats; at a high temperature\n",
    "you will see more invented words. Compare with Lesson LM5-T5."
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "for T in (0.3, 1.0, 1.8):\n",
    "    s = sample(model, \"the robot \", 160, temperature=T, seed=3)\n",
    "    print(f\"T={T}: word quality {word_quality(s):.2f}   {s[:90]!r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "## 8. TODO: diagnose a run\n",
    "\n",
    "Write `diagnose(history)`, which takes the list of `(step, train, val)` and returns one of the strings:\n",
    "\n",
    "* `\"diverged\"`: any loss is not finite, or any training loss is above twice the initial training loss;\n",
    "* `\"underfitting\"`: the final training loss is still above half of the initial training loss;\n",
    "* `\"overfitting\"`: the final validation loss is at least 0.25 above the final training loss **and** at least 0.05 above the lowest validation loss seen;\n",
    "* `\"healthy\"`: otherwise.\n",
    "\n",
    "Check the order: test for divergence first, then underfitting, then overfitting."
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def diagnose(history):\n",
    "    tr0 = history[0][1]\n",
    "    if any(not math.isfinite(a) or not math.isfinite(b) for _, a, b in history) or max(a for _, a, _ in history) > 2 * tr0:\n",
    "        return \"diverged\"\n",
    "    final_tr, final_va = history[-1][1], history[-1][2]\n",
    "    if final_tr > 0.5 * tr0:\n",
    "        return \"underfitting\"\n",
    "    if final_va - final_tr >= 0.25 and final_va >= min(v for _, _, v in history) + 0.05:\n",
    "        return \"overfitting\"\n",
    "    return \"healthy\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "### Self-check 5: one healthy run and two broken ones\n",
    "\n",
    "* **Learning rate far too high** (3.0, no clipping): training blows up or stalls.\n",
    "* **Memorisation**: a bigger model trained for a long time on only 600 characters; the training loss falls, the validation loss does not."
   ]
  },
  {
   "cell_type": "code",
   "id": "c028",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "assert diagnose(history) == \"healthy\", \"the run of section 6 should be diagnosed as healthy\"\n",
    "\n",
    "torch.manual_seed(0)\n",
    "bad_lr = train(TinyGPT(VOCAB), train_data, val_data, steps=200, lr=3.0, clip=None)\n",
    "print(\"lr=3.0:\", [(s, round(a, 2)) for s, a, _ in bad_lr])\n",
    "assert diagnose(bad_lr) in (\"diverged\", \"underfitting\"), \"a learning rate of 3.0 should give a run that diverges or does not learn\"\n",
    "\n",
    "torch.manual_seed(0)\n",
    "t0 = time.time()\n",
    "overfit = train(TinyGPT(VOCAB, d=96), train_data[:600], val_data, steps=400, lr=3e-3, wd=0.0)\n",
    "print(f\"memorisation run ({time.time() - t0:.0f} s):\", [(s, round(a, 2), round(b, 2)) for s, a, b in overfit[::2]])\n",
    "assert diagnose(overfit) == \"overfitting\", \"training loss low, validation loss high and no longer improving: that is overfitting\"\n",
    "print(\"self-check 5 passed: diagnosis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c029",
   "metadata": {},
   "source": [
    "## 9. Reflect\n",
    "\n",
    "1. The first loss was close to ln(22) = 3.09. Why? What would you suspect if it had been 0.5 at step 0?\n",
    "2. The healthy run's loss stops at about 0.2, not 0. Why can't it reach 0 on this text? (Hint: look at how the sentences are generated.)\n",
    "3. In the memorisation run, which curve do you trust, training or validation, and what would you change: more data, a smaller model, more regularisation?\n",
    "4. What does gradient clipping protect against, and why was it not enough in the learning-rate-3.0 run?\n",
    "\n",
    "**Done when** all five self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}