{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# Step 7 \u2014 Tokenizer and Sampler\n",
    "\n",
    "**Solution notebook.** Replaces the two toy ends of the pipeline: a\n",
    "byte-pair-encoding tokenizer going in, and a proper sampler coming\n",
    "out. BPE training on 1.1M characters is a few minutes of pure Python;\n",
    "everything else is seconds. Requires `step6_model.pt` from the Step 6\n",
    "notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c002",
   "metadata": {},
   "source": [
    "## Setup: data and tokenizer (Step 1's solution, reproduced)\n",
    "\n",
    "Every notebook in this series is self-contained: it re-creates what it\n",
    "needs from earlier steps in one compact cell, so you can run it top to\n",
    "bottom without opening the others. On Colab, uncomment the download."
   ]
  },
  {
   "cell_type": "code",
   "id": "c003",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "import torch\n",
    "\n",
    "with open('input.txt') as f:\n",
    "    text = f.read()\n",
    "vocab = sorted(set(text))\n",
    "V = len(vocab)\n",
    "stoi = {ch: i for i, ch in enumerate(vocab)}\n",
    "itos = {i: ch for i, ch in enumerate(vocab)}\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda ids: ''.join(itos[i] for i in ids)\n",
    "ids = torch.tensor(encode(text), dtype=torch.long)\n",
    "\n",
    "assert len(text) == 1_115_394 and V == 65\n",
    "print(f'{len(text):,} characters, vocab {V}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c004",
   "metadata": {},
   "source": [
    "## 7.1 BPE trainer\n",
    "\n",
    "Greedy compression (Lecture 7, Prop. 1.2): repeatedly merge the most\n",
    "frequent adjacent pair. Two helpers \u2014 count pairs, and replace every\n",
    "occurrence of one pair in a single left-to-right pass."
   ]
  },
  {
   "cell_type": "code",
   "id": "c005",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import time\n",
    "\n",
    "def pair_counts(seq):\n",
    "    return Counter(zip(seq, seq[1:]))\n",
    "\n",
    "def merge(seq, pair, new_id):\n",
    "    out, i = [], 0\n",
    "    while i < len(seq):\n",
    "        if i < len(seq) - 1 and (seq[i], seq[i+1]) == pair:\n",
    "            out.append(new_id); i += 2\n",
    "        else:\n",
    "            out.append(seq[i]); i += 1\n",
    "    return out\n",
    "\n",
    "NUM_MERGES = 256\n",
    "base_ids = encode(text)\n",
    "seq = list(base_ids)\n",
    "merges = {}                         # (id, id) -> new_id, in training order\n",
    "vocab_s = {i: ch for i, ch in enumerate(vocab)}\n",
    "sizes = {}\n",
    "\n",
    "t0 = time.time()\n",
    "for m in range(NUM_MERGES):\n",
    "    counts = pair_counts(seq)\n",
    "    pair, cnt = counts.most_common(1)[0]\n",
    "    new_id = V + m\n",
    "    merges[pair] = new_id\n",
    "    vocab_s[new_id] = vocab_s[pair[0]] + vocab_s[pair[1]]\n",
    "    seq = merge(seq, pair, new_id)\n",
    "    if m < 10:\n",
    "        print(f'  merge {m:>3}: {vocab_s[new_id]!r:<8} ({cnt:,} occurrences)')\n",
    "    sizes[m+1] = len(seq)\n",
    "print(f'\\ntrained {NUM_MERGES} merges in {time.time()-t0:.0f}s')\n",
    "print(f'{len(base_ids):,} chars -> {len(seq):,} tokens '\n",
    "      f'= {len(base_ids)/len(seq):.3f} chars/token')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c006",
   "metadata": {},
   "source": [
    "Reference first ten merges: `'e '`, `'th'`, `'t '`, `'s '`, `'d '`,\n",
    "`', '`, `'ou'`, `'er'`, `'in'`, `'y '` \u2014 then `an`, `:\\n`, `or`,\n",
    "`o `, `en`, `\\n\\n`. Three things worth a sentence each:\n",
    "\n",
    "1. **Word-final patterns dominate the top.** The space is the most\n",
    "   predictable character in English, so pairs ending in it are the most\n",
    "   frequent. Frequency alone finds the word boundary.\n",
    "2. **`:\\n` and `\\n\\n` are the document format** \u2014 the `SPEAKER:`\n",
    "   convention and the blank line between speeches, discovered with no\n",
    "   notion of what a play is.\n",
    "3. **By merge 28 whole words appear** (`'and '`). BPE crosses from\n",
    "   morphology into vocabulary on its own.\n",
    "\n",
    "Reference compression: 1.500 / **1.963** / 2.286 / 2.692 chars per\n",
    "token at 64 / 256 / 512 / 1024 merges \u2014 strongly diminishing returns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 7.2 Encoder and decoder\n",
    "\n",
    "Encoding applies merges **in training order**; doing them in any other\n",
    "order can give a different tokenization, which is why `merges` is an\n",
    "ordered dict and not a set."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def bpe_encode(s):\n",
    "    out = [stoi[c] for c in s]\n",
    "    for pair, new_id in merges.items():      # insertion order = training order\n",
    "        if len(out) < 2:\n",
    "            break\n",
    "        out = merge(out, pair, new_id)\n",
    "    return out\n",
    "\n",
    "def bpe_decode(toks):\n",
    "    return ''.join(vocab_s[t] for t in toks)\n",
    "\n",
    "unseen = ('Shall I compare thee to a summer\\'s day?\\n'\n",
    "          'Thou art more lovely and more temperate:')\n",
    "toks = bpe_encode(unseen)\n",
    "assert bpe_decode(toks) == unseen                  # exact round-trip\n",
    "print(f'{len(unseen)} chars -> {len(toks)} tokens')\n",
    "print([vocab_s[t] for t in toks[:14]])\n",
    "print('round-trip on unseen text: exact')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 7.3 The unit problem\n",
    "\n",
    "Per-token loss is **not** comparable across tokenizers. Lecture 7's\n",
    "Proposition 2.1 gives the invariant. Our Step 6 character model scored\n",
    "1.5178 nats/token, but each token was one character; a BPE model's\n",
    "tokens carry ~1.96 characters each, so its per-token loss must be\n",
    "*higher* to break even."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "def bits_per_char(nats_per_tok, n_tok, n_chars):\n",
    "    return nats_per_tok / math.log(2) * n_tok / n_chars\n",
    "\n",
    "print(f\"{'model':<26}{'nats/tok':>10}{'bits/char':>11}\")\n",
    "print(f\"{'Step 1 bigram (chars)':<26}{2.4549:>10.4f}\"\n",
    "      f\"{bits_per_char(2.4549, len(base_ids), len(text)):>11.4f}\")\n",
    "print(f\"{'Step 6 GPT (chars)':<26}{1.5178:>10.4f}\"\n",
    "      f\"{bits_per_char(1.5178, len(base_ids), len(text)):>11.4f}\")\n",
    "print(f\"\\nA BPE model would need {1.5178 * len(base_ids)/len(seq):.4f} \"\n",
    "      f\"nats/token to match Step 6's bits/char.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "**Retraining on BPE tokens** is the same Step 6 loop with $V' = 321$\n",
    "and the retokenized stream \u2014 about an hour of CPU, so it is left as an\n",
    "exercise rather than executed here. Convert both runs to bits/char\n",
    "before declaring a winner. A second effect partly offsets the cost:\n",
    "with the same $T_{\\max}$, a BPE model's window spans roughly twice as\n",
    "much text."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c012",
   "metadata": {},
   "source": [
    "## 7.4 The sampler\n",
    "\n",
    "Three knobs. Temperature is entropy-regularized soft argmax (Lecture 7,\n",
    "Prop. 3.1); top-$k$ truncates to a fixed count; nucleus truncates to an\n",
    "*adaptive* count \u2014 the smallest set carrying probability $p$."
   ]
  },
  {
   "cell_type": "code",
   "id": "c013",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Head(nn.Module):\n",
    "    def __init__(self, d, dh, T):\n",
    "        super().__init__()\n",
    "        self.key   = nn.Linear(d, dh, bias=False)\n",
    "        self.query = nn.Linear(d, dh, bias=False)\n",
    "        self.value = nn.Linear(d, dh, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(T, T)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, _ = x.shape\n",
    "        q, k, v = self.query(x), self.key(x), self.value(x)\n",
    "        s = q @ k.transpose(-2, -1) / k.shape[-1]**0.5\n",
    "        s = s.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        A = F.softmax(s, dim=-1)\n",
    "        self.A = A.detach()          # stashed for Step 8 interpretability\n",
    "        return A @ v\n",
    "\n",
    "class MHA(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList(Head(d, d // H, T) for _ in range(H))\n",
    "        self.proj = nn.Linear(d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.fc, self.proj = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(F.gelu(self.fc(x)))\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.ln1, self.attn = nn.LayerNorm(d), MHA(d, H, T)\n",
    "        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        return x + self.mlp(self.ln2(x))\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    def __init__(self, V, T, d, H, L):\n",
    "        super().__init__()\n",
    "        self.T = T\n",
    "        self.tok = nn.Embedding(V, d)\n",
    "        self.pos = nn.Embedding(T, d)\n",
    "        self.blocks = nn.ModuleList(Block(d, H, T) for _ in range(L))\n",
    "        self.lnf = nn.LayerNorm(d)\n",
    "        self.head = nn.Linear(d, V, bias=False)\n",
    "        for blk in self.blocks:   # GPT-2 residual scaling (Lecture 6 \u00a71)\n",
    "            nn.init.normal_(blk.attn.proj.weight, std=0.02 / math.sqrt(2 * L))\n",
    "            nn.init.normal_(blk.mlp.proj.weight,  std=0.02 / math.sqrt(2 * L))\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        x = self.tok(idx) + self.pos(torch.arange(T))\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "        logits = self.head(self.lnf(x))\n",
    "        loss = None if targets is None else F.cross_entropy(\n",
    "            logits.view(-1, logits.size(-1)), targets.reshape(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, idx, n_new, temperature=1.0):\n",
    "        for _ in range(n_new):\n",
    "            logits, _ = self(idx[:, -self.T:])\n",
    "            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)\n",
    "            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "model = GPT(V=65, T=64, d=128, H=4, L=4)\n",
    "model.load_state_dict(torch.load('step6_model.pt'))\n",
    "model.eval()\n",
    "print(f'loaded trained model, {sum(p.numel() for p in model.parameters()):,} parameters')\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "c015",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "def sample_next(logits, temperature=1.0, top_k=None, top_p=None):\n",
    "    logits = logits.clone()\n",
    "    if temperature <= 0:                      # greedy = tau -> 0 limit\n",
    "        return logits.argmax(-1, keepdim=True)\n",
    "    logits = logits / temperature\n",
    "    if top_k is not None:\n",
    "        kth = logits.topk(top_k, dim=-1).values[..., -1:]\n",
    "        logits = logits.masked_fill(logits < kth, float('-inf'))\n",
    "    if top_p is not None:\n",
    "        srt, idx = logits.sort(dim=-1, descending=True)\n",
    "        cum = srt.softmax(-1).cumsum(-1)\n",
    "        drop = cum - srt.softmax(-1) >= top_p   # keep the token that crosses p\n",
    "        srt = srt.masked_fill(drop, float('-inf'))\n",
    "        logits = torch.full_like(logits, float('-inf')).scatter(-1, idx, srt)\n",
    "    return torch.multinomial(logits.softmax(-1), 1)\n",
    "\n",
    "@torch.no_grad()\n",
    "def generate(prompt, n_new=250, **kw):\n",
    "    idx = torch.tensor([encode(prompt)])\n",
    "    for _ in range(n_new):\n",
    "        logits, _ = model(idx[:, -model.T:])\n",
    "        idx = torch.cat([idx, sample_next(logits[:, -1, :], **kw)], dim=1)\n",
    "    return decode(idx[0].tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c016",
   "metadata": {},
   "source": [
    "### The decoding gallery\n",
    "\n",
    "Same prompt, same model, seven decoding rules. Watch greedy fall into\n",
    "a loop and $\\tau=1.5$ lose English."
   ]
  },
  {
   "cell_type": "code",
   "id": "c017",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "settings = [('greedy (tau->0)', dict(temperature=0.0)),\n",
    "            ('tau = 0.3',       dict(temperature=0.3)),\n",
    "            ('tau = 0.8',       dict(temperature=0.8)),\n",
    "            ('tau = 1.0',       dict(temperature=1.0)),\n",
    "            ('tau = 1.5',       dict(temperature=1.5)),\n",
    "            ('top-k = 5',       dict(temperature=1.0, top_k=5)),\n",
    "            ('nucleus p = 0.9', dict(temperature=1.0, top_p=0.9))]\n",
    "for name, kw in settings:\n",
    "    torch.manual_seed(0)\n",
    "    print('=' * 70); print(name); print('=' * 70)\n",
    "    print(generate('ROMEO:', 200, **kw)); print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c018",
   "metadata": {},
   "source": [
    "**Greedy degenerates** \u2014 it locks into a repeated phrase and stays\n",
    "there. This is Holtzman et al.'s neural text degeneration, on your own\n",
    "model. The explanation (Lecture 7 \u00a73) is that the *mode* of a\n",
    "high-dimensional distribution is wildly unrepresentative of it: natural\n",
    "text is **typical**, not most-probable, and maximizing likelihood lands\n",
    "you outside the region where real text lives.\n",
    "\n",
    "The objective we train and the objective we want at decode time\n",
    "genuinely differ. That seam is what Lecture 8's alignment methods work\n",
    "on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "### Entropy of the decoding distributions\n",
    "\n",
    "Quantifying what each knob does: the mean entropy of the next-token\n",
    "distribution, and how many tokens the nucleus actually keeps."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "idx = torch.tensor([encode('ROMEO:\\nWhat light through yonder')])\n",
    "with torch.no_grad():\n",
    "    logits = model(idx)[0][0, -1, :]\n",
    "print(f\"{'setting':<20}{'entropy (nats)':>16}{'effective support':>20}\")\n",
    "for name, kw in settings[1:]:\n",
    "    lg = logits / kw.get('temperature', 1.0)\n",
    "    if kw.get('top_k'):\n",
    "        kth = lg.topk(kw['top_k']).values[-1]\n",
    "        lg = lg.masked_fill(lg < kth, float('-inf'))\n",
    "    if kw.get('top_p'):\n",
    "        srt, i2 = lg.sort(descending=True)\n",
    "        cum = srt.softmax(-1).cumsum(-1)\n",
    "        srt = srt.masked_fill(cum - srt.softmax(-1) >= kw['top_p'], float('-inf'))\n",
    "        lg = torch.full_like(lg, float('-inf')).scatter(-1, i2, srt)\n",
    "    p = lg.softmax(-1)\n",
    "    H = -(p[p > 0] * p[p > 0].log()).sum()\n",
    "    print(f'{name:<20}{H:>16.4f}{int((p > 1e-6).sum()):>20}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "\u2192 Continue with [Step 8](../project/step-8): the capstone."
   ]
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 }
}