1865 lines
69 KiB
Plaintext
1865 lines
69 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "ef460528",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Let's build GPT — from scratch\n",
|
||
"\n",
|
||
"Tiny GPT on *tiny Shakespeare* — **concept -> code -> Your turn** each step.\n",
|
||
"\n",
|
||
"**Red thread:** every section has **-> Training:** plus a **goal stack**\n",
|
||
"so you always know *why* this step exists and what it unlocks next.\n",
|
||
"\n",
|
||
"**Kernel:** `gpt_from_scratch/.venv`. Run cells top to bottom.\n",
|
||
"\n",
|
||
"**Ultimate goal:** weight matrices that score the **true** next character\n",
|
||
"highly on Shakespeare — so generation looks like language, not noise.\n",
|
||
"\n",
|
||
"**Read the chain bottom-up.** Each line needs the line below it first:\n",
|
||
"```\n",
|
||
"GOAL better weights -> better text generation\n",
|
||
" ^ step 4 optimizer.step() uses gradients to nudge weights\n",
|
||
" ^ step 3 loss.backward() computes gradient per weight\n",
|
||
" ^ step 2 loss = cross_entropy(...) one number: how wrong we were\n",
|
||
" ^ step 1 forward: model(xb, yb) logits compared to answer key yb\n",
|
||
" ^ step 0 get_batch -> xb, yb one practice batch from train_data\n",
|
||
" ^ ... tokenize, tensor tape everything in sections 0-5\n",
|
||
"```\n",
|
||
"\n",
|
||
"| Sec | You build | So that later... |\n",
|
||
"|-----|-----------|------------------|\n",
|
||
"| 0 | PyTorch + GPU | tensors & weights can run fast |\n",
|
||
"| 1-2 | text -> ids | model eats numbers |\n",
|
||
"| 3 | tensor tape + split | sample train/val batches |\n",
|
||
"| 4-5 | x/y windows, batches | one `forward` has questions + answers |\n",
|
||
"| 6 | logits, loss, weights | `loss.backward()` has something to fix |\n",
|
||
"| 7 | generate (untrained) | baseline before learning |\n",
|
||
"| 8 | forward/backward/step | **weights actually change** |\n",
|
||
"| 9 | generate (trained) | see the goal — better patterns |\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "33e0cea0",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 0. Setup (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 0):** install the math engine (PyTorch + GPU).\n",
|
||
"\n",
|
||
"**Needs:** nothing in this notebook yet.\n",
|
||
"**Unlocks:** tensors on `device` (section 3 batches, section 6 weights all live here).\n",
|
||
"We use **PyTorch**: tensors (arrays of numbers) + `torch.nn` (neural nets).\n",
|
||
"Training is heavy math on millions of numbers — a **GPU** speeds that up 10–100×.\n",
|
||
"\n",
|
||
"**→ Training:** Learning = huge amounts of matrix math. Setup ensures PyTorch, tensors, and the GPU are ready to run that math millions of times.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "91321dde",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 0.1 Imports\n",
|
||
"\n",
|
||
"These three lines appear in almost every PyTorch project.\n",
|
||
"\n",
|
||
"**→ Training:** `torch` runs the math, `nn` holds the **weight matrices** we train, `F` computes **loss** (how wrong we are) from predictions.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d99d4439",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import torch\n",
|
||
"import torch.nn as nn\n",
|
||
"from torch.nn import functional as F\n",
|
||
"print(\"PyTorch version:\", torch.__version__)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "57084056",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 0.2 CPU vs GPU (`device`)\n",
|
||
"\n",
|
||
"Tensors can live on **CPU** (always works) or **CUDA/GPU** (fast). We pick once\n",
|
||
"and reuse `device` so every tensor and model lands in the same place.\n",
|
||
"\n",
|
||
"If `CUDA available` is False, you installed CPU-only PyTorch — training still\n",
|
||
"works but will be slow.\n",
|
||
"\n",
|
||
"**→ Training:** Model **weights** and each **batch** must sit on the same device. Training on GPU is the same logic, just far faster.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c3f68853",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"print(\"CUDA available:\", torch.cuda.is_available())\n",
|
||
"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
|
||
"print(\"using device:\", device)\n",
|
||
"if device == \"cuda\":\n",
|
||
" print(\"GPU:\", torch.cuda.get_device_name(0))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "6126b5a6",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 0.3 What is a tensor?\n",
|
||
"\n",
|
||
"A **tensor** is an n-dimensional array. `shape (5,)` = 5 numbers in a row.\n",
|
||
"`shape (2, 3)` = 2 rows × 3 columns. Token ids will be 1-D tensors.\n",
|
||
"\n",
|
||
"**→ Training:** Training data (`xb`, `yb`), model **weights**, and **gradients** are all tensors. Slicing `train_data` builds each practice batch.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "2b37eccd",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"row = torch.tensor([10, 20, 30])\n",
|
||
"grid = torch.tensor([[1, 2], [3, 4]])\n",
|
||
"print(\"row shape:\", row.shape, \" values:\", row.tolist())\n",
|
||
"print(\"grid shape:\", grid.shape)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8a1bb533",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 0** — Create a 2×2 tensor of zeros on `device`.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "25810c05",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"practice = torch.zeros(2, 2, device=device)\n",
|
||
"print(practice)\n",
|
||
"print(\"on device:\", practice.device.type == device)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bad266c2",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 1. The data (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 1):** load the raw training material.\n",
|
||
"\n",
|
||
"**Needs:** section 0 (imports).\n",
|
||
"**Unlocks:** `text` -> tokenization (section 2) -> the pool of next-token questions.\n",
|
||
"A **language model** learns patterns in text by repeatedly asking: *given what\n",
|
||
"came before, what comes next?* No separate label file — the text **is** the labels.\n",
|
||
"\n",
|
||
"We use ~1M characters of Shakespeare in `input.txt`.\n",
|
||
"\n",
|
||
"**→ Training:** The text is our only ingredient. Each training step asks: \"given chars before, predict the next\" — labels are **inside** the text.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f7039a6a",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 1.1 Load the file\n",
|
||
"\n",
|
||
"Plain Python: read the whole file into one Python string `text`.\n",
|
||
"\n",
|
||
"**→ Training:** `text` is the raw pool of examples. Nothing trains until we encode it and sample windows from it.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "9e4fa669",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
|
||
" text = f.read()\n",
|
||
"print(\"characters in file:\", len(text))\n",
|
||
"print(\"---- first 200 chars ----\")\n",
|
||
"print(text[:200])\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "83f3fde7",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 1.2 What's inside?\n",
|
||
"\n",
|
||
"The model sees **every** character: letters, spaces, punctuation, newlines.\n",
|
||
"Newlines matter — they mark line breaks in the plays.\n",
|
||
"\n",
|
||
"**→ Training:** The model learns the **statistics** of these characters (spaces, newlines, letter pairs). It does not store the file verbatim in weights.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "31f45b4a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"print(\"newlines:\", text.count(\"\\n\"))\n",
|
||
"print(\"unique character count:\", len(set(text)))\n",
|
||
"print(\"first line:\", repr(text.split(\"\\n\")[0][:60]))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9bfde4e7",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 1** — Print the last 100 characters of `text`.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "dc5dfe07",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"print(repr(text[-100:]))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f9a73d64",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 2. Tokenization (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 2):** text -> integer token ids.\n",
|
||
"\n",
|
||
"**Needs:** section 1 (`text`).\n",
|
||
"**Unlocks:** `encode` / `vocab_size` -> tensor tape (section 3) -> model input/output size (section 6).\n",
|
||
"Neural nets need **integers**. **Tokenization** = chop text into **tokens** and\n",
|
||
"assign each an id.\n",
|
||
"\n",
|
||
"| Approach | Token | Used by |\n",
|
||
"|----------|-------|---------|\n",
|
||
"| Character | one letter | us (simple) |\n",
|
||
"| Sub-word (BPE) | pieces like \"ing\" | GPT, Llama |\n",
|
||
"\n",
|
||
"We use **one char = one token** — tiny vocab, longer sequences.\n",
|
||
"\n",
|
||
"**→ Training:** The network only sees integers. Tokenization turns Shakespeare into the id sequences that every forward pass and loss calculation uses.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "78bd31aa",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 2.1 Vocabulary\n",
|
||
"\n",
|
||
"The **vocabulary** = all distinct characters in the training text. `vocab_size`\n",
|
||
"is how many different tokens the model can output.\n",
|
||
"\n",
|
||
"**→ Training:** `vocab_size` fixes the width of the output scores. Each training step: model outputs `vocab_size` logits; truth is **one** of those ids.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "07af82cf",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"chars = sorted(list(set(text)))\n",
|
||
"vocab_size = len(chars)\n",
|
||
"print(\"vocab_size:\", vocab_size)\n",
|
||
"print(\"all chars:\", repr(\"\".join(chars)))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9f03ff92",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 2.2 Lookup tables: `stoi` and `itos`\n",
|
||
"\n",
|
||
"- `stoi` (**s**tring **to** **i**nt): character → id\n",
|
||
"- `itos` (**i**nt **to** **s**tring): id → character\n",
|
||
"\n",
|
||
"**→ Training:** `stoi` converts batches of text to `xb`/`yb` tensors. `itos` is for us humans when inspecting generated output.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "3e0b4759",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"stoi = {ch: i for i, ch in enumerate(chars)}\n",
|
||
"itos = {i: ch for i, ch in enumerate(chars)}\n",
|
||
"print('stoi[h] =', stoi['h'])\n",
|
||
"print('itos[0] =', repr(itos[0]))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8f5a3085",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 2.3 `encode` and `decode`\n",
|
||
"\n",
|
||
"Encode: string → list of ids. Decode: list of ids → string. Round-trip should match.\n",
|
||
"\n",
|
||
"**→ Training:** `encode(text)` built the master tape `data`. Every training batch is a slice of those ids — never raw strings inside the model.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "14c4e3a3",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"encode = lambda s: [stoi[c] for c in s]\n",
|
||
"decode = lambda l: \"\".join(itos[i] for i in l)\n",
|
||
"demo = \"hello there\"\n",
|
||
"ids = encode(demo)\n",
|
||
"print(demo, \"->\", ids)\n",
|
||
"print(\"back:\", decode(ids))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "40dd1ac1",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 2** — Encode a short phrase. Flip one id and decode — see one wrong letter.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0b50c865",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"phrase = \"First Citizen\"\n",
|
||
"ids = encode(phrase)\n",
|
||
"ids[3] = (ids[3] + 1) % vocab_size\n",
|
||
"print(\"corrupted:\", decode(ids))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f90ea698",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 3. Tensor tape + train/val split (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 3):** whole book as one integer tape + train/val split.\n",
|
||
"\n",
|
||
"**Needs:** section 2 (`encode`, `vocab_size`).\n",
|
||
"**Unlocks:** `train_data` to sample batches (section 5); `val_data` to check generalisation (section 8).\n",
|
||
"Section 2 left us with Python lists of ids. Here we wrap the whole book in a\n",
|
||
"**tensor** (3.1), then split it into train and validation (3.2–3.4).\n",
|
||
"\n",
|
||
"**→ Training:** `train_data` is what we sample from when updating weights. `val_data` checks whether those weights work on **held-out** text.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "5a1caac1",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 3.1 What is a tensor?\n",
|
||
"\n",
|
||
"In section 2 we had a Python **list** of token ids: `[18, 47, 56, ...]`.\n",
|
||
"PyTorch uses a **tensor** for the same idea — a block of numbers the library\n",
|
||
"can slice, stack, and run on a GPU very fast.\n",
|
||
"\n",
|
||
"**Think of it like this:** a list is a handwritten shopping list. A tensor is\n",
|
||
"the same list typed into a spreadsheet — same items, but the computer can do\n",
|
||
"math on the whole column at once.\n",
|
||
"\n",
|
||
"#### How many dimensions? (`shape`)\n",
|
||
"\n",
|
||
"| Example | shape | Meaning |\n",
|
||
"|---------|-------|---------|\n",
|
||
"| `torch.tensor(5)` | `()` | one number (scalar) |\n",
|
||
"| `torch.tensor([1,2,3])` | `(3,)` | one row of 3 ids — **our book tape** |\n",
|
||
"| `torch.tensor([[1,2],[3,4]])` | `(2, 2)` | 2 rows × 2 cols — later `(B, T)` batches |\n",
|
||
"\n",
|
||
"The comma in `(3,)` means \"one dimension of length 3\" — not a tuple pair.\n",
|
||
"\n",
|
||
"#### `dtype` — what kind of number?\n",
|
||
"\n",
|
||
"- `torch.long` — **integers** (token ids). No decimals.\n",
|
||
"- `torch.float32` — **decimals** (model weights, logits, loss).\n",
|
||
"\n",
|
||
"Token ids are always integers, so we use `dtype=torch.long`.\n",
|
||
"\n",
|
||
"#### Why not stay with Python lists?\n",
|
||
"\n",
|
||
"1. **Speed** — millions of multiply/adds per second on GPU.\n",
|
||
"2. **Slicing** — `data[1000:1008]` gives another tensor, no copy of the whole book.\n",
|
||
"3. **Gradients** — later, PyTorch tracks how each weight affects loss (training).\n",
|
||
"\n",
|
||
"Run the cell below: same ids as a list vs as a tensor.\n",
|
||
"\n",
|
||
"**→ Training:** We store the whole book as one tensor so we can grab random windows fast and (later) move batches to the GPU in one shot.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "21051ca8",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"ids_list = encode(\"First\")\n",
|
||
"ids_tensor = torch.tensor(ids_list, dtype=torch.long)\n",
|
||
"\n",
|
||
"print(\"Python list:\", ids_list)\n",
|
||
"print(\"tensor: \", ids_tensor)\n",
|
||
"print(\"shape:\", ids_tensor.shape)\n",
|
||
"print(\"dtype:\", ids_tensor.dtype)\n",
|
||
"print()\n",
|
||
"print(\"slice [1:3]:\", ids_tensor[1:3])\n",
|
||
"print(\"decode slice:\", decode(ids_tensor[1:3].tolist()))\n",
|
||
"print(\"back to list:\", ids_tensor.tolist())\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7b498e68",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 3.1** — Encode a word of your choice. Build a tensor, print `.shape`\n",
|
||
"and `.dtype`. Slice the middle two ids and decode them.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c40b416b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"word = \"Citizen\" # edit me\n",
|
||
"t = torch.tensor(encode(word), dtype=torch.long)\n",
|
||
"print(\"shape:\", t.shape, \" dtype:\", t.dtype)\n",
|
||
"mid = t[1:3]\n",
|
||
"print(\"middle ids:\", mid.tolist(), \"->\", decode(mid.tolist()))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "d1c84cec",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 3.2 String → tensor (the whole book)\n",
|
||
"\n",
|
||
"Now we wrap **every** token id in the dataset into one 1-D tensor `data`.\n",
|
||
"Same idea as 3.1, but ~1 million ids long — our **tape**.\n",
|
||
"\n",
|
||
"`dtype=torch.long` = integers only. Shape `(N,)` = one long row of length `N`.\n",
|
||
"\n",
|
||
"**→ Training:** `data` / `train_data` never change during training. Only the **model weights** change. The tape is fixed questions; weights are learned answers.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d419690c",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"data = torch.tensor(encode(text), dtype=torch.long)\n",
|
||
"print(\"shape:\", data.shape, \" dtype:\", data.dtype)\n",
|
||
"print(\"first 10 ids:\", data[:10].tolist())\n",
|
||
"print(\"first 10 chars:\", decode(data[:10].tolist()))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "5450cf66",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 3.3 Train / validation split\n",
|
||
"\n",
|
||
"- **train_data** — model learns from this\n",
|
||
"- **val_data** — later we check loss on unseen tail of the book\n",
|
||
"\n",
|
||
"If train loss falls but val loss stays high → **memorising**, not generalising.\n",
|
||
"\n",
|
||
"**→ Training:** We **only** run `backward()` on train batches. Val loss is read-only: if train improves but val does not, weights are memorising.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "bbc77a95",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"n = int(0.9 * len(data))\n",
|
||
"train_data = data[:n]\n",
|
||
"val_data = data[n:]\n",
|
||
"print(\"train:\", len(train_data), \" val:\", len(val_data))\n",
|
||
"print(\"sum check:\", len(train_data) + len(val_data) == len(data))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "dcf68150",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 3.4 The split boundary in plain text\n",
|
||
"\n",
|
||
"Train ends and val begins at the same place in the original book — no gap, no overlap.\n",
|
||
"\n",
|
||
"**→ Training:** Confirms train and val are contiguous slices — so val loss is a fair \"unseen tail of the book\" check.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e9bc8a68",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"print(\"...end of train:\", repr(decode(train_data[-40:].tolist())))\n",
|
||
"print(\"start of val: \", repr(decode(val_data[:40].tolist())))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "51b55dd8",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 3** — What fraction of the book is validation? *(~10%.)*\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "70e5f2ac",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"frac_val = len(val_data) / len(data)\n",
|
||
"print(f\"validation fraction: {frac_val:.1%}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "fe71c7d8",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 4. Predicting the next token (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 4):** define ONE training question (context -> next token).\n",
|
||
"\n",
|
||
"**Needs:** section 3 (`train_data` tape).\n",
|
||
"**Unlocks:** `x` / `y` / `block_size` -> batched training examples (section 5) -> loss targets (section 6).\n",
|
||
"A language model is a **next-token guesser**. You show it some text; it scores\n",
|
||
"every possible next token; you compare to the real answer and nudge its weights.\n",
|
||
"\n",
|
||
"Section 4 breaks that idea into small pieces. **Read one part, run the code\n",
|
||
"right below it**, then move on. By the end you'll understand chunks, `block_size`,\n",
|
||
"`x`/`y`, and the prefix trick — the foundation for batching in section 5.\n",
|
||
"\n",
|
||
"**→ Training:** The **objective** for the whole notebook: minimise surprise at the true next token. Everything in §4 defines one training example.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "e7a8f7c0",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.1 The next-token game (one step)\n",
|
||
"\n",
|
||
"The training question is always the same:\n",
|
||
"\n",
|
||
"> Given what came before → what is the **single** next character?\n",
|
||
"\n",
|
||
"There is no separate labels file. The label is hiding **in the text itself**:\n",
|
||
"the next character in the sequence.\n",
|
||
"\n",
|
||
"**Analogy:** a friend starts a sentence and you blurt out the next word. They\n",
|
||
"tell you if you were right. That's one training step.\n",
|
||
"\n",
|
||
"**→ Training:** One row of the loss compares model scores to **one** correct next id. A batch contains `B × T` such comparisons at once.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6adf5e2f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# tiny phrase from our vocabulary — watch one step at a time\n",
|
||
"phrase = \"First Citizen\"\n",
|
||
"ids = encode(phrase)\n",
|
||
"print(\"text: \", phrase)\n",
|
||
"print(\"ids: \", ids)\n",
|
||
"print()\n",
|
||
"for i in range(len(ids) - 1):\n",
|
||
" context = decode(ids[: i + 1])\n",
|
||
" nxt = itos[ids[i + 1]]\n",
|
||
" print(f\" seen {context!r:20s} -> predict {nxt!r}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "537d244e",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 4.1** — Change `phrase` to another short string (letters/spaces only).\n",
|
||
"How many training steps does a string of length *L* give? *(Answer: L − 1.)*\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "49b69d41",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"phrase = \"To be\" # edit me\n",
|
||
"ids = encode(phrase)\n",
|
||
"steps = len(ids) - 1\n",
|
||
"print(f\"length {len(ids)} -> {steps} next-token questions\")\n",
|
||
"for i in range(steps):\n",
|
||
" print(f\" {decode(ids[:i+1])!r} -> {itos[ids[i+1]]!r}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a5039f5f",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.2 The book is one long tape\n",
|
||
"\n",
|
||
"After section 3, the whole dataset is a single 1-D tensor `train_data` — a tape\n",
|
||
"of token ids with no breaks between chapters:\n",
|
||
"\n",
|
||
"```\n",
|
||
"train_data: [18, 47, 56, 57, ...] ← millions of ids, one after another\n",
|
||
"```\n",
|
||
"\n",
|
||
"Training never asks \"predict the whole book.\" It only ever asks about **local**\n",
|
||
"stretches of that tape.\n",
|
||
"\n",
|
||
"**→ Training:** Every example is a contiguous window on this tape. Random `start` indices = seeing different parts of the book each step.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d217e75d",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"print(\"train_data length:\", len(train_data))\n",
|
||
"print(\"first 15 ids: \", train_data[:15].tolist())\n",
|
||
"print(\"first 15 chars: \", decode(train_data[:15].tolist()))\n",
|
||
"print(\"chars 15-30: \", decode(train_data[15:30].tolist()))\n",
|
||
"print()\n",
|
||
"print(\"The tape is continuous — char 15 follows char 14 in the real book.\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0ecf904b",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.3 Why we use chunks (`block_size`)\n",
|
||
"\n",
|
||
"We could slide one character at a time, but in practice we grab a **window** of\n",
|
||
"`block_size` characters at once. Three reasons:\n",
|
||
"\n",
|
||
"1. **Memory** — models store state per position; a 1M-wide window would not fit on a GPU.\n",
|
||
"2. **Fixed context** — Transformers are built to look at at most `block_size` tokens.\n",
|
||
"3. **Efficiency** — one forward pass on 8 positions trains 8 questions in parallel (next part).\n",
|
||
"\n",
|
||
"`block_size` is the **context window**: how far back the model is *allowed* to look.\n",
|
||
"We use 8 here so the numbers fit on screen; GPT-4 uses thousands.\n",
|
||
"\n",
|
||
"**→ Training:** `block_size` = how many next-token questions per row (`T`). Also caps context length the architecture can use.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "7b330af9",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"block_size = 8\n",
|
||
"start = 1000 # where we cut the tape\n",
|
||
"chunk = train_data[start : start + block_size]\n",
|
||
"\n",
|
||
"print(f\"block_size = {block_size}\")\n",
|
||
"print(f\"slice from index {start} to {start + block_size - 1}\")\n",
|
||
"print(\"chunk ids: \", chunk.tolist())\n",
|
||
"print(\"chunk text: \", repr(decode(chunk.tolist())))\n",
|
||
"print()\n",
|
||
"print(\"book length:\", len(train_data), \" window:\", block_size,\n",
|
||
" \" ratio:\", round(len(train_data) / block_size))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f19d6e14",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 4.3** — Change `start`. The chunk should always be exactly `block_size`\n",
|
||
"characters (except if you pick a start too close to the end).\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8528eb17",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"start = 25000 # try different values\n",
|
||
"chunk = train_data[start : start + block_size]\n",
|
||
"print(\"len(chunk):\", len(chunk), \" text:\", repr(decode(chunk.tolist())))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "342c105f",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.4 Inputs `x` and targets `y` (shift by one)\n",
|
||
"\n",
|
||
"From one slice we build two equal-length tensors:\n",
|
||
"\n",
|
||
"- **`x`** = the window `train_data[start : start + block_size]`\n",
|
||
"- **`y`** = the same window, shifted **one step forward** in time:\n",
|
||
" `train_data[start + 1 : start + block_size + 1]`\n",
|
||
"\n",
|
||
"So `y` is always \"what comes next\" relative to `x`, position by position:\n",
|
||
"\n",
|
||
"```\n",
|
||
"x: h e l l o _ t h\n",
|
||
"y: e l l o _ t h ? ← ? is the 9th char (first char AFTER the window)\n",
|
||
"```\n",
|
||
"\n",
|
||
"The model never sees `y` during the forward pass — only during loss calculation.\n",
|
||
"\n",
|
||
"#### Tie-in to the training goal\n",
|
||
"\n",
|
||
"**Ultimate goal:** nudge **weights** so the model scores the true next character highly.\n",
|
||
"\n",
|
||
"`x` and `y` are how we package **one practice round** for that goal:\n",
|
||
"\n",
|
||
"| Tensor | Role in training | Analogy |\n",
|
||
"|--------|------------------|---------|\n",
|
||
"| **`x`** (later `xb`) | **Question** — tokens the model reads to produce logits | exam paper |\n",
|
||
"| **`y`** (later `yb`) | **Answer key** — the correct next token at each position | mark scheme |\n",
|
||
"\n",
|
||
"At column `t`: model reads `x[t]`, predicts a distribution; loss compares it to **`y[t]`** (the char that really followed on the tape). Wrong guess → higher loss → `backward()` pushes weights to do better next time.\n",
|
||
"\n",
|
||
"**In the chain (why this section exists):**\n",
|
||
"```\n",
|
||
"GOAL: better weights\n",
|
||
" ^ needs backward()\n",
|
||
" ^ needs loss (one number)\n",
|
||
" ^ needs yb — the CORRECT next tokens ← y is the template for yb\n",
|
||
" ^ needs forward on xb — model predictions ← x is the template for xb\n",
|
||
"```\n",
|
||
"\n",
|
||
"Without the shift-by-one split you have text but **no labelled examples** — nothing to score, no loss, no training.\n",
|
||
"\n",
|
||
"**→ Training:** `get_batch` builds matrices `xb`/`yb` the same way as `x`/`y` here. Every training step is: predict from `xb`, grade against `yb`, update weights.\n",
|
||
"\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0dec560b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"start = 1000\n",
|
||
"x = train_data[start : start + block_size]\n",
|
||
"y = train_data[start + 1 : start + block_size + 1]\n",
|
||
"\n",
|
||
"print(\"x:\", decode(x.tolist()))\n",
|
||
"print(\"y:\", decode(y.tolist()))\n",
|
||
"print()\n",
|
||
"for i in range(block_size):\n",
|
||
" print(f\" x[{i}]={repr(itos[x[i].item()])} y[{i}]={repr(itos[y[i].item()])}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "70613e82",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 4.4** — Confirm `y[i]` is always the id **right after** `x[i]` on the tape.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "5be4764d",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"ok = all(y[t].item() == train_data[start + t + 1].item() for t in range(block_size))\n",
|
||
"print(\"every y[t] is the next char after x[t] on the tape:\", ok)\n",
|
||
"if not ok:\n",
|
||
" for t in range(block_size):\n",
|
||
" print(t, x[t].item(), y[t].item(), train_data[start + t + 1].item())\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "6fecc0fb",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.5 The prefix trick (one chunk → many lessons)\n",
|
||
"\n",
|
||
"Here is the key idea that confuses many beginners:\n",
|
||
"\n",
|
||
"An 8-character chunk is **not** one training example. It is **eight** examples\n",
|
||
"packed together — one for each prefix length.\n",
|
||
"\n",
|
||
"| step `t` | context (prefix of `x`) | must predict `y[t]` |\n",
|
||
"|----------|---------------------------|---------------------|\n",
|
||
"| 0 | 1st char only | 2nd char |\n",
|
||
"| 1 | 1st + 2nd | 3rd char |\n",
|
||
"| 2 | first 3 chars | 4th char |\n",
|
||
"| … | … | … |\n",
|
||
"| 7 | all 8 chars | 9th char (just past the window) |\n",
|
||
"\n",
|
||
"**Analogy:** teaching spelling by asking \"after **h**?\", then \"after **he**?\",\n",
|
||
"then \"after **hel**?\" — same word, harder each time.\n",
|
||
"\n",
|
||
"**→ Training:** One slice yields `block_size` supervised examples per forward pass — more learning signal per batch.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e4559d0f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"start = 1000\n",
|
||
"x = train_data[start : start + block_size]\n",
|
||
"y = train_data[start + 1 : start + block_size + 1]\n",
|
||
"\n",
|
||
"print(\"full chunk:\", repr(decode(x.tolist())))\n",
|
||
"print()\n",
|
||
"for t in range(block_size):\n",
|
||
" prefix = decode(x[: t + 1].tolist())\n",
|
||
" target = itos[y[t].item()]\n",
|
||
" print(f\"t={t} prefix {prefix!r:12s} -> {target!r}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "3be6fbb3",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 4.5** — At `t=0`, how many characters does the model see? At `t=7`?\n",
|
||
"*(0: one char; 7: all eight.)*\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "898927dc",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"for t in [0, 3, 7]:\n",
|
||
" prefix_len = t + 1\n",
|
||
" print(f\"t={t}: prefix length = {prefix_len}, target = {repr(itos[y[t].item()])}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "36cde17a",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.6 Column alignment (what `x[t]` and `y[t]` mean together)\n",
|
||
"\n",
|
||
"When we later batch data into a matrix, **column `t`** means:\n",
|
||
"\n",
|
||
"- input token at that step: `x[t]`\n",
|
||
"- correct next token: `y[t]`\n",
|
||
"\n",
|
||
"Important: `y[t]` is the character that follows **`x[t]`** in the original book,\n",
|
||
"not \"the char after the whole prefix.\" (A Transformer uses the full prefix via\n",
|
||
"attention; our Bigram model will only look at `x[t]` — more in section 6.)\n",
|
||
"\n",
|
||
"**→ Training:** At column `t`, loss asks: \"after token `xb[b,t]`, was your distribution close to `yb[b,t]`?\"\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "876ad4ce",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"start = 1000\n",
|
||
"x = train_data[start : start + block_size]\n",
|
||
"y = train_data[start + 1 : start + block_size + 1]\n",
|
||
"\n",
|
||
"for t in range(block_size):\n",
|
||
" on_tape = train_data[start + t].item()\n",
|
||
" on_tape_next = train_data[start + t + 1].item()\n",
|
||
" match = x[t].item() == on_tape and y[t].item() == on_tape_next\n",
|
||
" print(f\"t={t} x={repr(itos[x[t].item()])} y={repr(itos[y[t].item()])} matches tape: {match}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "78bef65a",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.7 Shape of one chunk\n",
|
||
"\n",
|
||
"Right now `x` and `y` are **1-D** tensors of shape `(block_size,)` — one chunk,\n",
|
||
"one row. In section 5 we stack many chunks:\n",
|
||
"\n",
|
||
"- `(B, T)` where **B** = batch size, **T** = `block_size`\n",
|
||
"- Later the model outputs `(B, T, C)` — **C** = `vocab_size` scores per step\n",
|
||
"\n",
|
||
"Same idea, just more rows for GPU parallelism.\n",
|
||
"\n",
|
||
"**→ Training:** `(B, T)` tells us we get **B×T** gradient signals per step. Later `(B, T, C)` is one score vector per position for cross-entropy.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8bb27b4f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"x = train_data[1000 : 1000 + block_size]\n",
|
||
"y = train_data[1001 : 1001 + block_size]\n",
|
||
"\n",
|
||
"print(\"x shape:\", x.shape, \" y shape:\", y.shape)\n",
|
||
"print(\"T (time steps) = block_size =\", block_size)\n",
|
||
"print()\n",
|
||
"# pretend we had batch_size=3 — same T, three independent rows\n",
|
||
"B = 3\n",
|
||
"print(f\"batched inputs would be shape ({B}, {block_size}) = (B, T)\")\n",
|
||
"print(f\"model logits later: ({B}, {block_size}, {vocab_size}) = (B, T, C)\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f86da52d",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.8 Put it together — build your own `x` and `y`\n",
|
||
"\n",
|
||
"Pick any valid `start` index, build `x` and `y`, decode them, and verify every\n",
|
||
"target. This is exactly what `get_batch` does — but with random `start` values.\n",
|
||
"\n",
|
||
"**→ Training:** Same logic as `get_batch` — you're manually doing what the training loop automates thousands of times.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c2ee21e5",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"start = 5000 # change me\n",
|
||
"x = train_data[start : start + block_size]\n",
|
||
"y = train_data[start + 1 : start + block_size + 1]\n",
|
||
"\n",
|
||
"print(\"start:\", start)\n",
|
||
"print(\"x:\", decode(x.tolist()))\n",
|
||
"print(\"y:\", decode(y.tolist()))\n",
|
||
"checks = [y[t].item() == train_data[start + t + 1].item() for t in range(block_size)]\n",
|
||
"print(\"all correct:\", all(checks))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c1d3c8ce",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 4.9 Why this matters for the Bigram model\n",
|
||
"\n",
|
||
"We feed 8 positions per row, and the loss averages over **all** of them — so one\n",
|
||
"chunk gives 8 gradients per row. But the Bigram architecture only uses the\n",
|
||
"**current** character at each position, not the full prefix. The data is ready for\n",
|
||
"long context; the model is not. Fixing that is what **self-attention** does later.\n",
|
||
"\n",
|
||
"**Section 4 checklist**\n",
|
||
"- [ ] Next-token = one step forward on the tape\n",
|
||
"- [ ] `block_size` = window length / context limit\n",
|
||
"- [ ] `y` is `x` shifted by one\n",
|
||
"- [ ] One chunk = `block_size` prefix drills\n",
|
||
"- [ ] Shapes: `(T,)` now → `(B, T)` in section 5\n",
|
||
"\n",
|
||
"**→ Training:** Data supplies full prefixes; Bigram weights only use the current char. Training still works, but capacity is limited until attention.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "45dafc47",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 5. Batching (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 5):** many `x`/`y` windows at once as matrices.\n",
|
||
"\n",
|
||
"**Needs:** sections 3-4 (`train_data`, `block_size`, shift-by-one).\n",
|
||
"**Unlocks:** `get_batch()` -> every `forward` in training (section 8) starts here.\n",
|
||
"Section 4 used **one** chunk (shape `(T,)`). GPUs want **many** chunks at once.\n",
|
||
"We stack `batch_size` independent windows → shape **`(B, T)`**.\n",
|
||
"\n",
|
||
"- **B** = `batch_size` (parallel snippets)\n",
|
||
"- **T** = `block_size` (time steps / context length)\n",
|
||
"\n",
|
||
"**→ Training:** Each optimizer step uses one batch `(xb, yb)`. Larger `B` = more examples contributing to the **same** weight update.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "ffc77ac6",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 5.1 Why batch?\n",
|
||
"\n",
|
||
"**Analogy:** grading one exam at a time vs a stack of 32 exams — same rules,\n",
|
||
"more throughput. Random start indices each step = model sees different parts of\n",
|
||
"the book every iteration.\n",
|
||
"\n",
|
||
"**→ Training:** Random windows each step ≈ shuffled exposure to the book. Parallel rows = faster steps on GPU.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0b02e589",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 5.2 Stack two chunks by hand\n",
|
||
"\n",
|
||
"Before `get_batch`, see how two slices become a 2×8 matrix.\n",
|
||
"\n",
|
||
"**→ Training:** `get_batch` does this with `batch_size` rows. One `optimizer.step()` updates weights using loss from **all** rows.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "11fc074b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"torch.manual_seed(0)\n",
|
||
"ix = torch.randint(len(train_data) - block_size, (2,))\n",
|
||
"chunk_a = train_data[ix[0] : ix[0] + block_size]\n",
|
||
"chunk_b = train_data[ix[1] : ix[1] + block_size]\n",
|
||
"stacked = torch.stack([chunk_a, chunk_b])\n",
|
||
"print(\"stacked shape:\", stacked.shape, \" (= 2 rows, block_size cols)\")\n",
|
||
"for row in range(2):\n",
|
||
" print(f\" row {row}:\", decode(stacked[row].tolist()))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c296b808",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 5.3 `get_batch` — random windows\n",
|
||
"\n",
|
||
"`xb` = inputs, `yb` = targets (each row shifted +1 on the tape). Same shapes.\n",
|
||
"\n",
|
||
"**→ Training:** Called every training iteration. Returns the `(xb, yb)` pair that `loss.backward()` will score and correct.\n",
|
||
"\n",
|
||
"**In the chain:** `get_batch` is **step 0** of every training iteration — supplies `xb` (questions) and `yb` (answers) before `model()` can run.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "67c1ce6f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"torch.manual_seed(1337)\n",
|
||
"batch_size = 4\n",
|
||
"\n",
|
||
"def get_batch(split):\n",
|
||
" d = train_data if split == \"train\" else val_data\n",
|
||
" ix = torch.randint(len(d) - block_size, (batch_size,))\n",
|
||
" x = torch.stack([d[i : i + block_size] for i in ix])\n",
|
||
" y = torch.stack([d[i + 1 : i + block_size + 1] for i in ix])\n",
|
||
" return x.to(device), y.to(device)\n",
|
||
"\n",
|
||
"xb, yb = get_batch(\"train\")\n",
|
||
"print(\"xb shape:\", xb.shape, \" yb shape:\", yb.shape)\n",
|
||
"print(xb)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "e72033cb",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 5.4 Decode a batch row\n",
|
||
"\n",
|
||
"Row `b`, column `t`: `yb[b,t]` is the next token after `xb[b,t]` on the tape.\n",
|
||
"\n",
|
||
"**→ Training:** Sanity check that `yb` is the answer key aligned with `xb` before trusting the loss number.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d8be50cf",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"b = 0\n",
|
||
"print(\"row\", b, \"input: \", decode(xb[b].tolist()))\n",
|
||
"print(\"row\", b, \"target:\", decode(yb[b].tolist()))\n",
|
||
"for t in range(block_size):\n",
|
||
" print(f\" t={t} x={repr(itos[xb[b,t].item()])} y={repr(itos[yb[b,t].item()])}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "3f69fc3b",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 5** — Call `get_batch(\"val\")` and decode row 2.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "cd1c51dc",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"xb_val, yb_val = get_batch(\"val\")\n",
|
||
"print(decode(xb_val[2].tolist()))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "883c2705",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 5.5 `.to(device)`\n",
|
||
"\n",
|
||
"Moves tensors to GPU so the model (also on GPU) can process them in one place.\n",
|
||
"\n",
|
||
"**Section 5 checklist**\n",
|
||
"- [ ] `(B, T)` = batch of `block_size`-long windows\n",
|
||
"- [ ] `yb` is `xb` shifted +1 on the tape\n",
|
||
"- [ ] `get_batch` randomises which part of the book each step\n",
|
||
"\n",
|
||
"**→ Training:** Moves batch tensors next to **weight matrices** on GPU so forward + backward run without device errors.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0e99e2c5",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 6. Logits, loss, and the Bigram model (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 6):** the **weights** + how we score wrongness.\n",
|
||
"\n",
|
||
"**Needs:** section 5 (`xb`, `yb` shapes); section 2 (`vocab_size`).\n",
|
||
"**Unlocks:** `model(xb,yb)->loss` for `backward()` (section 8); `model(xb)->sample` for generation (section 7).\n",
|
||
"The **Bigram** model: next character depends **only** on the current character.\n",
|
||
"Before the full class, we unpack **logits**, **softmax**, and **cross-entropy** —\n",
|
||
"the same three ideas GPT uses at scale.\n",
|
||
"\n",
|
||
"**→ Training:** Forward produces **logits**; loss measures wrongness; backward rewrites the **embedding matrix** — the only weights in this model.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f6bb711d",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 6.1 Logits — raw scores\n",
|
||
"\n",
|
||
"For each possible next token, the model outputs a **logit** (unnormalised score).\n",
|
||
"Higher = more likely. Length = `vocab_size`.\n",
|
||
"\n",
|
||
"**→ Training:** Logits are not trained directly — **weights** produce logits. Gradients flow back through logits into those weights.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e7428034",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"toy_logits = torch.tensor([2.0, 0.5, 0.1]) # pretend vocab_size=3\n",
|
||
"print(\"logits:\", toy_logits.tolist())\n",
|
||
"print(\"argmax (favourite class):\", toy_logits.argmax().item())\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "ff289741",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 6.2 Softmax — scores → probabilities\n",
|
||
"\n",
|
||
"Softmax squashes logits into positive numbers that **sum to 1** — a distribution.\n",
|
||
"\n",
|
||
"**→ Training:** Used heavily in **generation** (sampling). Training uses cross-entropy on logits directly (softmax is inside that).\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "71e447c8",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"probs = F.softmax(toy_logits, dim=0)\n",
|
||
"print(\"probs:\", [round(p, 3) for p in probs.tolist()])\n",
|
||
"print(\"sum:\", probs.sum().item())\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "e60d5dd1",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 6.3 Cross-entropy loss — how wrong were we?\n",
|
||
"\n",
|
||
"Compare model distribution to the **true** next token id.\n",
|
||
"- right answer gets high prob → low loss\n",
|
||
"- wrong answer → high loss\n",
|
||
"\n",
|
||
"Random guessing over `vocab_size` tokens → loss ≈ `-ln(1/vocab_size)`.\n",
|
||
"\n",
|
||
"**→ Training:** The single number `loss` that `backward()` uses. Lower loss ⇒ nudge weights so true next tokens get higher score.\n",
|
||
"\n",
|
||
"**In the chain:** cross-entropy turns logits + `yb` into the **single number** `loss`. Section 8 cannot run `backward()` until this exists.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c9166861",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"for target in [0, 1, 2]:\n",
|
||
" loss = F.cross_entropy(toy_logits.unsqueeze(0), torch.tensor([target]))\n",
|
||
" print(f\"true class {target} -> loss {loss.item():.3f}\")\n",
|
||
"expected = -torch.log(torch.tensor(1.0 / vocab_size))\n",
|
||
"print(f\"random baseline (vocab {vocab_size}): {expected.item():.3f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7b45f3a7",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 6** — Make class 2 strongest; loss for target 2 should drop.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "dfe98faf",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"strong = torch.tensor([0.1, 0.5, 2.5])\n",
|
||
"for target in [0, 2]:\n",
|
||
" loss = F.cross_entropy(strong.unsqueeze(0), torch.tensor([target]))\n",
|
||
" print(f\"target {target} -> {loss.item():.3f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7926ed10",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 6.4 Bigram = one lookup table\n",
|
||
"\n",
|
||
"`nn.Embedding(vocab_size, vocab_size)` stores a matrix: **row i** = logits for\n",
|
||
"\"what follows token i?\" No memory of earlier tokens.\n",
|
||
"\n",
|
||
"**Analogy:** a cheat sheet — \"after `t` usually comes `h` or `e`\" — but only\n",
|
||
"looking at the **last** letter.\n",
|
||
"\n",
|
||
"**→ Training:** The entire learned model is this **one matrix** (~vocab² numbers). Training adjusts every row: \"what follows char i?\"\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "f705fa50",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"table = nn.Embedding(vocab_size, vocab_size)\n",
|
||
"ch = 't'\n",
|
||
"tid = stoi[ch]\n",
|
||
"row = table.weight[tid]\n",
|
||
"print(\"row for\", repr(ch), \"has shape:\", row.shape)\n",
|
||
"top_p, top_i = F.softmax(row, dim=0).topk(3)\n",
|
||
"print(\"top-3 followers (random weights):\")\n",
|
||
"for p, i in zip(top_p, top_i):\n",
|
||
" print(f\" {repr(itos[i.item()])}: {p.item():.3f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "e5208051",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 6.5 The full `BigramLanguageModel`\n",
|
||
"\n",
|
||
"Forward: `idx` `(B,T)` → logits `(B,T,C)`. Flatten to `(B*T, C)` for one big\n",
|
||
"cross-entropy over every position in the batch.\n",
|
||
"\n",
|
||
"**→ Training:** `forward(xb, yb)` = predict on all positions, compare to `yb`, return scalar loss. `generate()` uses same weights without `yb`.\n",
|
||
"\n",
|
||
"**In the chain:** `forward(xb, yb)` is **step 1** of training. Without `yb`, you get logits only (generation path) — no loss, no backward.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "f7cc9f29",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class BigramLanguageModel(nn.Module):\n",
|
||
" def __init__(self, vocab_size):\n",
|
||
" super().__init__()\n",
|
||
" self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
|
||
"\n",
|
||
" def forward(self, idx, targets=None):\n",
|
||
" logits = self.token_embedding_table(idx)\n",
|
||
" if targets is None:\n",
|
||
" return logits, None\n",
|
||
" B, T, C = logits.shape\n",
|
||
" loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))\n",
|
||
" return logits, loss\n",
|
||
"\n",
|
||
" def generate(self, idx, max_new_tokens):\n",
|
||
" for _ in range(max_new_tokens):\n",
|
||
" logits, _ = self(idx)\n",
|
||
" logits = logits[:, -1, :]\n",
|
||
" probs = F.softmax(logits, dim=-1)\n",
|
||
" idx_next = torch.multinomial(probs, num_samples=1)\n",
|
||
" idx = torch.cat((idx, idx_next), dim=1)\n",
|
||
" return idx\n",
|
||
"\n",
|
||
"model = BigramLanguageModel(vocab_size).to(device)\n",
|
||
"xb, yb = get_batch(\"train\")\n",
|
||
"logits, loss = model(xb, yb)\n",
|
||
"print(\"logits shape (B, T, C):\", logits.shape)\n",
|
||
"print(f\"initial loss: {loss.item():.4f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1614313d",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 6b** — Initial loss should be near `-ln(1/vocab_size)`. Compare.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "25b7140e",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"baseline = -torch.log(torch.tensor(1.0 / vocab_size)).item()\n",
|
||
"print(f\"loss {loss.item():.4f} vs random baseline {baseline:.4f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "23f37864",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Section 6 checklist**\n",
|
||
"- [ ] Logits → softmax → probs (generation)\n",
|
||
"- [ ] Cross-entropy compares preds to true `yb` (training)\n",
|
||
"- [ ] Bigram row `i` = next-token scores after char `i` only\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8f6036e2",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 7. Generation before training (overview)\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 7):** use weights **without** training (no `yb`, no `backward`).\n",
|
||
"\n",
|
||
"**Needs:** section 6 (`model`, random weights).\n",
|
||
"**Unlocks:** baseline gibberish to compare after section 9 — proof that training changed weights.\n",
|
||
"**Generation** = autoregressive sampling: predict one token, append it, repeat.\n",
|
||
"With random weights, output is gibberish — our **before** snapshot.\n",
|
||
"\n",
|
||
"**→ Training:** With **random** weights, generation is nonsense. This is the baseline before any `optimizer.step()` changes the matrices.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f51c9655",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 7.1 One generation step by hand\n",
|
||
"\n",
|
||
"Context → forward → logits at **last** position → softmax → sample one id.\n",
|
||
"\n",
|
||
"**→ Training:** Same forward path as training, but no `yb` and no `backward()`. Shows what one row of trained weights **does** at inference.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "f9cdda7f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"context = torch.zeros((1, 1), dtype=torch.long, device=device) # start: newline\n",
|
||
"logits, _ = model(context)\n",
|
||
"last_logits = logits[0, -1, :]\n",
|
||
"probs = F.softmax(last_logits, dim=0)\n",
|
||
"sampled = torch.multinomial(probs, num_samples=1)\n",
|
||
"print(\"context:\", repr(itos[context[0,0].item()]))\n",
|
||
"print(\"sampled next char:\", repr(itos[sampled.item()]))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f8461a58",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 7.2 Full generation (200 chars, untrained)\n",
|
||
"\n",
|
||
"Start from newline (id 0). Expect nonsense.\n",
|
||
"\n",
|
||
"**→ Training:** Save mentally (or in output) for contrast with section 9 after weights have been updated.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "4b5c3e45",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
|
||
"out = model.generate(context, max_new_tokens=200)\n",
|
||
"print(decode(out[0].tolist()))\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "eb818ac2",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 8. Training (overview)\n",
|
||
"\n",
|
||
"**You are here (section 8):** run the full chain thousands of times.\n",
|
||
"\n",
|
||
"**Needs:** ALL above — especially `get_batch` (5), `model`+loss (6).\n",
|
||
"**Unlocks:** updated weights — the **result** of training.\n",
|
||
"\n",
|
||
"One iteration — **do not reorder**:\n",
|
||
"\n",
|
||
"```\n",
|
||
"prereq: xb,yb = get_batch() # step 0 — sections 4-5\n",
|
||
"step 1: logits, loss = model(xb,yb) # forward — section 6\n",
|
||
"step 2: (loss is ready) # scalar surprise\n",
|
||
"step 3: loss.backward() # gradients per weight\n",
|
||
"step 4: optimizer.step() # nudge weights\n",
|
||
"```\n",
|
||
"\n",
|
||
"Plus `zero_grad` before step 1 each time — wipe last batch's gradients.\n",
|
||
"\n",
|
||
"**-> Training:** This section is where the ultimate goal happens: weights change.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4ac65ed4",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.0 What training is — fundamentals\n",
|
||
"\n",
|
||
"Strip away the jargon and training is four ideas:\n",
|
||
"\n",
|
||
"1. **Weights** — numbers in matrices (our Bigram: one `vocab_size × vocab_size` table).\n",
|
||
"2. **Forward** — multiply/lookup through weights → predictions (logits).\n",
|
||
"3. **Loss** — one number: how wrong predictions are vs true next tokens `yb`.\n",
|
||
"4. **Update** — nudge every weight a tiny bit to lower loss (`backward` + `step`).\n",
|
||
"\n",
|
||
"**Analogy:** weights = knobs on a mixer; loss = how bad the song sounds; training turns knobs repeatedly until the mix fits the data.\n",
|
||
"\n",
|
||
"The model does **not** store the book inside weights. It stores **habits** like \"after `t` often comes `h`.\"\n",
|
||
"\n",
|
||
"**→ Training:** This is the goal every earlier section prepared: data → batches → loss → changed weights.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0881b2c7",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# THE trainable artifact: model parameters (weight matrices)\n",
|
||
"n_params = sum(p.numel() for p in model.parameters())\n",
|
||
"print(\"trainable parameters:\", n_params)\n",
|
||
"print(\"embedding weight shape:\", model.token_embedding_table.weight.shape)\n",
|
||
"print(\"dtype:\", model.token_embedding_table.weight.dtype)\n",
|
||
"print()\n",
|
||
"print(\"These numbers ARE what training rewrites.\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "362158ed",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.0b One step = forward, loss, backward, update\n",
|
||
"\n",
|
||
"| Step | What happens | Uses from earlier sections |\n",
|
||
"|------|----------------|---------------------------|\n",
|
||
"| forward | `model(xb, yb)` → logits, loss | `get_batch`, `xb`/`yb`, embedding |\n",
|
||
"| zero_grad | clear old gradients | — |\n",
|
||
"| backward | `loss.backward()` computes ∂loss/∂weight | cross-entropy on logits |\n",
|
||
"| step | `optimizer.step()` adds small change to weights | AdamW |\n",
|
||
"\n",
|
||
"One step touches the **same** weight matrix thousands of times indirectly (via every position in the batch), then applies one averaged update.\n",
|
||
"\n",
|
||
"Run **8.1** next (creates `optimizer`), then the cell below watches one weight change.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "eb72fa20",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.0d Backward (`loss.backward`) — why it exists in the chain\n",
|
||
"\n",
|
||
"**Goal of backward:** figure out *how each weight should move* to make loss smaller.\n",
|
||
"\n",
|
||
"You cannot run backward until the steps below exist:\n",
|
||
"\n",
|
||
"| Prerequisite | Built in | Provides |\n",
|
||
"|--------------|----------|----------|\n",
|
||
"| **Step 1** `xb`, `yb` | sections 4-5 | input + answer key |\n",
|
||
"| **Step 2** `logits`, `loss` | section 6 forward | one scalar \"how wrong\" |\n",
|
||
"| **Step 3** `loss.backward()` | **this step** | `.grad` on each weight |\n",
|
||
"| **Step 4** `optimizer.step()` | section 8.1+ | weights actually change |\n",
|
||
"\n",
|
||
"**Analogy:** loss = \"how off-key is the song?\" Backward = \"for each knob,\n",
|
||
"which direction fixes pitch?\" Step = \"turn the knobs slightly.\"\n",
|
||
"\n",
|
||
"Without step 2 (loss), step 3 has nothing to differentiate — PyTorch needs a\n",
|
||
"single number that depends on the weights.\n",
|
||
"\n",
|
||
"**-> Training:** backward is the bridge between \"we were wrong\" (loss) and\n",
|
||
"\"here is how to fix each weight\" (gradients).\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "019a14ab",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# trace the chain on one batch (steps 1 -> 2 -> 3)\n",
|
||
"xb, yb = get_batch(\"train\")\n",
|
||
"print(\"STEP 1 prerequisite: batch shapes\", xb.shape, yb.shape)\n",
|
||
"\n",
|
||
"model.zero_grad(set_to_none=True)\n",
|
||
"logits, loss = model(xb, yb)\n",
|
||
"print(\"STEP 2: forward produced loss =\", round(loss.item(), 4), \"(scalar)\")\n",
|
||
"\n",
|
||
"loss.backward() # STEP 3\n",
|
||
"w = model.token_embedding_table.weight\n",
|
||
"print(\"STEP 3: backward filled w.grad with shape\", w.grad.shape)\n",
|
||
"print(\" example grad[t,h]:\", w.grad[stoi[\"t\"], stoi[\"h\"]].item())\n",
|
||
"print(\"STEP 4: run optimizer.step() to add grad into weights (see 8.1)\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "843b881d",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 8d** — Run the cell above, then answer:\n",
|
||
"1. What shape is `w.grad`? (Same as weights — one partial derivative per entry.)\n",
|
||
"2. Why must `loss` be a single number, not a vector?\n",
|
||
"3. What section would break if `yb` were missing?\n",
|
||
"\n",
|
||
"*(2: backward needs one objective to differentiate. 3: section 6 loss — no answer key.)*\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "80359cd3",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.1 One training step\n",
|
||
"\n",
|
||
"Run this once to see loss change after a single update.\n",
|
||
"\n",
|
||
"**→ Training:** Minimal version of the full loop. If loss changes after `step()`, weights moved — training is working.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8949b082",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
|
||
"xb, yb = get_batch(\"train\")\n",
|
||
"_, loss_before = model(xb, yb)\n",
|
||
"optimizer.zero_grad(set_to_none=True)\n",
|
||
"_, loss = model(xb, yb)\n",
|
||
"loss.backward()\n",
|
||
"optimizer.step()\n",
|
||
"_, loss_after = model(xb, yb)\n",
|
||
"print(f\"before: {loss_before.item():.4f} after: {loss_after.item():.4f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8d2cc5da",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# watch ONE matrix entry change (run after 8.1 created optimizer)\n",
|
||
"w = model.token_embedding_table.weight\n",
|
||
"before = w[stoi['t'], stoi['h']].item()\n",
|
||
"\n",
|
||
"xb, yb = get_batch(\"train\")\n",
|
||
"_, loss = model(xb, yb)\n",
|
||
"optimizer.zero_grad(set_to_none=True)\n",
|
||
"loss.backward()\n",
|
||
"optimizer.step()\n",
|
||
"\n",
|
||
"after = w[stoi['t'], stoi['h']].item()\n",
|
||
"print(f\"loss: {loss.item():.4f}\")\n",
|
||
"print(f\"weight score for t->h: {before:.6f} -> {after:.6f}\")\n",
|
||
"print(\"changed:\", before != after)\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4b41e5fc",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.0c Where the result of training is stored\n",
|
||
"\n",
|
||
"| Location | What | Persists? |\n",
|
||
"|----------|------|-----------|\n",
|
||
"| `model.token_embedding_table.weight` | learned matrix in RAM/VRAM | until kernel restarts |\n",
|
||
"| `model.state_dict()` | name → tensor map of all weights | same |\n",
|
||
"| disk (`torch.save`) | checkpoint file | yes — **not done by default** |\n",
|
||
"\n",
|
||
"Re-run the notebook from the top → new random weights → must train again unless you load a saved checkpoint.\n",
|
||
"\n",
|
||
"**→ Training:** After section 8.2, `state_dict()` holds the **result**. Section 9 reads those same weights to generate text.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "66fd4109",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"sd = model.state_dict()\n",
|
||
"print(\"keys in state_dict:\", list(sd.keys()))\n",
|
||
"print(\"weight tensor shape:\", sd[\"token_embedding_table.weight\"].shape)\n",
|
||
"\n",
|
||
"# optional: save / load (uncomment to use across sessions)\n",
|
||
"# torch.save(sd, \"bigram_shakespeare.pt\")\n",
|
||
"# model.load_state_dict(torch.load(\"bigram_shakespeare.pt\", map_location=device))\n",
|
||
"print(\"trained weights live in memory inside model until you torch.save them.\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "603bd9b0",
|
||
"metadata": {},
|
||
"source": [
|
||
"### 8.2 Full training loop + val loss\n",
|
||
"\n",
|
||
"Every 1000 steps we average loss on train **and** val batches. Val should track\n",
|
||
"train; if train drops but val does not → overfitting.\n",
|
||
"\n",
|
||
"**→ Training:** 10,000 updates to the embedding matrix. Printed val loss checks that weights generalise, not only memorise train slices.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "b34ee95e",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"batch_size = 32\n",
|
||
"\n",
|
||
"@torch.no_grad()\n",
|
||
"def estimate_loss(eval_iters=100):\n",
|
||
" out = {}\n",
|
||
" model.eval()\n",
|
||
" for split in [\"train\", \"val\"]:\n",
|
||
" losses = torch.zeros(eval_iters)\n",
|
||
" for k in range(eval_iters):\n",
|
||
" X, Y = get_batch(split)\n",
|
||
" _, L = model(X, Y)\n",
|
||
" losses[k] = L.item()\n",
|
||
" out[split] = losses.mean().item()\n",
|
||
" model.train()\n",
|
||
" return out\n",
|
||
"\n",
|
||
"for step in range(10000):\n",
|
||
" xb, yb = get_batch(\"train\")\n",
|
||
" _, loss = model(xb, yb)\n",
|
||
" optimizer.zero_grad(set_to_none=True)\n",
|
||
" loss.backward()\n",
|
||
" optimizer.step()\n",
|
||
" if step % 1000 == 0:\n",
|
||
" L = estimate_loss()\n",
|
||
" print(f\"step {step:5d} train {L['train']:.4f} val {L['val']:.4f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "69ff0dfc",
|
||
"metadata": {},
|
||
"source": [
|
||
"**Your turn 8** — After training, is val loss below the random baseline (~4.17)?\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "23be596a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"L = estimate_loss(50)\n",
|
||
"baseline = -torch.log(torch.tensor(1.0 / vocab_size)).item()\n",
|
||
"print(f\"train {L['train']:.4f} val {L['val']:.4f} baseline {baseline:.4f}\")\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1a49aef3",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 9. Generation after training + what's next\n",
|
||
"\n",
|
||
"\n",
|
||
"**You are here (section 9):** see the **goal** achieved (partially — bigram is limited).\n",
|
||
"\n",
|
||
"**Needs:** section 8 (trained weights in `model`).\n",
|
||
"**Unlocks:** next notebook — self-attention for longer context.\n",
|
||
"Same `generate` as section 7 — but weights learned letter/space patterns.\n",
|
||
"Still **not** real Shakespeare: bigram only sees **one** character back.\n",
|
||
"\n",
|
||
"Next notebook: **self-attention** so each token can use the full prefix.\n",
|
||
"\n",
|
||
"**Whole notebook checklist**\n",
|
||
"- [ ] Text → tokens → 1-D tensor tape\n",
|
||
"- [ ] `x`/`y` shift, prefix trick, `(B,T)` batching\n",
|
||
"- [ ] Logits, softmax, cross-entropy\n",
|
||
"- [ ] Train loop + val loss\n",
|
||
"- [ ] Generate before vs after\n",
|
||
"\n",
|
||
"**→ Training:** Proof that weights changed. Better letter patterns = the matrix learned co-occurrence stats. Weights live in `model` until you `torch.save` them.\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "2ef869fb",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
|
||
"print(decode(model.generate(context, max_new_tokens=400)[0].tolist()))\n"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"name": "python",
|
||
"version": "3.12"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|