{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM6-L2: Fine-tune a language model for classification\n",
    "\n",
    "You will take a small GPT, give it a **classification head** instead of the vocabulary head, and train it on a labelled task with different parts frozen. You will also write\n",
    "the metrics (confusion matrix, precision, recall, F1) and look at the errors.\n",
    "\n",
    "**The task is synthetic.** Each example is a sentence such as \"the drone moves to the red box.\" and the label is 1 (\"mobile unit is moving\") only when the subject is the drone or the cart\n",
    "**and** the verb is \"moves to\"; otherwise 0. About 12% of the examples are positive, so the classes are **imbalanced**: a classifier that always says 0 is right 88% of the time and useless.\n",
    "The rule is simple on purpose, so that the whole notebook runs in about a minute on a CPU.\n",
    "\n",
    "**Honest scope.** The backbone is pretrained here, for a few seconds, on the unlabelled sentences themselves, not on a large corpus. The results show the mechanics (head, pooling,\n",
    "freezing, metrics); they do not show how much a real pretrained model helps on real tasks. The notebook prints a comparison with training from scratch; on this\n",
    "easy toy task and a single seed, the difference is small.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, about one minute. No downloads. The book's version of this chapter fine-tunes the pretrained GPT-2 small on spam detection; see the link on the lab page.\n",
    "\n",
    "Cells marked `TODO` raise `NotImplementedError` until you write the code. Each is followed by a self-check (`assert`). Theory: Lessons LM6-T1 to LM6-T5."
   ]
  },
  {
   "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__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. Given: data\n",
    "\n",
    "An unlabelled stream of sentences (for pretraining) and a labelled set. Characters are tokens; index `PAD` is used to pad short sentences to a common length."
   ]
  },
  {
   "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 sentence(r):\n",
    "    return r.choice(SUBJ), r.choice(VERB), f\"the {r.choice(COLOR)} {r.choice(NOUN)}\"\n",
    "\n",
    "def make_stream(n=1200, seed=0):\n",
    "    r = random.Random(seed)\n",
    "    return \" \".join(f\"{s} {v} {o}.\" for s, v, o in (sentence(r) for _ in range(n)))\n",
    "\n",
    "def make_labelled(n, seed):\n",
    "    r = random.Random(seed)\n",
    "    texts, labels = [], []\n",
    "    for _ in range(n):\n",
    "        s, v, o = sentence(r)\n",
    "        texts.append(f\"{s} {v} {o}.\")\n",
    "        labels.append(int(s in (\"the drone\", \"the cart\") and v == \"moves to\"))\n",
    "    return texts, labels\n",
    "\n",
    "stream = make_stream()\n",
    "chars = sorted(set(stream))\n",
    "stoi = {c: i for i, c in enumerate(chars)}\n",
    "VOCAB = len(chars)\n",
    "PAD = VOCAB                                   # extra index for padding\n",
    "stream_ids = torch.tensor([stoi[c] for c in stream])\n",
    "\n",
    "texts, labels = make_labelled(1200, seed=5)\n",
    "def encode(batch_texts):\n",
    "    L = max(len(t) for t in batch_texts)\n",
    "    X = torch.full((len(batch_texts), L), PAD)\n",
    "    lengths = torch.tensor([len(t) for t in batch_texts])\n",
    "    for i, t in enumerate(batch_texts):\n",
    "        X[i, : len(t)] = torch.tensor([stoi[c] for c in t])\n",
    "    return X, lengths\n",
    "\n",
    "X_all, len_all = encode(texts)\n",
    "y_all = torch.tensor(labels)\n",
    "X_tr, len_tr, y_tr = X_all[:200], len_all[:200], y_all[:200]            # 200 labelled examples for training\n",
    "X_va, len_va, y_va = X_all[700:], len_all[700:], y_all[700:]            # 500 held-out examples\n",
    "print(\"positive rate:\", round(y_all.float().mean().item(), 3), \"| train\", len(y_tr), \"| validation\", len(y_va))\n",
    "print(\"accuracy of always answering 0 on validation:\", round(1 - y_va.float().mean().item(), 3))\n",
    "print(texts[0], labels[0], \"|\", next(t for t, l in zip(texts, labels) if l == 1), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. Given: a small GPT backbone and a short pretraining\n",
    "\n",
    "The same architecture as before (width 64, 4 heads, 2 blocks, context 64) with a method `features` that returns the final hidden states `(B, n, d)` **without** the vocabulary head."
   ]
  },
  {
   "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 Backbone(nn.Module):\n",
    "    def __init__(self, vocab, context=64, d=64, heads=4, layers=2):\n",
    "        super().__init__()\n",
    "        self.context, self.d = context, d\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 = nn.LayerNorm(d)\n",
    "        self.lm_head = nn.Linear(d, vocab - 1, bias=False)     # used only for pretraining (no logit for PAD)\n",
    "\n",
    "    def features(self, idx):\n",
    "        return self.norm(self.blocks(self.tok(idx) + self.pos(torch.arange(idx.shape[1]))))\n",
    "\n",
    "    def forward(self, idx):                                     # language-model logits, for pretraining\n",
    "        return self.lm_head(self.features(idx))\n",
    "\n",
    "def pretrain(model, steps=300, lr=3e-3, seed=1):\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.1)\n",
    "    cut = int(0.9 * len(stream_ids))\n",
    "    data = stream_ids[:cut]\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(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()\n",
    "    return loss.item()\n",
    "\n",
    "torch.manual_seed(0)\n",
    "pretrained = Backbone(VOCAB + 1)\n",
    "t0 = time.time()\n",
    "print(f\"pretraining loss {pretrain(pretrained):.3f} after {time.time() - t0:.0f} s (ln(vocab) was {math.log(VOCAB):.2f})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "## 3. TODO: the classification model\n",
    "\n",
    "Wrap a backbone and add `self.cls = nn.Linear(d, n_classes)`. In `forward(idx, lengths)`:\n",
    "\n",
    "1. get the hidden states `h = self.backbone.features(idx)`, shape `(B, n, d)`;\n",
    "2. pick, for each sequence `b`, the hidden state of its **last real token** (position `lengths[b] - 1`, not the last padded position): shape `(B, d)`;\n",
    "3. return `self.cls(...)`, shape `(B, n_classes)`.\n",
    "\n",
    "Why the last token? With the causal mask, position `i` sees only positions `<= i`; only the last real token has read the whole sentence (Lesson LM6-T2).\n",
    "`pool=\"first\"` is provided so that you can see what goes wrong when the first token is used."
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class Classifier(nn.Module):\n",
    "    def __init__(self, backbone, n_classes=2, pool=\"last\"):\n",
    "        super().__init__()\n",
    "        self.backbone, self.pool = backbone, pool\n",
    "        self.cls = nn.Linear(backbone.d, n_classes)\n",
    "\n",
    "    def forward(self, idx, lengths):\n",
    "        h = self.backbone.features(idx)\n",
    "        pos = lengths - 1 if self.pool == \"last\" else torch.zeros_like(lengths)\n",
    "        pooled = h[torch.arange(idx.shape[0]), pos]\n",
    "        return self.cls(pooled)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "### Self-check 1"
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "clf = Classifier(copy.deepcopy(pretrained))\n",
    "logits = clf(X_tr[:6], len_tr[:6])\n",
    "assert logits.shape == (6, 2), f\"logits should be (batch, n_classes), got {tuple(logits.shape)}\"\n",
    "h = clf.backbone.features(X_tr[:6])\n",
    "manual = torch.stack([clf.cls(h[i, len_tr[i] - 1]) for i in range(6)])\n",
    "assert torch.allclose(logits, manual, atol=1e-6), \"use the hidden state at position lengths-1 for each sequence\"\n",
    "altered = X_tr[:6].clone()\n",
    "for i in range(6):\n",
    "    altered[i, len_tr[i]:] = 3                      # change only the padding after the end of each sentence\n",
    "assert torch.allclose(logits, clf(altered, len_tr[:6]), atol=1e-6), \"the output must not depend on the padding after the last real token\"\n",
    "print(\"self-check 1 passed: classification head and pooling\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "## 4. TODO: freezing\n",
    "\n",
    "`set_trainable(model, mode)` sets `requires_grad` on the parameters:\n",
    "\n",
    "* `\"head\"`: only `model.cls` is trainable;\n",
    "* `\"last\"`: `model.cls`, the last Transformer block (`model.backbone.blocks[-1]`) and the final norm (`model.backbone.norm`);\n",
    "* `\"all\"`: everything except `model.backbone.lm_head` (the language-model head is not used for classification).\n",
    "\n",
    "Return the number of trainable parameters."
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def set_trainable(model, mode):\n",
    "    for p in model.parameters():\n",
    "        p.requires_grad = False\n",
    "    groups = [model.cls]\n",
    "    if mode in (\"last\", \"all\"):\n",
    "        groups += [model.backbone.blocks[-1], model.backbone.norm]\n",
    "    if mode == \"all\":\n",
    "        groups += [model.backbone.tok, model.backbone.pos, model.backbone.blocks]\n",
    "    for g in groups:\n",
    "        for p in g.parameters():\n",
    "            p.requires_grad = True\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "### Self-check 2"
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "n_head = set_trainable(clf, \"head\")\n",
    "assert n_head == 64 * 2 + 2, f\"head only should train 130 parameters, got {n_head}\"\n",
    "n_last = set_trainable(clf, \"last\")\n",
    "block_params = sum(p.numel() for p in clf.backbone.blocks[-1].parameters())\n",
    "assert n_last == 130 + block_params + 128, \"last = head + last block + final norm (128 parameters)\"\n",
    "n_all = set_trainable(clf, \"all\")\n",
    "total_no_lmhead = sum(p.numel() for p in clf.parameters()) - sum(p.numel() for p in clf.backbone.lm_head.parameters())\n",
    "assert n_all == total_no_lmhead, \"all = every parameter except the language-model head\"\n",
    "print(f\"self-check 2 passed: trainable parameters head {n_head:,} / last block {n_last:,} / all {n_all:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "## 5. TODO: metrics\n",
    "\n",
    "`confusion(y_true, y_pred)` returns a dict with the counts `tp`, `fp`, `fn`, `tn` (class 1 is the positive class).\n",
    "`scores(c)` takes that dict and returns `(accuracy, precision, recall, f1)`. Define precision, recall and F1 as 0 when their denominator is 0\n",
    "(a model that never predicts 1 has precision 0 here, by convention)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def confusion(y_true, y_pred):\n",
    "    tp = sum(1 for a, b in zip(y_true, y_pred) if a == 1 and b == 1)\n",
    "    fp = sum(1 for a, b in zip(y_true, y_pred) if a == 0 and b == 1)\n",
    "    fn = sum(1 for a, b in zip(y_true, y_pred) if a == 1 and b == 0)\n",
    "    tn = sum(1 for a, b in zip(y_true, y_pred) if a == 0 and b == 0)\n",
    "    return {\"tp\": tp, \"fp\": fp, \"fn\": fn, \"tn\": tn}\n",
    "\n",
    "def scores(c):\n",
    "    total = c[\"tp\"] + c[\"fp\"] + c[\"fn\"] + c[\"tn\"]\n",
    "    acc = (c[\"tp\"] + c[\"tn\"]) / total\n",
    "    prec = c[\"tp\"] / (c[\"tp\"] + c[\"fp\"]) if c[\"tp\"] + c[\"fp\"] else 0.0\n",
    "    rec = c[\"tp\"] / (c[\"tp\"] + c[\"fn\"]) if c[\"tp\"] + c[\"fn\"] else 0.0\n",
    "    f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0\n",
    "    return acc, prec, rec, f1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "### Self-check 3"
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "yt = [1, 1, 1, 0, 0, 0, 0, 0]\n",
    "yp = [1, 1, 0, 1, 0, 0, 0, 0]\n",
    "c = confusion(yt, yp)\n",
    "assert c == {\"tp\": 2, \"fp\": 1, \"fn\": 1, \"tn\": 4}, f\"confusion matrix is wrong: {c}\"\n",
    "acc, prec, rec, f1 = scores(c)\n",
    "assert abs(acc - 0.75) < 1e-9 and abs(prec - 2 / 3) < 1e-9 and abs(rec - 2 / 3) < 1e-9 and abs(f1 - 2 / 3) < 1e-9\n",
    "acc0, p0, r0, f0 = scores(confusion(yt, [0] * 8))\n",
    "assert abs(acc0 - 5 / 8) < 1e-9 and (p0, r0, f0) == (0.0, 0.0, 0.0), \"always answering 0: accuracy 5/8, but precision, recall and F1 are 0\"\n",
    "print(\"self-check 3 passed: metrics (and why accuracy alone misleads)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "## 6. Given: train the variants and compare\n",
    "\n",
    "`fit` trains with AdamW for 150 steps on the 200 labelled training examples, only on the parameters that have `requires_grad=True`. `evaluate` returns the four scores on the held-out examples.\n",
    "Read the table; the assertions only check what must hold."
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def fit(model, X, lengths, y, steps=150, lr=1e-3, seed=2):\n",
    "    g = torch.Generator().manual_seed(seed)\n",
    "    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)\n",
    "    for _ in range(steps):\n",
    "        ix = torch.randint(0, len(X), (32,), generator=g)\n",
    "        loss = F.cross_entropy(model(X[ix], lengths[ix]), y[ix])\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "\n",
    "def predict(model, X, lengths):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        p = model(X, lengths).argmax(-1).tolist()\n",
    "    model.train()\n",
    "    return p\n",
    "\n",
    "def evaluate(model):\n",
    "    return scores(confusion(y_va.tolist(), predict(model, X_va, len_va)))\n",
    "\n",
    "results = {}\n",
    "for name, mode, init, pool, lr in [(\"head only (pretrained)\", \"head\", \"pre\", \"last\", 3e-3), (\"last block + head\", \"last\", \"pre\", \"last\", 3e-3),\n",
    "                                   (\"all layers (pretrained)\", \"all\", \"pre\", \"last\", 1e-3), (\"all layers (from scratch)\", \"all\", \"scratch\", \"last\", 1e-3),\n",
    "                                   (\"all layers, FIRST token pooled\", \"all\", \"pre\", \"first\", 1e-3)]:\n",
    "    torch.manual_seed(0)\n",
    "    model = Classifier(copy.deepcopy(pretrained) if init == \"pre\" else Backbone(VOCAB + 1), pool=pool)\n",
    "    n_train = set_trainable(model, mode)\n",
    "    fit(model, X_tr, len_tr, y_tr, lr=lr)\n",
    "    results[name] = (n_train, *evaluate(model))\n",
    "print(f\"{'variant':34s} {'trainable':>9s} {'acc':>6s} {'prec':>6s} {'recall':>6s} {'F1':>6s}\")\n",
    "for name, (n, a, p, r, f) in results.items():\n",
    "    print(f\"{name:34s} {n:9,d} {a:6.3f} {p:6.3f} {r:6.3f} {f:6.3f}\")\n",
    "base_acc = 1 - y_va.float().mean().item()\n",
    "print(f\"{'always answer 0':34s} {'0':>9s} {base_acc:6.3f} {0:6.3f} {0:6.3f} {0:6.3f}\")\n",
    "\n",
    "assert results[\"all layers (pretrained)\"][4] >= 0.9, \"fine-tuning all layers should reach F1 >= 0.9 on this task\"\n",
    "assert results[\"all layers, FIRST token pooled\"][4] < 0.2, \"the first token has seen only itself: pooling it cannot solve the task\"\n",
    "assert results[\"head only (pretrained)\"][0] == 130\n",
    "print(\"self-check 4 passed: the variants behave as expected\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 7. TODO: look at the errors\n",
    "\n",
    "`find_errors(model, texts, X, lengths, y)` returns a list of `(text, true_label, predicted_label)` for every example the model gets wrong.\n",
    "Apply it to the \"last block + head\" variant, which is not perfect, and look at which sentences it misses."
   ]
  },
  {
   "cell_type": "code",
   "id": "c022",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def find_errors(model, texts_, X, lengths, y):\n",
    "    pred = predict(model, X, lengths)\n",
    "    return [(t, int(a), b) for t, a, b in zip(texts_, y.tolist(), pred) if a != b]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "### Self-check 5"
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "partial = Classifier(copy.deepcopy(pretrained)); set_trainable(partial, \"last\"); fit(partial, X_tr, len_tr, y_tr, lr=3e-3)\n",
    "errs = find_errors(partial, texts[700:], X_va, len_va, y_va)\n",
    "n_wrong = sum(a != b for a, b in zip(y_va.tolist(), predict(partial, X_va, len_va)))\n",
    "assert len(errs) == n_wrong and all(t != p for _, t, p in errs), \"every returned example must be a real error\"\n",
    "print(f\"{len(errs)} errors out of {len(y_va)}\")\n",
    "for t, true, pred in errs[:6]:\n",
    "    print(f\"  true {true} predicted {pred}: {t}\")\n",
    "fn = sum(1 for _, t, p in errs if t == 1); fp = len(errs) - fn\n",
    "print(f\"missed positives (false negatives): {fn}, false alarms (false positives): {fp}\")\n",
    "print(\"self-check 5 passed: error analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c025",
   "metadata": {},
   "source": [
    "## 8. Reflect\n",
    "\n",
    "1. Training only the head (130 parameters) did no better than answering 0 for everything, although its accuracy looks high. Why? What would the F1 score have told you that the accuracy did not?\n",
    "2. Why does pooling the first token fail while pooling the last real token works? What changes if you use a bidirectional model?\n",
    "3. Pretraining did not help much here. Which properties of this task explain that, and what would you expect on a real task with 200 labelled examples?\n",
    "4. The errors of the partly fine-tuned model: do they look random, or do they share a pattern? What would you do about it (more data of that kind, a different loss weighting, a different threshold)?\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
}