{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c001",
   "metadata": {},
   "source": [
    "# LM12-L1: Capstone milestone 1: dataset, tokenizer and pipeline\n",
    "\n",
    "The capstone builds a complete small LLM system in four milestones. The **default project** used in all notebooks is a *robot-command parser*:\n",
    "a tiny GPT that turns a short natural-language command (\"head to bay 17 quickly\") into a JSON tool call\n",
    "(`{\"tool\":\"move_to\",\"bay\":17,\"speed\":0.8}`) that the validators and safety layer of Module 11 can check.\n",
    "The data are **synthetic** (generated by code in this notebook): no personal data, no licence questions, fully reproducible.\n",
    "\n",
    "**This milestone:** generate the data, split it *without leakage* (including a held-out slice of unseen bay numbers to test generalisation),\n",
    "build a character tokenizer from the training split only, build padded batches with a loss mask, and write a dataset report.\n",
    "\n",
    "**Your own project:** you may replace `make_example` by your own data source (and keep every check). Read Lesson LM12-T1 first on scoping.\n",
    "\n",
    "**Hardware:** CPU only, about 2 GB of RAM, a few seconds. Needs PyTorch (CPU build).\n",
    "**How it works:** cells marked `TODO` raise `NotImplementedError` until you write the code; each is followed by a self-check (`assert`)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c002",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import json, math, random\n",
    "import torch\n",
    "\n",
    "torch.manual_seed(0)\n",
    "print(\"torch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c003",
   "metadata": {},
   "source": [
    "## 1. The data generator (given)\n",
    "\n",
    "1,000 distinct `move_to` commands plus 56 commands for the three rare tools. `bay` is stored with every example so the split can hold out whole bay numbers.\n",
    "Every input string occurs exactly once (no duplicates), so a text can never be in two splits by repetition alone."
   ]
  },
  {
   "cell_type": "code",
   "id": "c004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "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}      # never seen in training: a test of generalisation\n",
    "\n",
    "def other_commands():\n",
    "    \"\"\"The 56 commands for the three rare tools; each is used exactly once.\"\"\"\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",
    "    \"\"\"n_move distinct move commands (verb, bay 1..60, speed phrase) plus all 56 other commands, in random order.\"\"\"\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",
    "pairs = generate_pairs(1000, seed=0)\n",
    "print(len(pairs), \"examples; first:\", pairs[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c005",
   "metadata": {},
   "source": [
    "## 2. TODO: split without leakage\n",
    "\n",
    "`split_dataset(pairs, seed, holdout)` returns a dict with four lists:\n",
    "\n",
    "- `test_ood`: every example whose `bay` is in `holdout` (the \"out-of-distribution\" slice);\n",
    "- the remaining examples are shuffled with `random.Random(seed)` and cut **80% / 10% / 10%** into `train`, `val`, `test`\n",
    "  (`n_train = int(0.8 * n)`, `n_val_end = int(0.9 * n)`).\n",
    "\n",
    "No example with a held-out bay may appear in `train`, `val` or `test`, and no input text may occur in two splits."
   ]
  },
  {
   "cell_type": "code",
   "id": "c006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "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",
    "    rng = random.Random(seed)\n",
    "    rng.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}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c007",
   "metadata": {},
   "source": [
    "### Self-check 1: sizes, held-out bays, leakage and determinism"
   ]
  },
  {
   "cell_type": "code",
   "id": "c008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "sp = split_dataset(pairs, seed=0)\n",
    "print({k: len(v) for k, v in sp.items()})\n",
    "assert set(sp) == {\"train\", \"val\", \"test\", \"test_ood\"}\n",
    "assert sum(len(v) for v in sp.values()) == len(pairs), \"every example must be in exactly one split\"\n",
    "assert all(p[\"bay\"] in HOLDOUT_BAYS for p in sp[\"test_ood\"]) and len(sp[\"test_ood\"]) > 50\n",
    "for name in (\"train\", \"val\", \"test\"):\n",
    "    assert all(p[\"bay\"] not in HOLDOUT_BAYS for p in sp[name]), f\"held-out bays leaked into {name}\"\n",
    "texts = {k: {p[\"text\"] for p in v} for k, v in sp.items()}\n",
    "for a in sp:\n",
    "    for b in sp:\n",
    "        if a < b:\n",
    "            assert not (texts[a] & texts[b]), f\"input texts shared between {a} and {b}: leakage\"\n",
    "n_rest = len(sp[\"train\"]) + len(sp[\"val\"]) + len(sp[\"test\"])\n",
    "assert len(sp[\"train\"]) == int(0.8 * n_rest) and len(sp[\"val\"]) == int(0.9 * n_rest) - int(0.8 * n_rest)\n",
    "assert split_dataset(pairs, seed=0)[\"train\"] == sp[\"train\"], \"the split must be deterministic for a fixed seed\"\n",
    "assert split_dataset(pairs, seed=1)[\"train\"] != sp[\"train\"], \"another seed should give another split\"\n",
    "print(\"self-check 1 passed: no leakage, deterministic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c009",
   "metadata": {},
   "source": [
    "## 3. TODO: the tokenizer\n",
    "\n",
    "A **character** tokenizer is enough for this closed domain (Module 2 explains why real models use subwords).\n",
    "\n",
    "- `build_vocab(train)` returns `(itos, stoi)`. `itos[0] = PAD`, `itos[1] = END`, `itos[2] = UNK`, followed by the sorted set of characters that occur in the **training** split's\n",
    "  `text + \"\\n=>\" + target` strings. `stoi` is the inverse mapping. Build it from the training split *only*, so the validation and test sets stay unseen.\n",
    "- `encode(p, stoi, mask_prompt=True)` returns `(ids, labels)`. The sequence is `text + \"\\n=>\"` (the **prompt**), then `target`, then `END`. Characters missing from `stoi` map to `UNK`.\n",
    "  `labels` equals `ids` except that, when `mask_prompt` is true, the prompt positions are `-100` (ignored by the loss): the model is trained to produce the target, not to repeat the prompt."
   ]
  },
  {
   "cell_type": "code",
   "id": "c010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "PAD, END, UNK = \"\\x00\", \"\\x01\", \"\\x02\"\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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c011",
   "metadata": {},
   "source": [
    "### Self-check 2: vocabulary and encoding"
   ]
  },
  {
   "cell_type": "code",
   "id": "c012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "itos, stoi = build_vocab(sp[\"train\"])\n",
    "print(\"vocabulary size\", len(itos))\n",
    "assert itos[:3] == [PAD, END, UNK] and 40 < len(itos) < 60\n",
    "assert all(stoi[c] == i for i, c in enumerate(itos))\n",
    "assert \"~\" not in stoi, \"the vocabulary must contain only characters that occur in the training split\"\n",
    "val_chars = {c for p in sp[\"val\"] + sp[\"test\"] + sp[\"test_ood\"] for c in p[\"text\"] + \"\\n=>\" + p[\"target\"]}\n",
    "assert val_chars <= set(itos), f\"characters outside the vocabulary: {val_chars - set(itos)}\"\n",
    "\n",
    "ex = {\"text\": \"stop\", \"target\": '{\"tool\":\"emergency_stop\"}', \"bay\": None}\n",
    "ids, lab = encode(ex, stoi)\n",
    "assert len(ids) == len(\"stop\\n=>\") + len(ex[\"target\"]) + 1 == len(lab)\n",
    "assert lab[: len(\"stop\\n=>\")] == [-100] * len(\"stop\\n=>\"), \"the prompt must be masked\"\n",
    "assert lab[len(\"stop\\n=>\"):] == ids[len(\"stop\\n=>\"):] and ids[-1] == stoi[END]\n",
    "ids2, lab2 = encode(ex, stoi, mask_prompt=False)\n",
    "assert ids2 == ids and lab2 == ids, \"without masking the labels equal the ids\"\n",
    "unk_ids, _ = encode({\"text\": \"stop~\", \"target\": \"{}\", \"bay\": None}, stoi)\n",
    "assert unk_ids[4] == stoi[UNK], \"an unseen character maps to UNK\"\n",
    "print(\"self-check 2 passed: tokenizer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c013",
   "metadata": {},
   "source": [
    "## 4. TODO: padded batches\n",
    "\n",
    "`make_batch(pairs, stoi, max_len, mask_prompt=True)` returns tensors `X` and `Y` of shape `(len(pairs), max_len)`, dtype `torch.long`.\n",
    "`X` holds the ids padded with `stoi[PAD]` (= 0); `Y` holds the labels padded with `-100`. Sequences must fit: `max_len` is the longest encoded sequence in the data."
   ]
  },
  {
   "cell_type": "code",
   "id": "c014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c015",
   "metadata": {},
   "source": [
    "### Self-check 3: shapes, padding and the number of supervised tokens"
   ]
  },
  {
   "cell_type": "code",
   "id": "c016",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "max_len = max(len(encode(p, stoi)[0]) for p in pairs)\n",
    "X, Y = make_batch(sp[\"train\"][:16], stoi, max_len)\n",
    "print(\"longest sequence\", max_len, \"tokens; batch\", tuple(X.shape))\n",
    "assert X.shape == Y.shape == (16, max_len) and X.dtype == Y.dtype == torch.long\n",
    "for i, p in enumerate(sp[\"train\"][:16]):\n",
    "    assert int((Y[i] != -100).sum()) == len(p[\"target\"]) + 1, \"supervised tokens = target characters + END\"\n",
    "    n = len(p[\"text\"]) + 3 + len(p[\"target\"]) + 1\n",
    "    assert torch.all(X[i, n:] == 0) and torch.all(Y[i, n:] == -100), \"padding\"\n",
    "Xu, Yu = make_batch(sp[\"train\"][:4], stoi, max_len, mask_prompt=False)\n",
    "assert int((Yu[0] != -100).sum()) == len(sp[\"train\"][0][\"text\"]) + 3 + len(sp[\"train\"][0][\"target\"]) + 1\n",
    "print(\"self-check 3 passed: batches\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c017",
   "metadata": {},
   "source": [
    "## 5. TODO: a dataset report\n",
    "\n",
    "Every report needs the facts about the data. `dataset_report(sp, stoi)` returns a dict with:\n",
    "\n",
    "- `sizes`: number of examples per split;\n",
    "- `tools`: counts of each `tool` name in the training split (parse the `target` JSON);\n",
    "- `max_len`, `p50_len`, `p95_len`: the maximum, median and 95th percentile of the encoded sequence length over **all** splits\n",
    "  (sort the lengths; percentile `q` is the element at index `ceil(q * n) - 1`)."
   ]
  },
  {
   "cell_type": "code",
   "id": "c018",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def dataset_report(sp, stoi):\n",
    "    tools = {}\n",
    "    for p in sp[\"train\"]:\n",
    "        t = json.loads(p[\"target\"])[\"tool\"]\n",
    "        tools[t] = tools.get(t, 0) + 1\n",
    "    lengths = sorted(len(encode(p, stoi)[0]) for v in sp.values() for p in v)\n",
    "    n = len(lengths)\n",
    "    pct = lambda q: lengths[math.ceil(q * n) - 1]\n",
    "    return {\"sizes\": {k: len(v) for k, v in sp.items()}, \"tools\": tools,\n",
    "            \"max_len\": lengths[-1], \"p50_len\": pct(0.5), \"p95_len\": pct(0.95)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c019",
   "metadata": {},
   "source": [
    "### Self-check 4: the report"
   ]
  },
  {
   "cell_type": "code",
   "id": "c020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "rep = dataset_report(sp, stoi)\n",
    "print(json.dumps(rep, indent=1))\n",
    "assert rep[\"sizes\"] == {k: len(v) for k, v in sp.items()}\n",
    "assert set(rep[\"tools\"]) == {\"move_to\", \"emergency_stop\", \"get_battery\", \"say\"} and sum(rep[\"tools\"].values()) == len(sp[\"train\"])\n",
    "assert rep[\"max_len\"] == max_len and rep[\"p50_len\"] <= rep[\"p95_len\"] <= rep[\"max_len\"]\n",
    "assert rep[\"tools\"][\"move_to\"] > 10 * rep[\"tools\"][\"say\"], \"the data are dominated by move_to commands: say so in your report\"\n",
    "print(\"self-check 4 passed: dataset report\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c021",
   "metadata": {},
   "source": [
    "## 6. Reflect and write the dataset card\n",
    "\n",
    "Copy these answers into your project notes; they become the *Data* section of the report (Lesson LM12-T4).\n",
    "\n",
    "1. **Source and licence:** where do the data come from and who may use them? (Here: generated by this notebook, synthetic, no personal data.)\n",
    "2. **What is in them and what is missing?** Which commands, phrasings, languages and typos does the generator never produce? What will the model therefore fail on?\n",
    "3. **Splits:** how were they made, and what does each measure? (`test` = same distribution; `test_ood` = bay numbers never seen in training.)\n",
    "4. **Class balance:** `move_to` dominates (the rare tools have only a handful of examples in each evaluation split); why can overall accuracy hide failures on the rare tools?\n",
    "5. **What would make a leak?** Name two ways in which test information could reach the training set in your own project.\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
}