{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM9-L3: Evaluate a tiny GPT end to end\n",
    "\n",
    "You will build a small, repeatable **evaluation harness** and use it to answer three questions about a tiny instruction-following GPT:\n",
    "\n",
    "1. Is the checkpoint trained longer really better than the one trained briefly, or could the difference be luck? (accuracy with confidence intervals and a **paired** test)\n",
    "2. How much do results change from run to run when the model **samples** instead of decoding greedily? (seeds and spread)\n",
    "3. What does inference **cost**? (time to first token, tokens per second, parameter and KV-cache memory)\n",
    "\n",
    "**The model and tasks are tiny and synthetic**: a 113,000-parameter character-level GPT tuned on three string instructions. The point is the *method*: criteria written down,\n",
    "fixed held-out items, seeds, intervals, paired comparison and measurement. The numbers do not transfer to large models; latency numbers depend on your machine.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, 2 to 4 minutes. No downloads.\n",
    "\n",
    "Cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`). Theory: Lessons LM9-T2, LM9-T3, LM9-T5 and LM9-T6."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import copy\n",
    "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__, \"| threads\", torch.get_num_threads())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Given: tasks, model and two checkpoints\n",
    "\n",
    "The model of the previous notebooks. Two checkpoints are made from the same start: `short` is instruction-tuned for 550 steps and `long` for 1,500 steps (about 35 seconds in total).\n",
    "Held-out words are random strings never used in training."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TASKS = {\"uppercase\": lambda w: w.upper(), \"repeat\": lambda w: w + \" \" + w, \"loud\": lambda w: w.upper() + \"!\"}\n",
    "INSTRUCTION = {\"uppercase\": \"write the word in capitals\", \"repeat\": \"say the word twice\", \"loud\": \"make the word loud\"}\n",
    "\n",
    "rng = random.Random(0)\n",
    "pool = set()\n",
    "while len(pool) < 460:\n",
    "    pool.add(\"\".join(rng.choice(\"abcdefghijklmnopqrstuvwxyz\") for _ in range(rng.randint(3, 5))))\n",
    "pool = sorted(pool)\n",
    "rng.shuffle(pool)\n",
    "train_words, eval_words = pool[:400], pool[400:]\n",
    "EVAL_ITEMS = [(t, w) for t in TASKS for w in eval_words[:20]]          # a fixed evaluation set: 60 items, defined once\n",
    "\n",
    "END, PADC = \"\\x03\", \"\\x00\"\n",
    "chars = sorted(set(\"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ #:\\n.-!\" + \"\".join(INSTRUCTION.values())) | {END, PADC})\n",
    "stoi = {c: i for i, c in enumerate(chars)}\n",
    "itos = {i: c for c, i in stoi.items()}\n",
    "VOCAB, PAD, END_ID = len(chars), stoi[PADC], stoi[END]\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "\n",
    "def prompt(task, word):\n",
    "    return f\"### Instruction:\\n{INSTRUCTION[task]}\\n\\n### Input:\\n{word}\\n\\n### Response:\\n\"\n",
    "\n",
    "def make_pair(p, response):\n",
    "    return torch.tensor(encode(p + response) + [END_ID]), torch.tensor([-100] * len(p) + encode(response) + [END_ID])\n",
    "\n",
    "def collate(pairs):\n",
    "    n = max(len(i) for i, _ in pairs)\n",
    "    ids = torch.full((len(pairs), n), PAD); lab = torch.full((len(pairs), n), -100)\n",
    "    for k, (i, l) in enumerate(pairs):\n",
    "        ids[k, : len(i)] = i; lab[k, : len(l)] = l\n",
    "    return ids[:, :-1], lab[:, 1:]\n",
    "\n",
    "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 GPT(nn.Module):\n",
    "    def __init__(self, vocab, context=96, d=64, heads=4, layers=2):\n",
    "        super().__init__()\n",
    "        self.context, self.d, self.layers = context, d, layers\n",
    "        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(context, d)\n",
    "        self.blocks = nn.ModuleList([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",
    "        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))\n",
    "        for b in self.blocks:\n",
    "            x = b(x)\n",
    "        return self.head(self.norm(x))\n",
    "\n",
    "def tune(model, steps, lr=1e-3, seed=2):\n",
    "    data = [make_pair(prompt(t, w), TASKS[t](w)) for t in TASKS for w in train_words]\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0)\n",
    "    for _ in range(steps):\n",
    "        ix = torch.randint(0, len(data), (32,), generator=g)\n",
    "        X, Y = collate([data[i] for i in ix])\n",
    "        loss = F.cross_entropy(model(X).flatten(0, 1), Y.flatten(), ignore_index=-100)\n",
    "        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()\n",
    "\n",
    "torch.manual_seed(0)\n",
    "start = GPT(VOCAB)\n",
    "text = \" \".join(random.Random(1).choice(pool) for _ in range(4000))\n",
    "stream = torch.tensor(encode(text))\n",
    "g = torch.Generator().manual_seed(1)\n",
    "opt = torch.optim.AdamW(start.parameters(), lr=3e-3, weight_decay=0.1)\n",
    "for _ in range(150):\n",
    "    s = torch.randint(0, len(stream) - 65, (32,), generator=g)\n",
    "    x = torch.stack([stream[i : i + 64] for i in s]); y = torch.stack([stream[i + 1 : i + 65] for i in s])\n",
    "    loss = F.cross_entropy(start(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()\n",
    "ckpt_short, ckpt_long = copy.deepcopy(start), copy.deepcopy(start)\n",
    "t0 = time.time()\n",
    "tune(ckpt_short, 550)\n",
    "tune(ckpt_long, 1500)\n",
    "print(f\"checkpoints ready in {time.time() - t0:.0f} s; {len(EVAL_ITEMS)} fixed evaluation items\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. TODO: the scoring function\n",
    "\n",
    "`is_correct(pred, ref)`: the normalised exact match used by the harness: both strings are stripped of leading and trailing whitespace; `pred` is correct if it equals `ref` exactly (case matters here, since \"BOX\" and \"box\" are different answers)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def is_correct(pred, ref):\n",
    "    return pred.strip() == ref.strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "assert is_correct(\"BOX \", \"BOX\") and not is_correct(\"box\", \"BOX\") and not is_correct(\"BOX!\", \"BOX\")\n",
    "print(\"self-check 1 passed: scoring\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 3. TODO: generation with a decoding setting\n",
    "\n",
    "`generate(model, p, temperature=None, seed=0, max_new=24)`: starting from the prompt string `p`, generate until the end token or `max_new` tokens and return the response string (without the end token).\n",
    "With `temperature=None` use greedy decoding (arg-max). Otherwise divide the logits by the temperature, take a softmax and draw with `torch.multinomial` using a `torch.Generator` seeded with `seed`.\n",
    "Evaluate in eval mode under `torch.no_grad()`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def generate(model, p, temperature=None, seed=0, max_new=24):\n",
    "    model.eval()\n",
    "    gen = torch.Generator().manual_seed(seed)\n",
    "    ids, out = encode(p), []\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_new):\n",
    "            logits = model(torch.tensor([ids[-model.context:]]))[0, -1]\n",
    "            if temperature is None:\n",
    "                nxt = int(logits.argmax())\n",
    "            else:\n",
    "                nxt = int(torch.multinomial(torch.softmax(logits / temperature, -1), 1, generator=gen))\n",
    "            if nxt == END_ID:\n",
    "                break\n",
    "            ids.append(nxt)\n",
    "            out.append(nxt)\n",
    "    model.train()\n",
    "    return \"\".join(itos[i] for i in out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "### Self-check 2"
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "pr = prompt(\"uppercase\", \"abc\")\n",
    "assert generate(ckpt_long, pr) == generate(ckpt_long, pr), \"greedy decoding is deterministic\"\n",
    "assert generate(ckpt_long, pr, temperature=1.0, seed=3) == generate(ckpt_long, pr, temperature=1.0, seed=3), \"the same seed gives the same sample\"\n",
    "outs = {generate(ckpt_short, prompt(\"repeat\", \"xyz\"), temperature=1.5, seed=s) for s in range(8)}\n",
    "assert len(outs) > 1, \"different seeds should give different samples at a high temperature (for an imperfect model)\"\n",
    "assert END not in generate(ckpt_long, pr)\n",
    "print(\"self-check 2 passed: decoding\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "## 4. TODO: the harness and an interval\n",
    "\n",
    "* `run_eval(model, items, temperature=None, seed=0)` returns a list of 0/1 results (one per item) using `generate` and `is_correct` (the reference is `TASKS[task](word)`).\n",
    "* `accuracy_ci(results)` returns `(p, low, high)`: the accuracy and a 95% normal-approximation interval `p ± 1.96 * sqrt(p * (1 - p) / n)`, clipped to `[0, 1]`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def run_eval(model, items, temperature=None, seed=0):\n",
    "    return [int(is_correct(generate(model, prompt(t, w), temperature, seed), TASKS[t](w))) for t, w in items]\n",
    "\n",
    "def accuracy_ci(results):\n",
    "    n = len(results)\n",
    "    p = sum(results) / n\n",
    "    half = 1.96 * math.sqrt(p * (1 - p) / n)\n",
    "    return p, max(0.0, p - half), min(1.0, p + half)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "### Self-check 3"
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "p, lo, hi = accuracy_ci([1] * 45 + [0] * 5)\n",
    "assert abs(p - 0.9) < 1e-12 and abs(lo - 0.8168) < 1e-3 and abs(hi - 0.9832) < 1e-3, f\"got {(p, lo, hi)}\"\n",
    "assert accuracy_ci([1] * 10) == (1.0, 1.0, 1.0)\n",
    "res = run_eval(ckpt_long, EVAL_ITEMS[:6])\n",
    "assert len(res) == 6 and set(res) <= {0, 1}\n",
    "print(\"self-check 3 passed: harness and interval\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "## 5. TODO: a paired comparison\n",
    "\n",
    "`sign_test(a, b)` takes two lists of 0/1 results on the same items and returns `(only_a, only_b, p)`: the number of items only A got right, only B got right, and the exact two-sided sign-test p-value\n",
    "`min(1, 2 * P(X <= min(only_a, only_b)))` for `X ~ Binomial(only_a + only_b, 1/2)` (1.0 if there are no disagreements). Use `math.comb`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def sign_test(a, b):\n",
    "    only_a = sum(1 for x, y in zip(a, b) if x == 1 and y == 0)\n",
    "    only_b = sum(1 for x, y in zip(a, b) if x == 0 and y == 1)\n",
    "    n = only_a + only_b\n",
    "    if n == 0:\n",
    "        return only_a, only_b, 1.0\n",
    "    k = min(only_a, only_b)\n",
    "    tail = sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n\n",
    "    return only_a, only_b, min(1.0, 2 * tail)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "### Self-check 4"
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "oa, ob, pv = sign_test([1] * 12 + [0] * 3 + [1] * 5, [0] * 12 + [1] * 3 + [1] * 5)\n",
    "assert (oa, ob) == (12, 3) and abs(pv - 0.03515625) < 1e-9, f\"got {(oa, ob, pv)}\"\n",
    "assert sign_test([1, 0], [1, 0]) == (0, 0, 1.0)\n",
    "print(\"self-check 4 passed: paired test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 6. TODO: inference cost\n",
    "\n",
    "* `measure_latency(model, p, n_tokens=16)` returns `(ttft, tps)`: the time in seconds for the **first** token (one forward pass over the prompt, and the arg-max of the last position) and the number of tokens per second when generating `n_tokens` further tokens one at a time by greedy decoding (do not stop at the end token). Use `time.perf_counter()` and `torch.no_grad()`.\n",
    "* `kv_cache_bytes(model, n_tokens, bytes_per=4)` returns the bytes of the key-value cache for `n_tokens`: `2 * layers * d * n_tokens * bytes_per`.\n",
    "* `param_bytes(model, bytes_per=4)` returns the bytes of the parameters."
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def measure_latency(model, p, n_tokens=16):\n",
    "    model.eval()\n",
    "    ids = encode(p)\n",
    "    with torch.no_grad():\n",
    "        t0 = time.perf_counter()\n",
    "        nxt = int(model(torch.tensor([ids]))[0, -1].argmax())\n",
    "        ttft = time.perf_counter() - t0\n",
    "        ids.append(nxt)\n",
    "        t1 = time.perf_counter()\n",
    "        for _ in range(n_tokens):\n",
    "            nxt = int(model(torch.tensor([ids[-model.context:]]))[0, -1].argmax())\n",
    "            ids.append(nxt)\n",
    "        tps = n_tokens / (time.perf_counter() - t1)\n",
    "    model.train()\n",
    "    return ttft, tps\n",
    "\n",
    "def kv_cache_bytes(model, n_tokens, bytes_per=4):\n",
    "    return 2 * model.layers * model.d * n_tokens * bytes_per\n",
    "\n",
    "def param_bytes(model, bytes_per=4):\n",
    "    return sum(p.numel() for p in model.parameters()) * bytes_per"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "### Self-check 5"
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ttft, tps = measure_latency(ckpt_long, prompt(\"loud\", \"abc\"))\n",
    "assert ttft > 0 and tps > 0\n",
    "assert kv_cache_bytes(ckpt_long, 100) == 2 * 2 * 64 * 100 * 4 == 102_400\n",
    "assert param_bytes(ckpt_long) == sum(p.numel() for p in ckpt_long.parameters()) * 4\n",
    "print(f\"self-check 5 passed: first token {ttft * 1000:.1f} ms, {tps:.0f} tokens/s (these numbers depend on your machine)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "## 7. Given: the experiments\n",
    "\n",
    "Greedy decoding of both checkpoints on the same 60 items; a paired test; then sampling at temperature 1.0 with five seeds for the long checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "res_short = run_eval(ckpt_short, EVAL_ITEMS)\n",
    "res_long = run_eval(ckpt_long, EVAL_ITEMS)\n",
    "for name, r in ((\"short (550 steps)\", res_short), (\"long (1500 steps)\", res_long)):\n",
    "    p, lo, hi = accuracy_ci(r)\n",
    "    print(f\"{name:18s} accuracy {p:.2f}   95% interval [{lo:.2f}, {hi:.2f}]   (n = {len(r)})\")\n",
    "oa, ob, pv = sign_test(res_short, res_long)\n",
    "print(f\"paired: only short right {oa}, only long right {ob}, p = {pv:.4f}\")\n",
    "\n",
    "sample_accs = [accuracy_ci(run_eval(ckpt_long, EVAL_ITEMS, temperature=1.0, seed=s))[0] for s in range(5)]\n",
    "mean_s = sum(sample_accs) / 5\n",
    "sd_s = (sum((a - mean_s) ** 2 for a in sample_accs) / 4) ** 0.5\n",
    "print(\"sampling T=1.0, five seeds:\", [round(a, 2) for a in sample_accs], f\"mean {mean_s:.2f} sd {sd_s:.3f}\")\n",
    "greedy_long = accuracy_ci(res_long)[0]\n",
    "\n",
    "by_task = {t: sum(r for (tt, _), r in zip(EVAL_ITEMS, res_long) if tt == t) / 20 for t in TASKS}\n",
    "print(\"long checkpoint by task:\", {k: round(v, 2) for k, v in by_task.items()})\n",
    "latency = {name: measure_latency(m, prompt(\"loud\", \"abc\")) for name, m in ((\"short\", ckpt_short), (\"long\", ckpt_long))}\n",
    "print(\"first-token time and tokens/s:\", {k: (round(v[0] * 1000, 1), round(v[1])) for k, v in latency.items()})\n",
    "print(f\"parameters {param_bytes(ckpt_long) / 1e6:.2f} MB (float32); KV cache for 96 tokens {kv_cache_bytes(ckpt_long, 96) / 1e3:.0f} kB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "### Self-check 6: what must hold"
   ]
  },
  {
   "cell_type": "code",
   "id": "c028",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "assert accuracy_ci(res_long)[0] > accuracy_ci(res_short)[0], \"the longer-trained checkpoint should score higher on these items\"\n",
    "assert accuracy_ci(res_long)[0] >= 0.6, \"the long checkpoint should solve most items\"\n",
    "assert pv < 0.05, \"the paired test should find the difference between the checkpoints real\"\n",
    "assert sd_s > 0 or mean_s == 1.0, \"different seeds should give different accuracies when sampling (unless every sample is perfect)\"\n",
    "assert mean_s <= greedy_long + 0.1, \"sampling at temperature 1.0 should not beat greedy decoding by more than noise on a deterministic task\"\n",
    "print(\"self-check 6 passed: the comparison is statistically supported and the variability of sampling is measured\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c029",
   "metadata": {},
   "source": [
    "## 8. Reflect\n",
    "\n",
    "1. The short checkpoint's interval is wide with only 60 items, and the long checkpoint's is [1.00, 1.00]. Why is a zero-width interval at 100% misleading (hint: the \"rule of three\": with no errors in n items, a 95% upper bound on the error rate is about 3/n)? How many items would you need to separate two checkpoints 3 points apart, and why did the paired test succeed here with so few?\n",
    "2. Sampling produced different accuracies for different seeds. What would you report for a sampled system: one number, a mean, or a mean with spread? How many seeds?\n",
    "3. Time to first token and tokens per second were measured on this CPU, with one request. Why can't you read the cost of serving a thousand users from them? What else would you measure?\n",
    "4. List three things this evaluation does *not* tell you about the model (think of prompts it has not seen, other inputs, safety).\n",
    "\n",
    "**Done when** all six self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}