{ "cells": [ { "cell_type": "markdown", "id": "c001", "metadata": {}, "source": [ "# LM12-L2: Capstone milestone 2: model, baselines and a training experiment\n", "\n", "Milestone 1 produced the data. Here you build the evidence that a model is worth having: **baselines**, a **training loop**, a **controlled ablation**\n", "(does masking the prompt in the loss matter?), **several seeds**, a **log record** per run and an honest **summary with intervals**.\n", "\n", "**Hardware:** CPU only, about 2 GB of RAM. The four training runs take roughly 20 to 30 seconds each on a laptop (about two minutes in all). PyTorch CPU build.\n", "**Default project:** the robot-command parser from milestone 1. Replace the data and the model size for your own project, but keep every check.\n", "**How it works:** cells marked `TODO` raise `NotImplementedError` until you write the code; each is followed by a self-check (`assert`).\n", "Theory: Lessons LM12-T2 (experiment design and reproducibility) and LM9-T2/T3 (statistics)." ] }, { "cell_type": "markdown", "id": "c002", "metadata": {}, "source": [ "## 1. Milestone 1, given\n", "\n", "The data code from milestone 1 (your solutions, or the reference ones) and a small GPT from Module 4 with the character tokenizer.\n", "Everything is seeded: running the notebook twice gives the same numbers." ] }, { "cell_type": "code", "id": "c003", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "import json, math, random, time\n", "import torch, torch.nn as nn, torch.nn.functional as F\n", "\n", "MOVE = [\"go to\", \"move to\", \"drive to\", \"head to\", \"navigate to\"]\n", "SPEED = [(\"slowly\", 0.2), (\"at half speed\", 0.5), (\"quickly\", 0.8), (\"\", 0.4)]\n", "MAX_BAY = 60\n", "HOLDOUT_BAYS = {56, 57, 58, 59, 60}\n", "PAD, END, UNK = \"\\x00\", \"\\x01\", \"\\x02\"\n", "\n", "def other_commands():\n", " out = []\n", " for word in [\"stop\", \"halt\", \"emergency stop\", \"freeze\"]:\n", " for suffix in [\"\", \" now\", \" please\", \" immediately\"]:\n", " out.append((word + suffix, {\"tool\": \"emergency_stop\"}, None))\n", " for word in [\"report battery\", \"battery level\", \"how much charge is left\", \"check battery\"]:\n", " for suffix in [\"\", \" please\", \" now\", \" again\"]:\n", " out.append((word + suffix, {\"tool\": \"get_battery\"}, None))\n", " for w in [\"hello\", \"ready\", \"done\", \"thanks\", \"goodbye\", \"waiting\", \"finished\", \"arrived\"]:\n", " for template in [\"say {}\", \"please say {}\", \"tell them {}\"]:\n", " out.append((template.format(w), {\"tool\": \"say\", \"text\": w}, None))\n", " return out\n", "\n", "def generate_pairs(n_move, seed):\n", " rng = random.Random(seed)\n", " moves = [(f\"{v} bay {b} {t}\".strip(), {\"tool\": \"move_to\", \"bay\": b, \"speed\": s}, b)\n", " for v in MOVE for b in range(1, MAX_BAY + 1) for t, s in SPEED]\n", " rng.shuffle(moves)\n", " items = moves[:n_move] + other_commands()\n", " rng.shuffle(items)\n", " return [{\"text\": t, \"target\": json.dumps(tg, separators=(\",\", \":\")), \"bay\": b} for t, tg, b in items]\n", "\n", "def split_dataset(pairs, seed=0, holdout=HOLDOUT_BAYS):\n", " ood = [p for p in pairs if p[\"bay\"] in holdout]\n", " rest = [p for p in pairs if p[\"bay\"] not in holdout]\n", " random.Random(seed).shuffle(rest)\n", " n = len(rest)\n", " a, b = int(0.8 * n), int(0.9 * n)\n", " return {\"train\": rest[:a], \"val\": rest[a:b], \"test\": rest[b:], \"test_ood\": ood}\n", "\n", "def build_vocab(train):\n", " chars = sorted({c for p in train for c in p[\"text\"] + \"\\n=>\" + p[\"target\"]})\n", " itos = [PAD, END, UNK] + chars\n", " return itos, {c: i for i, c in enumerate(itos)}\n", "\n", "def encode(p, stoi, mask_prompt=True):\n", " prompt = [stoi.get(c, stoi[UNK]) for c in p[\"text\"] + \"\\n=>\"]\n", " resp = [stoi.get(c, stoi[UNK]) for c in p[\"target\"]] + [stoi[END]]\n", " labels = ([-100] * len(prompt) if mask_prompt else prompt) + resp\n", " return prompt + resp, labels\n", "\n", "def make_batch(pairs, stoi, max_len, mask_prompt=True):\n", " X = torch.zeros(len(pairs), max_len, dtype=torch.long)\n", " Y = torch.full((len(pairs), max_len), -100, dtype=torch.long)\n", " for i, p in enumerate(pairs):\n", " ids, lab = encode(p, stoi, mask_prompt)\n", " X[i, :len(ids)] = torch.tensor(ids)\n", " Y[i, :len(lab)] = torch.tensor(lab)\n", " return X, Y\n", "\n", "pairs = generate_pairs(1000, seed=0)\n", "sp = split_dataset(pairs, seed=0)\n", "itos, stoi = build_vocab(sp[\"train\"])\n", "max_len = max(len(encode(p, stoi)[0]) for p in pairs)\n", "print({k: len(v) for k, v in sp.items()}, \"vocabulary\", len(itos), \"longest sequence\", max_len)" ] }, { "cell_type": "code", "id": "c004", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "class Block(nn.Module):\n", " def __init__(self, d, h):\n", " super().__init__()\n", " self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)\n", " self.attn = nn.MultiheadAttention(d, h, batch_first=True)\n", " self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))\n", " def forward(self, x):\n", " T = x.shape[1]\n", " mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)\n", " h = self.ln1(x)\n", " x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]\n", " return x + self.ff(self.ln2(x))\n", "\n", "class TinyGPT(nn.Module):\n", " def __init__(self, vocab, d=64, layers=2, heads=4, ctx=96):\n", " super().__init__()\n", " self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(ctx, d)\n", " self.blocks = nn.ModuleList(Block(d, heads) for _ in range(layers))\n", " self.ln, self.head = nn.LayerNorm(d), nn.Linear(d, vocab, bias=False)\n", " def forward(self, x):\n", " h = self.tok(x) + self.pos(torch.arange(x.shape[1]))\n", " for b in self.blocks:\n", " h = b(h)\n", " return self.head(self.ln(h))\n", "\n", "@torch.no_grad()\n", "def generate(model, text, max_new=60):\n", " ids = [stoi.get(c, stoi[UNK]) for c in text + \"\\n=>\"]\n", " n_prompt = len(ids)\n", " for _ in range(max_new):\n", " nxt = model(torch.tensor([ids]))[0, -1].argmax().item()\n", " if nxt == stoi[END]:\n", " break\n", " ids.append(nxt)\n", " return \"\".join(itos[i] for i in ids[n_prompt:])\n", "\n", "def exact_match(model, examples):\n", " model.eval()\n", " return [generate(model, p[\"text\"]) == p[\"target\"] for p in examples]" ] }, { "cell_type": "markdown", "id": "c005", "metadata": {}, "source": [ "## 2. TODO: two baselines you must beat\n", "\n", "A model result means nothing without baselines. Implement both, using only the **training** split:\n", "\n", "- `majority_exact(train, examples)`: the fraction of `examples` whose `target` equals the most frequent target in `train`.\n", "- `nearest_neighbour_exact(train, examples)`: for each example find the training example whose **text** has the highest Jaccard similarity of character bigrams\n", " (`|A ∩ B| / |A ∪ B|` of the sets of two-character substrings; the first training example wins ties) and predict its `target`; return the fraction correct.\n", "\n", "A third baseline, a hand-written **rule-based parser**, is given. Compare honestly." ] }, { "cell_type": "code", "id": "c006", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "def majority_exact(train, examples):\n", " raise NotImplementedError(\"TODO: predict the most common training target for every example\")\n", "\n", "def bigrams(t):\n", " return {t[i:i + 2] for i in range(len(t) - 1)}\n", "\n", "def nearest_neighbour_exact(train, examples):\n", " raise NotImplementedError(\"TODO: copy the target of the most similar training text\")\n", "\n", "def rule_parser(text):\n", " \"\"\"Given: a hand-written parser. It knows the phrasing of the generator exactly.\"\"\"\n", " words = text.split()\n", " if \"bay\" in words:\n", " bay = int(words[words.index(\"bay\") + 1])\n", " speed = 0.2 if \"slowly\" in words else 0.5 if \"half\" in words else 0.8 if \"quickly\" in words else 0.4\n", " return json.dumps({\"tool\": \"move_to\", \"bay\": bay, \"speed\": speed}, separators=(\",\", \":\"))\n", " if any(w in words for w in (\"stop\", \"halt\", \"freeze\")):\n", " return '{\"tool\":\"emergency_stop\"}'\n", " if \"battery\" in words or \"charge\" in words:\n", " return '{\"tool\":\"get_battery\"}'\n", " word = [w for w in words if w in (\"hello\", \"ready\", \"done\", \"thanks\", \"goodbye\", \"waiting\", \"finished\", \"arrived\")][0]\n", " return json.dumps({\"tool\": \"say\", \"text\": word}, separators=(\",\", \":\"))" ] }, { "cell_type": "markdown", "id": "c007", "metadata": {}, "source": [ "### Self-check 1: the baselines" ] }, { "cell_type": "code", "id": "c008", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "tiny = [{\"text\": \"aa\", \"target\": \"X\"}, {\"text\": \"ab\", \"target\": \"Y\"}, {\"text\": \"ac\", \"target\": \"X\"}]\n", "assert majority_exact(tiny, tiny) == 2 / 3\n", "assert nearest_neighbour_exact([{\"text\": \"stop\", \"target\": \"S\"}, {\"text\": \"go to bay 3\", \"target\": \"G\"}], [{\"text\": \"go to bay 3 now\", \"target\": \"G\"}, {\"text\": \"stop now\", \"target\": \"G\"}]) == 0.5\n", "\n", "base = {}\n", "for name in (\"val\", \"test\", \"test_ood\"):\n", " base[name] = {\"majority\": majority_exact(sp[\"train\"], sp[name]),\n", " \"nearest_neighbour\": nearest_neighbour_exact(sp[\"train\"], sp[name]),\n", " \"rule\": sum(rule_parser(p[\"text\"]) == p[\"target\"] for p in sp[name]) / len(sp[name])}\n", "print(json.dumps(base, indent=1))\n", "assert base[\"test\"][\"majority\"] < 0.05 and base[\"test\"][\"nearest_neighbour\"] < 0.2\n", "assert base[\"test\"][\"rule\"] == 1.0 and base[\"test_ood\"][\"rule\"] == 1.0\n", "print(\"self-check 1 passed: trivial baselines are weak; the rule-based parser is perfect on this synthetic task (remember this in the report)\")" ] }, { "cell_type": "markdown", "id": "c009", "metadata": {}, "source": [ "## 3. TODO: the training loop\n", "\n", "`train(model, X, Y, steps, lr=3e-3, bs=32, seed=0)` trains with AdamW (weight decay 0.01) for `steps` steps and returns the list of the loss of every step.\n", "\n", "- Batch: `idx = torch.randint(0, X.shape[0], (bs,), generator=g)` with `g = torch.Generator().manual_seed(seed)` created **once** before the loop.\n", "- Inputs `X[idx][:, :-1]`, targets `Y[idx][:, 1:]` (next-token prediction); cross-entropy over all positions with `ignore_index=-100`.\n", "- Learning rate at step `s`: `lr * min(1, (s + 1) / 30) * 0.5 * (1 + cos(pi * s / steps))` (warm-up over 30 steps, then cosine decay).\n", "- Zero the gradients, `backward`, `step`; append `loss.item()`." ] }, { "cell_type": "code", "id": "c010", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "def train(model, X, Y, steps, lr=3e-3, bs=32, seed=0):\n", " raise NotImplementedError(\"TODO: AdamW loop with warm-up, cosine decay and a seeded batch generator\")" ] }, { "cell_type": "markdown", "id": "c011", "metadata": {}, "source": [ "### Self-check 2: the loss falls and a run is reproducible" ] }, { "cell_type": "code", "id": "c012", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "X, Y = make_batch(sp[\"train\"], stoi, max_len)\n", "def fresh(seed):\n", " torch.manual_seed(seed)\n", " return TinyGPT(len(itos), d=32, layers=1, heads=4, ctx=max_len)\n", "a = train(fresh(0), X, Y, 40, seed=0)\n", "b = train(fresh(0), X, Y, 40, seed=0)\n", "c = train(fresh(1), X, Y, 40, seed=1)\n", "assert a == b, \"the same seeds must give exactly the same loss curve on CPU\"\n", "assert a != c, \"different seeds should give different curves\"\n", "assert a[0] > 3 and a[-1] < 0.9 * a[0], f\"the loss should already fall in 40 steps: {a[0]:.2f} -> {a[-1]:.2f}\"\n", "print(f\"loss {a[0]:.2f} -> {a[-1]:.2f} in 40 steps; identical for identical seeds\")\n", "print(\"self-check 2 passed: training loop\")" ] }, { "cell_type": "markdown", "id": "c013", "metadata": {}, "source": [ "## 4. TODO: log every run, summarise with intervals\n", "\n", "- `make_record(name, config, losses, val_correct, ood_correct, seconds)` returns a dict with the keys\n", " `name`, `config` (a copy), `final_loss` (mean of the last 20 losses), `val_exact` and `ood_exact` (fractions of `True` in the two lists),\n", " `n_val`, `n_ood`, `seconds` (rounded to 1 decimal) and `torch` (the version string).\n", "- `wilson(k, n, z=1.96)` returns the lower and upper end of the Wilson score interval (rounded to 3 decimals) for `k` successes in `n` trials (Lesson LM11-T5).\n", "- `summarise(records)` groups the records by `name` and returns, per name, `{\"runs\", \"val_mean\", \"ood_mean\", \"ood_ci\"}`: the number of runs, the mean validation and OOD exact match\n", " over runs, and the Wilson interval of the **pooled** OOD results (total correct over total examples of all runs of that name)." ] }, { "cell_type": "code", "id": "c014", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "def make_record(name, config, losses, val_correct, ood_correct, seconds):\n", " raise NotImplementedError(\"TODO: collect everything needed to reproduce and compare a run\")\n", "\n", "def wilson(k, n, z=1.96):\n", " raise NotImplementedError(\"TODO: Wilson score interval\")\n", "\n", "def summarise(records):\n", " raise NotImplementedError(\"TODO: group by name; means and the pooled Wilson interval of the OOD results\")" ] }, { "cell_type": "markdown", "id": "c015", "metadata": {}, "source": [ "### Self-check 3: the bookkeeping on made-up numbers" ] }, { "cell_type": "code", "id": "c016", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "r = make_record(\"t\", {\"a\": 1}, [1.0] * 30 + [0.5] * 20, [True, False, True, True], [True, False], 3.14159)\n", "assert r[\"final_loss\"] == 0.5 and r[\"val_exact\"] == 0.75 and r[\"ood_exact\"] == 0.5 and r[\"n_val\"] == 4 and r[\"seconds\"] == 3.1 and r[\"torch\"] == torch.__version__\n", "assert r[\"config\"] == {\"a\": 1} and r[\"config\"] is not None\n", "assert wilson(18, 20) == (0.699, 0.972) and wilson(90, 100) == (0.826, 0.945)\n", "fake = [make_record(\"A\", {}, [0.1] * 20, [True] * 10, [True] * 6 + [False] * 4, 1), make_record(\"A\", {}, [0.1] * 20, [True] * 10, [True] * 8 + [False] * 2, 1), make_record(\"B\", {}, [0.1] * 20, [True] * 9 + [False], [True] * 5 + [False] * 5, 1)]\n", "sm = summarise(fake)\n", "assert sm[\"A\"][\"runs\"] == 2 and abs(sm[\"A\"][\"ood_mean\"] - 0.7) < 1e-12 and sm[\"A\"][\"ood_ci\"] == wilson(14, 20)\n", "assert sm[\"B\"][\"runs\"] == 1 and sm[\"B\"][\"val_mean\"] == 0.9\n", "print(\"self-check 3 passed: records, Wilson interval, summary\")" ] }, { "cell_type": "markdown", "id": "c017", "metadata": {}, "source": [ "## 5. The experiment (given): a controlled ablation with two seeds\n", "\n", "**Question:** does masking the prompt in the loss (training only on the response) change the result? Everything else is identical: model (2 layers, width 64, about 100,000 parameters),\n", "1,000 steps, batch 32, the same data. Two seeds per variant. Evaluation: **exact match** of the generated JSON string with the target, greedy decoding, on `val` (same distribution)\n", "and `test_ood` (bay numbers never seen in training)." ] }, { "cell_type": "code", "id": "c018", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "def run(name, mask_prompt, seed, steps=1000):\n", " X, Y = make_batch(sp[\"train\"], stoi, max_len, mask_prompt)\n", " torch.manual_seed(seed)\n", " model = TinyGPT(len(itos), d=64, layers=2, heads=4, ctx=max_len)\n", " t0 = time.time()\n", " losses = train(model, X, Y, steps, seed=seed)\n", " val, ood = exact_match(model, sp[\"val\"]), exact_match(model, sp[\"test_ood\"])\n", " return make_record(name, {\"mask_prompt\": mask_prompt, \"steps\": steps, \"d\": 64, \"layers\": 2, \"seed\": seed}, losses, val, ood, time.time() - t0), model\n", "\n", "records = []\n", "for name, mask in ((\"masked\", True), (\"unmasked\", False)):\n", " for seed in (0, 1):\n", " rec, _ = run(name, mask, seed)\n", " records.append(rec)\n", " print(f\"{name:9s} seed {seed}: final loss {rec['final_loss']:.4f} val {rec['val_exact']:.3f} ood {rec['ood_exact']:.3f} ({rec['seconds']} s)\")" ] }, { "cell_type": "markdown", "id": "c019", "metadata": {}, "source": [ "### Self-check 4: the model beats the learned baselines, and the summary" ] }, { "cell_type": "code", "id": "c020", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "summary = summarise(records)\n", "print(json.dumps(summary, indent=1))\n", "for rec in records:\n", " assert rec[\"val_exact\"] >= 0.9, f\"{rec['name']} seed {rec['config']['seed']}: validation exact match only {rec['val_exact']:.3f}\"\n", " assert rec[\"val_exact\"] > base[\"val\"][\"majority\"] + 0.5 and rec[\"val_exact\"] > base[\"val\"][\"nearest_neighbour\"] + 0.5\n", " assert rec[\"ood_exact\"] >= 0.4, \"the model should generalise at least partly to unseen bay numbers\"\n", " assert rec[\"final_loss\"] < 0.2\n", "lo, hi = summary[\"masked\"][\"ood_ci\"], summary[\"unmasked\"][\"ood_ci\"]\n", "overlap = lo[0] <= hi[1] and hi[0] <= lo[1]\n", "print(\"OOD intervals masked\", lo, \"unmasked\", hi, \"-> overlap\" if overlap else \"-> separated\")\n", "print(\"self-check 4 passed: experiment logged and summarised\")" ] }, { "cell_type": "markdown", "id": "c021", "metadata": {}, "source": [ "## 6. Read the evidence, then write it down\n", "\n", "Use the printed table. Suggested wording for the report (Lesson LM12-T4), to be adapted to **your** numbers:\n", "\n", "- *Baselines:* majority and nearest-neighbour baselines score near 0; a hand-written rule parser scores 1.0 on this synthetic task, **so a learned model is not needed here**; the exercise is about the method.\n", "- *Main result:* exact match on `val` and `test_ood` for each variant, with the number of runs and the Wilson interval of the pooled `test_ood` results.\n", "- *Ablation:* compare the intervals **and** the spread between seeds. The pooled Wilson interval treats the runs as one big sample and ignores that two seeds can disagree: in the reference run the unmasked variant\n", " scored 0.787 and 0.551 on the OOD slice with the two seeds, a spread as large as the gap between the variants. With two seeds, report the per-seed values, say that the evidence is suggestive, and run more seeds before claiming an effect.\n", "- *Limits:* one data generator, one model size, two seeds, character-level model, synthetic data.\n", "\n", "**Reflect:** (1) Which baseline would a sceptical reviewer demand that you did not run? (2) What would you change to test whether masking matters for *larger* models? (3) What is the smallest change in\n", "the setup that could invalidate the comparison (hint: the seed also controls the batch order)?\n", "\n", "**Done when** all four self-checks print \"passed\"." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }