{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM7-L3: Instruction-tune a tiny model and evaluate it\n",
    "\n",
    "A pretrained (\"base\") model continues text; it does not answer. In this notebook you will see that with a small GPT, then turn it into an instruction follower with **supervised\n",
    "instruction fine-tuning**: format examples with a template, mask the loss so that only the response (and the end-of-turn token) is learned, train, generate until the end token,\n",
    "and evaluate on **held-out** examples with an exact-match criterion. You will also check an evaluation set for **contamination**.\n",
    "\n",
    "**The skills are tiny on purpose.** The instructions are \"write the word in capitals\", \"say the word twice\" and \"spell the word backwards\", applied to random letter strings.\n",
    "A 350,000-parameter character-level model can learn these in under a minute on a CPU. This shows the mechanics of instruction tuning; it says nothing about the quality of\n",
    "an instruction-tuned language model on real instructions.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, 1 to 2 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 LM7-T1 to LM7-T5."
   ]
  },
  {
   "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: tasks, data and tokenizer\n",
    "\n",
    "Each example has a task (the instruction), an input word, and the correct response. Training words and evaluation words are **different random strings**, so\n",
    "the model must learn the rule, not memorise answers. Characters are tokens; two special tokens are added: `END` (end of the response) and `PAD`."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TASKS = {\"uppercase\": lambda w: w.upper(), \"repeat\": lambda w: w + \" \" + w, \"reverse\": lambda w: w[::-1]}\n",
    "INSTRUCTION = {\"uppercase\": \"write the word in capitals\", \"repeat\": \"say the word twice\", \"reverse\": \"spell the word backwards\"}\n",
    "\n",
    "rng = random.Random(0)\n",
    "def random_word():\n",
    "    return \"\".join(rng.choice(\"abcdefghijklmnopqrstuvwxyz\") for _ in range(rng.randint(3, 5)))\n",
    "pool = set()\n",
    "while len(pool) < 460:\n",
    "    pool.add(random_word())\n",
    "pool = sorted(pool)\n",
    "rng.shuffle(pool)\n",
    "train_words, eval_words = pool[:400], pool[400:]\n",
    "print(len(train_words), \"training words,\", len(eval_words), \"held-out words, e.g.\", eval_words[:5])\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",
    "decode = lambda ids: \"\".join(itos[i] for i in ids)\n",
    "print(\"vocabulary\", VOCAB)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. TODO: the prompt template\n",
    "\n",
    "`build_prompt(task, word)` returns exactly this text (three sections, blank lines between them, and the response section left open):\n",
    "\n",
    "```\n",
    "### Instruction:\n",
    "<instruction text of the task>\n",
    "\n",
    "### Input:\n",
    "<word>\n",
    "\n",
    "### Response:\n",
    "```\n",
    "\n",
    "(it ends with a newline after `### Response:`; the response will be appended to it)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def build_prompt(task, word):\n",
    "    raise NotImplementedError(\"TODO: the template above, using INSTRUCTION[task] and word\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "p = build_prompt(\"repeat\", \"box\")\n",
    "assert p == \"### Instruction:\\nsay the word twice\\n\\n### Input:\\nbox\\n\\n### Response:\\n\", \"the prompt does not match the template\"\n",
    "assert p.endswith(\"### Response:\\n\")\n",
    "print(\"self-check 1 passed: template\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 3. TODO: one training example with a masked label sequence\n",
    "\n",
    "`make_pair(prompt, response)` returns `(ids, labels)`, two 1-D integer tensors of the same length:\n",
    "\n",
    "* `ids` = the characters of `prompt + response`, followed by the `END_ID` token;\n",
    "* `labels` = `-100` for every prompt position, then the response ids, then `END_ID`. The value `-100` means \"ignore this position in the loss\" (Lesson LM7-T4)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def make_pair(prompt, response):\n",
    "    raise NotImplementedError(\"TODO: ids and labels with -100 over the prompt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "### Self-check 2"
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ids, labels = make_pair(p, \"box box\")\n",
    "assert ids.shape == labels.shape and ids.dtype == torch.long\n",
    "assert (labels[: len(p)] == -100).all(), \"all prompt positions must be -100\"\n",
    "assert labels[-1].item() == END_ID and ids[-1].item() == END_ID, \"the end-of-turn token must be learned too\"\n",
    "assert decode(labels[len(p):-1].tolist()) == \"box box\", \"after the prompt come the response ids, then END\"\n",
    "assert decode(ids[:-1].tolist()) == p + \"box box\"\n",
    "print(\"self-check 2 passed: masked labels\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "## 4. TODO: batching with padding\n",
    "\n",
    "`collate(pairs)` takes a list of `(ids, labels)` of different lengths and returns `(X, Y)` for next-token training:\n",
    "\n",
    "* pad `ids` on the right with `PAD` and `labels` on the right with `-100` to the longest length in the batch;\n",
    "* the model input is `X = ids[:, :-1]` and the target is `Y = labels[:, 1:]` (each position predicts the next token)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def collate(pairs):\n",
    "    raise NotImplementedError(\"TODO: pad, then shift by one\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "### Self-check 3"
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "a = make_pair(build_prompt(\"uppercase\", \"ab\"), \"AB\")\n",
    "b = make_pair(build_prompt(\"uppercase\", \"abcde\"), \"ABCDE\")\n",
    "X, Y = collate([a, b])\n",
    "assert X.shape == Y.shape == (2, len(b[0]) - 1), f\"unexpected shapes {tuple(X.shape)} {tuple(Y.shape)}\"\n",
    "assert (X[0, len(a[0]):] == PAD).all(), \"the shorter example is padded with PAD\"\n",
    "assert (Y[0, len(a[0]) - 1:] == -100).all(), \"padding positions must be ignored in the loss\"\n",
    "first_response = len(build_prompt(\"uppercase\", \"ab\"))\n",
    "assert Y[0, first_response - 2].item() == -100 and Y[0, first_response - 1].item() == stoi[\"A\"], \"the first response token is the first target that counts\"\n",
    "print(\"self-check 3 passed: collate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "## 5. TODO: the masked loss\n",
    "\n",
    "`masked_loss(logits, targets)` with logits `(B, n, V)` and targets `(B, n)` (containing `-100`) returns the mean cross-entropy over the positions whose target is not `-100`\n",
    "(use `F.cross_entropy(..., ignore_index=-100)` after flattening)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def masked_loss(logits, targets):\n",
    "    raise NotImplementedError(\"TODO: cross-entropy that ignores -100\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "### Self-check 4"
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "g = torch.Generator().manual_seed(0)\n",
    "logits = torch.randn(2, 6, 7, generator=g)\n",
    "targets = torch.tensor([[-100, -100, 3, 1, 5, -100], [-100, 2, 2, -100, -100, -100]])\n",
    "manual = torch.stack([F.cross_entropy(logits[b, t:t + 1], targets[b, t:t + 1]) for b in range(2) for t in range(6) if targets[b, t] != -100]).mean()\n",
    "assert torch.allclose(masked_loss(logits, targets), manual, atol=1e-6), \"the loss must be the mean over the non-ignored positions only\"\n",
    "changed = logits.clone(); changed[0, 0] += 100; changed[1, 5] -= 50                  # change logits at ignored positions\n",
    "assert torch.allclose(masked_loss(changed, targets), masked_loss(logits, targets)), \"ignored positions must not influence the loss\"\n",
    "print(\"self-check 4 passed: masked loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 6. Given: the model and a \"base\" model\n",
    "\n",
    "A small GPT (3 blocks, width 96). The **base model** is pretrained for a few seconds on plain text, a stream of random words: it has never seen an instruction."
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "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 GPT(nn.Module):\n",
    "    def __init__(self, vocab, context=96, d=96, heads=4, layers=3):\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",
    "def pretrain(model, steps=200):\n",
    "    text = \" \".join(random.Random(1).choice(pool) for _ in range(4000))\n",
    "    data = torch.tensor(encode(text))\n",
    "    g = torch.Generator().manual_seed(1)\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.1)\n",
    "    for _ in range(steps):\n",
    "        s = torch.randint(0, len(data) - 65, (32,), generator=g)\n",
    "        x = torch.stack([data[i : i + 64] for i in s]); y = torch.stack([data[i + 1 : i + 65] for i in s])\n",
    "        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "    return loss.item()\n",
    "\n",
    "torch.manual_seed(0)\n",
    "base = GPT(VOCAB)\n",
    "print(\"parameters:\", sum(p.numel() for p in base.parameters()))\n",
    "t0 = time.time()\n",
    "print(f\"base model pretraining loss {pretrain(base):.2f} in {time.time() - t0:.0f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "## 7. TODO: generate a response\n",
    "\n",
    "`generate_response(model, prompt, max_new=24)`: greedy decoding (arg-max) starting from the prompt; stop when the model emits `END_ID` (do not include it) or after `max_new` tokens;\n",
    "return the generated characters as a string. Put the model in evaluation mode and use `torch.no_grad()`, and feed at most `model.context` tokens."
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def generate_response(model, prompt, max_new=24):\n",
    "    raise NotImplementedError(\"TODO: greedy loop that stops at END_ID\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "### Self-check 5: what does the base model do with an instruction?"
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "sample = build_prompt(\"uppercase\", \"box\")\n",
    "base_out = generate_response(base, sample)\n",
    "print(\"base model:\", repr(base_out))\n",
    "assert isinstance(base_out, str) and END not in base_out and len(base_out) <= 24\n",
    "assert base_out != \"BOX\", \"a base model that has never seen an instruction should not answer it\"\n",
    "assert generate_response(base, sample) == base_out, \"greedy decoding is deterministic\"\n",
    "print(\"self-check 5 passed: the base model does not follow the instruction\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "## 8. TODO: evaluation and contamination\n",
    "\n",
    "An evaluation item is a pair `(task, word)`; the correct response is `TASKS[task](word)`.\n",
    "\n",
    "* `evaluate(model, items)` returns a dict `task -> exact-match accuracy` (the share of items for which `generate_response` equals the correct response exactly).\n",
    "* `overlap(train_items, eval_items)` returns the list of evaluation items that also occur in the training items (contamination)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c028",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def evaluate(model, items):\n",
    "    raise NotImplementedError(\"TODO: exact match per task\")\n",
    "\n",
    "def overlap(train_items, eval_items):\n",
    "    raise NotImplementedError(\"TODO: evaluation items that appear in the training items\")\n",
    "\n",
    "train_items = [(t, w) for t in TASKS for w in train_words]\n",
    "clean_items = [(t, w) for t in TASKS for w in eval_words]\n",
    "small_words = train_words[:40]                                        # a small training set, used later to build a model that memorises\n",
    "small_items = [(t, w) for t in TASKS for w in small_words]\n",
    "leaky_items = clean_items[::3] + small_items[::2]                     # a \"benchmark\": 60 held-out items plus 60 items that were in the small training set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c029",
   "metadata": {},
   "source": [
    "### Self-check 6"
   ]
  },
  {
   "cell_type": "code",
   "id": "c030",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "assert overlap(train_items, clean_items) == [], \"the held-out words never appear in training\"\n",
    "assert len(overlap(small_items, leaky_items)) == 60, \"the leaky set contains 60 items of the small training set\"\n",
    "acc_base = evaluate(base, clean_items[::3])\n",
    "print(\"base model accuracy:\", {k: round(v, 2) for k, v in acc_base.items()})\n",
    "assert all(0.0 <= v <= 1.0 for v in acc_base.values()) and max(acc_base.values()) <= 0.1, \"the base model should be near 0 on every task\"\n",
    "print(\"self-check 6 passed: evaluation and contamination check\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c031",
   "metadata": {},
   "source": [
    "## 9. Given: instruction-tune and compare\n",
    "\n",
    "Fine-tune the base model on 1,200 examples (400 words x 3 tasks) with the masked loss for 1,000 steps (about 30 seconds on a CPU), then evaluate on held-out words.\n",
    "Then build a second model that **memorises**: the same recipe on only 120 examples (40 words). It will score well on the leaky benchmark and poorly on clean held-out words."
   ]
  },
  {
   "cell_type": "code",
   "id": "c032",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def instruction_tune(model, items, steps=1000, lr=1e-3, batch=32, seed=2):\n",
    "    data = [make_pair(build_prompt(t, w), TASKS[t](w)) for t, w in items]\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n",
    "    for step in range(1, steps + 1):\n",
    "        ix = torch.randint(0, len(data), (batch,), generator=g)\n",
    "        X, Y = collate([data[i] for i in ix])\n",
    "        loss = masked_loss(model(X), Y)\n",
    "        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()\n",
    "        if step % 250 == 0:\n",
    "            print(f\"  step {step:4d}  masked loss {loss.item():.4f}\")\n",
    "    return model\n",
    "\n",
    "import copy\n",
    "tuned = copy.deepcopy(base)\n",
    "t0 = time.time()\n",
    "instruction_tune(tuned, train_items)\n",
    "print(f\"instruction tuning took {time.time() - t0:.0f} s\")\n",
    "for prompt_task, w in [(\"uppercase\", \"zqk\"), (\"repeat\", \"zqk\"), (\"reverse\", \"zqk\")]:\n",
    "    print(f\"  {prompt_task:9s} zqk ->\", repr(generate_response(tuned, build_prompt(prompt_task, w))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c033",
   "metadata": {},
   "source": [
    "### Self-check 7: held-out accuracy, and what a leak does to a memorising model"
   ]
  },
  {
   "cell_type": "code",
   "id": "c034",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "acc_clean = evaluate(tuned, clean_items)\n",
    "mean = lambda d: sum(d.values()) / len(d)\n",
    "print(\"tuned model, held-out words:\", {k: round(v, 2) for k, v in acc_clean.items()}, \" mean\", round(mean(acc_clean), 2))\n",
    "assert min(acc_clean.values()) >= 0.4 and mean(acc_clean) >= 0.6, \"instruction tuning should reach a mean held-out exact-match of at least 0.6\"\n",
    "assert mean(acc_clean) > max(acc_base.values()) + 0.4, \"the tuned model should be far better than the base model\"\n",
    "\n",
    "memoriser = copy.deepcopy(base)\n",
    "instruction_tune(memoriser, small_items, steps=600)\n",
    "acc_mem_clean = evaluate(memoriser, clean_items[::3])\n",
    "acc_mem_leaky = evaluate(memoriser, leaky_items)\n",
    "print(\"memorising model, clean held-out :\", {k: round(v, 2) for k, v in acc_mem_clean.items()}, \" mean\", round(mean(acc_mem_clean), 2))\n",
    "print(\"memorising model, leaky benchmark:\", {k: round(v, 2) for k, v in acc_mem_leaky.items()}, \" mean\", round(mean(acc_mem_leaky), 2))\n",
    "assert mean(acc_mem_leaky) > mean(acc_mem_clean) + 0.2, \"a leak should inflate the score of a model that memorised the leaked items\"\n",
    "print(\"self-check 7 passed: the tuned model follows the instructions on unseen words; leaked items inflate a memorising model's score\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c035",
   "metadata": {},
   "source": [
    "## 10. Reflect\n",
    "\n",
    "1. Why was the base model unable to answer, although it had learned the characters and words? What exactly did instruction tuning add?\n",
    "2. What would have happened to the loss and to generation if the prompt tokens had **not** been masked? What if the `END` token had been left out of the labels?\n",
    "3. The memorising model scored much higher on the leaky benchmark than on clean words. Why? What does a leak do to a score, and how would you detect one in a benchmark you did not build yourself?\n",
    "4. The three skills are rule-based, and the model got them right on unseen words only partly. What would you change (data, steps, model size) to improve the weakest task, and how would you check it without contaminating the evaluation?\n",
    "\n",
    "**Done when** all seven self-checks print \"passed\"."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}