tiny_llm_from_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d6dbef8d",
   "metadata": {},
   "source": [
    "# Tiny Storyteller: building a language model from scratch\n",
    "\n",
    "**Author:** singhpratech. **Runs on:** Google Colab (free T4 GPU), or a Mac with Apple Silicon, or CPU.\n",
    "\n",
    "## Purpose\n",
    "Learn how a language model (an \"LLM\") is really built, by building one from nothing.\n",
    "No pretrained weights (dials borrowed from a model someone else trained), no fine-tuning\n",
    "(nudging borrowed dials), no magic. We start with raw PyTorch tensors, build the simplest\n",
    "possible model, then add one idea at a time until we have a small GPT-style transformer that\n",
    "we train ourselves. If those words mean nothing yet, good: the next section explains every\n",
    "one of them.\n",
    "\n",
    "## Use case\n",
    "A **bedtime story starter** for young kids. Give it the first few words, like\n",
    "*\"Once upon a time, a little monkey\"*, and it continues with a short, simple story.\n",
    "\n",
    "## Two characters, one story\n",
    "The whole notebook is told through **Momo**, a young monkey who is learning to tell stories,\n",
    "and **Ella**, an elephant who never forgets and reads to him every night. Momo is the model.\n",
    "Ella is the teacher and the memory. Every idea below, and every clickable picture, uses the\n",
    "same two characters, so nothing has to be re-learned from scratch.\n",
    "\n",
    "## The road map\n",
    "| Part | What we build | Why |\n",
    "|---|---|---|\n",
    "| 0 | Momo and Ella's story, a word list, a live flow diagram, then setup | Every idea and every strange word in plain language |\n",
    "| 1 | Tensors and gradients | The two things every neural net is made of |\n",
    "| 2 | Data: TinyStories | Simple English stories, made for small models |\n",
    "| 3 | A bigram model | The simplest language model that can \"write\" |\n",
    "| 4 | A real tokenizer (BPE) | Words instead of letters |\n",
    "| 5 | Attention, by hand | The idea that makes transformers work |\n",
    "| 6 | A tiny GPT | Stack the pieces into a full model |\n",
    "| 7 | Training | Teach it with a few minutes on a GPU |\n",
    "| 8 | Generate stories | Try the use case |\n",
    "| 9 | Publish on Hugging Face | Save the model, privately at first |\n",
    "| 10 | Load it back, and what we built | Use Momo anywhere |\n",
    "\n",
    "Run every cell top to bottom. Each cell is small on purpose."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52bce3d0",
   "metadata": {},
   "source": [
    "## 0. Before any code: Momo, Ella, and what we are about to build\n",
    "\n",
    "*No code in this section. It is written so that a five-year-old could follow the story and a\n",
    "grown-up could follow the meaning. Every strange word is explained the first time it appears,\n",
    "and all of them are collected in the word list at the end of this section.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18badc67",
   "metadata": {},
   "source": [
    "### Meet the characters\n",
    "\n",
    "**Momo** is a young monkey. He is quick and curious and wants to learn to tell bedtime\n",
    "stories. Inside Momo's head are millions of tiny **dials**. Right now every dial points\n",
    "somewhere random, so when Momo tries to tell a story it comes out as nonsense.\n",
    "\n",
    "**Ella** is an elephant, and elephants never forget. Every night Ella reads a story to Momo\n",
    "one word at a time. Before each word she stops and lets Momo guess what comes next. Then she\n",
    "tells him how wrong he was. Ella also keeps a shelf where she stores everything already said,\n",
    "so nobody has to work it out twice.\n",
    "\n",
    "In grown-up words: Momo is the **model**. His dials are the **weights**, also called\n",
    "**parameters**: the numbers that get adjusted during learning. Ella's stories are the\n",
    "**training data**. Her \"how wrong were you\" is the **loss**. Her shelf is the **KV cache**.\n",
    "Ella plays two roles on purpose: the teacher while Momo learns (training), and the memory\n",
    "while Momo writes (generation). Each of these gets its own section below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b124e5",
   "metadata": {},
   "source": [
    "### 0.1 What is PyTorch?\n",
    "\n",
    "Momo needs a place to keep his dials and do his sums very fast. That place is **PyTorch**,\n",
    "Momo's workshop. It has three things:\n",
    "\n",
    "- **Boxes for numbers**, called **tensors**. A tensor is a grid of numbers with a shape: a\n",
    "  row of baskets, a tray of baskets, or a stack of trays. Everything in Momo's head is a\n",
    "  tensor: the words coming in, the dials, the guesses going out.\n",
    "- **A very fast counting table**, the **GPU** (graphics processing unit: a chip first built\n",
    "  for video games that can do millions of small sums at the same time). On Colab we borrow\n",
    "  one called a T4.\n",
    "- **A tape recorder**, called **autograd** (automatic gradients). While Momo works it writes\n",
    "  down every step. Afterwards it reads the tape backwards and tells Momo which way to turn\n",
    "  every dial, and by how much. Ordinary programs do not do this, and it is the reason modern AI\n",
    "  is practical at all. Part 1 has a game that shows exactly how the tape is read.\n",
    "\n",
    "PyTorch was built at Meta's AI lab and is behind most famous models (the GPT family, Llama,\n",
    "Stable Diffusion, Whisper). **Hugging Face**, where we publish at the end, is the public\n",
    "library where people share models built this way."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8557a5e",
   "metadata": {},
   "source": [
    "### 0.2 What is a model, really?\n",
    "\n",
    "A **model** is Momo: a big pile of dials plus a fixed recipe for how to use them. Our Momo\n",
    "has about 5.9 million dials. Famous big models have hundreds of billions. At the start every\n",
    "dial is random. **Training** means turning the dials a tiny bit at a time, thousands of\n",
    "times, until the recipe gives good guesses. Momo never writes down a rulebook. He tunes his\n",
    "instincts. The dials *are* the instincts. (He can still end up repeating sentences he has\n",
    "heard many times, as we will see in Part 8.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f05878b3",
   "metadata": {},
   "source": [
    "### 0.3 What does a language model actually do?\n",
    "\n",
    "One thing: **given some words, guess the next word.** That is the whole job.\n",
    "\n",
    "Ella reads: \"Once upon a time, a little...\" Momo guesses: \"monkey!\" (or \"girl\", or \"fox\").\n",
    "\n",
    "Telling a whole story is just that guess over and over: guess a word, add it, guess the next\n",
    "word, add it. Chatting with an AI assistant is the same trick. Your question is the words so\n",
    "far, and the answer is a long run of next-word guesses.\n",
    "\n",
    "A **language model** is any machine that does this. An **LLM** is a *large* language model,\n",
    "one with billions of dials. Ours is tiny, but it works the same way."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8774b740",
   "metadata": {},
   "source": [
    "### 0.4 What happens when you ask Momo a question?\n",
    "\n",
    "Say you ask: *\"What is the capital of France?\"* Here is the journey.\n",
    "\n",
    "**Step 1, biting: words become tokens.** Momo cannot swallow a whole sentence. He bites it\n",
    "into pieces: \"What\", \" is\", \" the\", \" capital\", \" of\", \" France\", \"?\". Each kind of piece\n",
    "has a number painted on it. A piece is called a **token**, the biter is the **tokenizer**,\n",
    "and the set of all pieces Momo knows is his **vocabulary** (ours has 4,096 pieces). After\n",
    "this step your question is a short list of numbers.\n",
    "\n",
    "**Step 2, the meaning map: each token becomes a point.** Momo keeps a big map. Every token\n",
    "has a spot on it, and tokens that mean similar things sit close together: \"France\" near\n",
    "\"Spain\", \"capital\" near \"city\". A spot on the map is written as a list of numbers (256 of\n",
    "them in our model). That list is called an **embedding**, and the whole map is learned\n",
    "during training.\n",
    "\n",
    "**Step 3, looking around: attention.** This is the heart of it. Each word looks back at all\n",
    "the words before it (and at itself) and asks \"which of you matter to me right now?\" The word\n",
    "\"capital\" looks around, finds \"France\", and pulls some of France's meaning into itself. Now\n",
    "\"capital\" means \"the capital of France\". This looking-around is called **attention**. It is\n",
    "repeated in several rounds, called **layers**. Picture a tall jungle tree that Momo climbs:\n",
    "each rung is one round of looking around, and the meaning gets sharper at every rung. Our\n",
    "Momo's tree has 6 rungs.\n",
    "\n",
    "**Step 4, the guess.** The last word's sharpened meaning is compared with every token in\n",
    "the vocabulary, giving each a score. The raw scores are called **logits**. They are squashed\n",
    "into percentages that add up to 100, and that squashing is called **softmax**. \"Paris\" gets,\n",
    "say, 92 percent. One token is picked.\n",
    "\n",
    "**Step 5, repeat.** \"Paris\" is added to the words so far and the whole journey runs again for\n",
    "the next token, until Momo picks a special \"the story is over\" token or reaches the length limit. You see words appear one by one\n",
    "on screen because that is literally how they are made."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a0ae3aa",
   "metadata": {},
   "source": [
    "### 0.5 How attention looks things up: the story circle\n",
    "\n",
    "Picture the jungle story circle. Every word in the sentence is an animal sitting in the\n",
    "circle. Each animal holds three cards:\n",
    "\n",
    "- a **query** card: \"here is what I am looking for\" (\"capital\" is looking for a country),\n",
    "- a **key** card: \"here is what I am\" (\"France\" says: I am a country),\n",
    "- a **value** card: \"here is what I will tell you if you pick me\" (France's actual meaning).\n",
    "\n",
    "Every animal holds up its query card and compares it with the key cards of the animals\n",
    "before it, and with its own. A good match scores high, a poor match scores low. The scores\n",
    "are squashed into percentages (softmax again) and the animal collects a blend of everyone's\n",
    "value cards, weighted by those percentages. \"capital\" might take 70 percent from \"France\",\n",
    "20 percent from \"of\", and crumbs from the rest. That blend becomes its new, sharper meaning.\n",
    "\n",
    "One more thing. While Ella is teaching, every animal does this at the same time. While Momo\n",
    "is writing, only the newest animal needs to.\n",
    "\n",
    "Two rules make this work for *writing*:\n",
    "\n",
    "- **No peeking ahead.** When Momo is guessing word 7 he may only look at words 1 to 6. In\n",
    "  code this is a triangle-shaped **mask** that blocks out the future.\n",
    "- **Many heads.** The circle runs several times at once with different cards. Each run is a\n",
    "  **head**. One head might track grammar, another names, another position. Momo has 4 heads.\n",
    "  Big models have dozens."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e786fabb",
   "metadata": {},
   "source": [
    "### 0.6 Ella's shelf: the KV cache\n",
    "\n",
    "Remember Step 5: for each new word the whole journey runs again. Done naively, writing the\n",
    "500th word means re-reading 499 words, and the 501st means re-reading 500. Slower and slower.\n",
    "\n",
    "But the key card and value card of a word already said **never change**. \"France\" holds the\n",
    "same cards whether the story is 7 words or 700 words long. So Ella, who never forgets, keeps\n",
    "every finished word's key and value cards on her shelf. That shelf is the **KV cache** (K for\n",
    "key, V for value; a cache is a place where things are kept for quick reuse). When a new word\n",
    "arrives only *its* cards are made. It compares its query with all the keys on the shelf,\n",
    "collects the values, done. Nothing old is redone.\n",
    "\n",
    "The shelf is why a long chat stays fast, and why long chats need a lot of GPU memory: every\n",
    "token keeps its cards there. A model's **context window** (how much text it can hold at\n",
    "once) is partly how big the shelf is allowed to grow.\n",
    "\n",
    "Part 5 shows the shelf in about a dozen lines of code. Two honest notes: the shelf is emptied at the\n",
    "start of every new story (it is a memory *within* one piece of writing, not between stories),\n",
    "and Momo's own `generate` in Part 6 does not use it, to keep the code short. It redoes the\n",
    "whole journey for every word, which is fine for 200-word stories."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e0da60",
   "metadata": {},
   "source": [
    "### 0.7 The strange names, decoded\n",
    "\n",
    "| Word | In Momo's world | Grown-up meaning |\n",
    "|---|---|---|\n",
    "| **Model** | Momo | A pile of dials plus a recipe for using them |\n",
    "| **Weights / parameters** | Momo's dials | The numbers that training adjusts |\n",
    "| **Training data** | Ella's bedtime stories | The text the model learns from |\n",
    "| **Token / tokenizer / vocabulary** | A bite / the biter / all the bites Momo knows | Text pieces, the splitter, the set of pieces |\n",
    "| **Embedding** | A spot on Momo's meaning map | A list of numbers standing for one token |\n",
    "| **Attention** | Looking around the story circle | Each token blending in earlier tokens by relevance |\n",
    "| **Query, key, value** | The three cards each animal holds | What I seek, what I am, what I give |\n",
    "| **Head** | One run of the circle | One attention pattern; several run in parallel |\n",
    "| **Mask** | No peeking ahead | Blocks future tokens |\n",
    "| **Softmax** | Turning scores into shares of 100 | Raise e to the power of each score, then divide so they add up to 1 |\n",
    "| **Logits** | Raw scores before the squash | Unnormalised outputs of the last layer |\n",
    "| **KV cache** | Ella's shelf | Stored keys and values of past tokens |\n",
    "| **Context window** | How long Ella's shelf can get | Maximum tokens the model holds at once |\n",
    "| **GPT** | Momo's kind of model | Generative Pre-trained Transformer: makes text, learned from text first, built from attention |\n",
    "| **Pre-training / fine-tuning** | Ella's bedtime stories / a few extra lessons on one topic | Training from scratch on lots of text / continuing to train an already-trained model on a small special set |\n",
    "| **Transformer** | The tall tree of story-circle rungs | The 2017 design that stacks attention layers |\n",
    "| **Neural network** | Any pile of dials in layers | The general family transformers belong to |\n",
    "| **Layer / block** | One rung of the tree | One attention round plus a think-alone round |\n",
    "| **Feed-forward network** | Thinking alone on a twig after the circle | A small two-layer network applied per token |\n",
    "| **`nn.Linear`** | One tray of dials | Multiplies its input by a weight matrix (the fruit widget in Part 1) |\n",
    "| **LayerNorm** | Ella steadying the volume | Rescales numbers so training stays stable |\n",
    "| **Residual connection** | Adding thoughts instead of replacing them | Each block adds to its input |\n",
    "| **Loss / cross-entropy** | Ella's \"how wrong were you\" | How surprised the model was by the true next token |\n",
    "| **Gradient** | The slope under Momo's feet | For each dial, how the loss changes if the dial moves |\n",
    "| **Gradient descent** | Walking downhill one stride at a time | Repeatedly step every weight against its gradient |\n",
    "| **Autograd** | Momo's tape recorder | PyTorch recording every step so it can replay it backwards |\n",
    "| **Backpropagation** | Reading the tape backwards, the blame game | Chain rule applied from the loss back to every weight |\n",
    "| **Chain rule** | Multiply by how hard each step pushes | The slope through a chain of steps is the product of each step's slope |\n",
    "| **Optimizer / AdamW** | The rule for turning dials | Uses gradients plus a memory of recent slopes |\n",
    "| **Weight decay** (the W in AdamW) | Ella nudging every dial back toward the middle | Shrinks every weight a little each step so none grows wild |\n",
    "| **Learning rate** | Momo's stride | Step size per update |\n",
    "| **Warm-up / cosine decay** | Tiptoe, stride, tiptoe | Learning-rate schedule: small, then big, then small |\n",
    "| **Batch** | How many story pieces per lesson | Examples processed before one update |\n",
    "| **Validation set** | The stories Ella hides for the quiz | Held-out text the model never trains on; shows real learning, not memorising |\n",
    "| **Step / epoch** | One dial turn / one pass over every story | One update / one pass over the data |\n",
    "| **Perplexity** | How many tokens Momo is choosing between | exp(loss); 1 is perfect, vocabulary size is random |\n",
    "| **Temperature / top-k** | How wild Momo is when picking | Sampling knobs at generation time |\n",
    "| **BPE** | How Momo decided where to bite | Byte-pair encoding: merge the most common pairs repeatedly |\n",
    "| **GPU** | The fast counting table | Graphics chip doing millions of sums at once |\n",
    "| **Hugging Face Hub / Colab** | The jungle library / a borrowed workshop | Public model library / Google's free GPU notebooks |\n",
    "| **Access token** | Momo's library card | A secret key that lets you upload; nothing to do with text tokens |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91659f55",
   "metadata": {},
   "source": [
    "### The one-sentence version\n",
    "\n",
    "Momo bites text into tokens, places them on his meaning map, lets each word look around the\n",
    "story circle to sharpen its meaning (attention), guesses the next token, and Ella keeps the\n",
    "cards on her shelf so nothing is redone.\n",
    "\n",
    "Now let's build Momo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eabf2c36",
   "metadata": {},
   "source": [
    "## Live flow diagram\n",
    "\n",
    "Run the cell below (it only draws, it trains nothing). Press **Play** and watch Momo produce\n",
    "one word at a time: the text is bitten into tokens, each token gets a spot on the meaning\n",
    "map, the newest token looks around the story circle (the orange arcs are attention), the\n",
    "vocabulary is scored, a word is picked, and it joins the text. On the right, Ella's shelf\n",
    "(the KV cache) fills up one card per token, and only the newest token ever needs new cards.\n",
    "One honest note: the diagram shows the shelf because real models use it. Our own short\n",
    "`generate` in Part 6 skips the shelf and redoes the whole journey for every word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1baea7bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML, display\n",
    "\n",
    "# Everything below is plain HTML + JavaScript, no libraries. It runs inside the notebook output.\n",
    "display(HTML(r\"\"\"<div id=\"llmflow\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:980px\">\n",
    "<style>\n",
    "#llmflow .stage{display:inline-block;padding:8px 12px;border-radius:8px;border:2px solid #555;margin:0 4px;min-width:110px;text-align:center;font-size:13px;transition:all .3s;background:#2a2a2a}\n",
    "#llmflow .stage.on{border-color:#ffb347;background:#3a2e1a;box-shadow:0 0 12px #ffb34788}\n",
    "#llmflow .arrow{display:inline-block;color:#888;font-size:18px;vertical-align:middle}\n",
    "#llmflow .chip{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f4f6f;border:1px solid #5a8ab8;font-size:13px;transition:all .3s}\n",
    "#llmflow .chip.new{background:#6a4a1a;border-color:#ffb347}\n",
    "#llmflow .chip.q{outline:2px solid #ffb347}\n",
    "#llmflow .slot{display:inline-block;width:22px;height:30px;margin:2px;border:1px dashed #666;border-radius:4px;vertical-align:middle;font-size:9px;text-align:center;line-height:30px;color:#aaa}\n",
    "#llmflow .slot.full{background:#2f6f4f;border:1px solid #6fcf97;color:#fff}\n",
    "#llmflow .bar{height:14px;background:#5a8ab8;border-radius:3px;margin:2px 0;transition:width .4s}\n",
    "#llmflow button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer;font-size:13px}\n",
    "#llmflow button:hover{background:#444}\n",
    "#llmflow .panel{background:#262626;border-radius:8px;padding:10px 12px;margin-top:10px;font-size:13px;line-height:1.45}\n",
    "#llmflow .lbl{font-size:11px;color:#aaa;text-transform:uppercase;letter-spacing:.05em;margin:8px 0 4px}\n",
    "</style>\n",
    "\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:8px\">How Momo writes a story: live flow</div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin-bottom:10px\">Press <b>Play</b> and watch one word get produced, then the next. Attention weights are illustrative.</div>\n",
    "\n",
    "<div id=\"stages\">\n",
    "  <span class=\"stage\" data-s=\"0\">1. Text so far</span><span class=\"arrow\">&#8594;</span>\n",
    "  <span class=\"stage\" data-s=\"1\">2. Tokenizer<br><small>chop into tokens</small></span><span class=\"arrow\">&#8594;</span>\n",
    "  <span class=\"stage\" data-s=\"2\">3. Meaning map<br><small>token &#8594; vector</small></span><span class=\"arrow\">&#8594;</span>\n",
    "  <span class=\"stage\" data-s=\"3\">4. Attention<br><small>look back</small></span><span class=\"arrow\">&#8594;</span>\n",
    "  <span class=\"stage\" data-s=\"4\">5. Next-word scores</span><span class=\"arrow\">&#8617;</span>\n",
    "  <span class=\"stage\" data-s=\"5\">6. Append word, repeat</span>\n",
    "</div>\n",
    "\n",
    "<div class=\"lbl\">Text so far (each chip is one token)</div>\n",
    "<div id=\"chips\" style=\"min-height:36px\"></div>\n",
    "<svg id=\"arcs\" height=\"70\" style=\"display:block;width:100%\"></svg>\n",
    "\n",
    "<div style=\"display:flex;gap:20px;flex-wrap:wrap\">\n",
    "  <div style=\"flex:1;min-width:280px\">\n",
    "    <div class=\"lbl\">Next-word scores (top 5)</div>\n",
    "    <div id=\"bars\"></div>\n",
    "  </div>\n",
    "  <div style=\"flex:1;min-width:280px\">\n",
    "    <div class=\"lbl\">Ella's shelf (the KV cache): stored key/value cards, one per token</div>\n",
    "    <div id=\"cache\"></div>\n",
    "    <div style=\"font-size:11px;color:#aaa;margin-top:4px\">Only the newest token computes its cards. Old ones are reused.</div>\n",
    "  </div>\n",
    "</div>\n",
    "\n",
    "<div class=\"panel\" id=\"explain\">Press Play.</div>\n",
    "<div style=\"margin-top:10px\">\n",
    "  <button onclick=\"llmPlay()\">&#9654; Play</button>\n",
    "  <button onclick=\"llmStep()\">Step</button>\n",
    "  <button onclick=\"llmReset()\">Reset</button>\n",
    "  <label style=\"font-size:12px;color:#aaa;margin-left:8px\">Speed <input id=\"spd\" type=\"range\" min=\"1\" max=\"5\" value=\"3\" style=\"vertical-align:middle\"></label>\n",
    "</div>\n",
    "</div>\n",
    "\n",
    "<script>\n",
    "(function(){\n",
    "const prompt = [\"Once\",\" upon\",\" a\",\" time\",\",\",\" a\",\" little\",\" monkey\"];\n",
    "const gen = [\n",
    "  {w:\" found\", top:[[\" found\",.41],[\" was\",.22],[\" saw\",.14],[\" lived\",.09],[\" had\",.06]], att:[.05,.03,.04,.06,.02,.08,.22,.50]},\n",
    "  {w:\" a\",     top:[[\" a\",.55],[\" the\",.18],[\" some\",.08],[\" her\",.06],[\" his\",.05]], att:[.03,.02,.05,.03,.02,.06,.12,.27,.40]},\n",
    "  {w:\" shiny\", top:[[\" shiny\",.31],[\" big\",.27],[\" red\",.16],[\" small\",.10],[\" little\",.06]], att:[.02,.02,.03,.02,.01,.04,.18,.33,.10,.25]},\n",
    "  {w:\" ball\",  top:[[\" ball\",.46],[\" stone\",.15],[\" box\",.12],[\" hat\",.08],[\" key\",.05]], att:[.02,.01,.02,.02,.01,.03,.08,.28,.06,.12,.35]},\n",
    "  {w:\".\",      top:[[\".\",.52],[\" in\",.14],[\" and\",.12],[\" on\",.08],[\" that\",.05]], att:[.02,.01,.02,.03,.02,.02,.05,.15,.05,.10,.18,.35]},\n",
    "];\n",
    "let tokens=[], gi=0, stage=-1, timer=null;\n",
    "const $=id=>document.getElementById(id);\n",
    "const txt={\n",
    " 0:\"<b>1. Text so far.</b> Everything written so far, prompt plus generated words, is the input. Nothing else.\",\n",
    " 1:\"<b>2. Tokenizer.</b> The text is chopped into tokens (roughly words, note the leading spaces) and each becomes a number from the model's dictionary.\",\n",
    " 2:\"<b>3. Meaning map.</b> Each token number is swapped for a vector, a long list of values, placing it on a learned map where similar words sit close together.\",\n",
    " 3:\"<b>4. Attention.</b> The newest token (orange outline) looks back at every earlier token. This is Momo looking around the story circle. Arc thickness shows how much it attends to each. It blends their value cards into its own meaning. Future tokens do not exist yet, so there is nothing to peek at. Repeated in 6 layers with 4 heads each.\",\n",
    " 4:\"<b>5. Next-word scores.</b> The refined meaning of the last token is compared with every word in the dictionary. Scores become probabilities; one word is sampled.\",\n",
    " 5:\"<b>6. Append and repeat.</b> The chosen word is added to the text and its key/value cards go onto Ella's shelf, the KV cache. The loop starts again from step 1, but only the new token needs computing.\",\n",
    "};\n",
    "function render(){\n",
    "  document.querySelectorAll('#llmflow .stage').forEach(s=>s.classList.toggle('on',+s.dataset.s===stage));\n",
    "  const c=$('chips'); c.innerHTML='';\n",
    "  tokens.forEach((t,i)=>{const s=document.createElement('span');s.className='chip'+(i===tokens.length-1&&gi>0&&stage<5?' new':'')+(stage===3&&i===tokens.length-1?' q':'');s.textContent=t.replace(/^ /,'␣');c.appendChild(s);});\n",
    "  const k=$('cache'); k.innerHTML='';\n",
    "  const nCached = stage>=5 ? tokens.length : Math.max(0, tokens.length-(gi>0?1:0));\n",
    "  tokens.forEach((t,i)=>{const s=document.createElement('span');s.className='slot'+(i<nCached?' full':'');s.textContent=i<nCached?'KV':'';k.appendChild(s);});\n",
    "  const b=$('bars'); b.innerHTML='';\n",
    "  if(stage>=4 && gi<gen.length){gen[gi].top.forEach(([w,p])=>{b.innerHTML+=`<div style=\"display:flex;align-items:center;gap:8px;font-size:12px\"><span style=\"width:60px;text-align:right\">${w.replace(/^ /,'␣')}</span><div class=\"bar\" style=\"width:${Math.round(p*260)}px\"></div><span style=\"color:#aaa\">${Math.round(p*100)}%</span></div>`;});}\n",
    "  else b.innerHTML='<div style=\"color:#666;font-size:12px\">(appears at step 5)</div>';\n",
    "  drawArcs();\n",
    "  $('explain').innerHTML = stage<0 ? 'Press Play.' : txt[stage];\n",
    "}\n",
    "function drawArcs(){\n",
    "  const svg=$('arcs'); svg.innerHTML='';\n",
    "  if(stage!==3||gi>=gen.length) return;\n",
    "  const chips=[...document.querySelectorAll('#chips .chip')]; if(chips.length<2) return;\n",
    "  const base=$('chips').getBoundingClientRect(); const last=chips[chips.length-1].getBoundingClientRect();\n",
    "  const x1=last.left-base.left+last.width/2; const w=gen[gi].att;\n",
    "  chips.slice(0,-1).forEach((ch,i)=>{const r=ch.getBoundingClientRect();const x0=r.left-base.left+r.width/2;const a=w[i]||0.02;\n",
    "    const p=document.createElementNS('http://www.w3.org/2000/svg','path');\n",
    "    p.setAttribute('d',`M${x1},2 Q${(x0+x1)/2},${20+Math.abs(x1-x0)/6} ${x0},2`);\n",
    "    p.setAttribute('fill','none');p.setAttribute('stroke','#ffb347');p.setAttribute('stroke-width',(0.5+a*10).toFixed(1));p.setAttribute('opacity',(0.25+a*1.5).toFixed(2));\n",
    "    svg.appendChild(p);\n",
    "    const t=document.createElementNS('http://www.w3.org/2000/svg','text');t.setAttribute('x',x0);t.setAttribute('y',60);t.setAttribute('fill','#ffb347');t.setAttribute('font-size','10');t.setAttribute('text-anchor','middle');t.textContent=Math.round(a*100)+'%';svg.appendChild(t);});\n",
    "}\n",
    "window.llmReset=function(){clearInterval(timer);timer=null;tokens=prompt.slice();gi=0;stage=-1;render();};\n",
    "window.llmStep=function(){\n",
    "  if(gi>=gen.length){$('explain').innerHTML='<b>Done.</b> The model wrote: <i>'+tokens.join('')+'</i>. Press Reset to watch again.';clearInterval(timer);timer=null;return;}\n",
    "  stage++;\n",
    "  if(stage===5){tokens.push(gen[gi].w);}\n",
    "  if(stage>5){stage=0;gi++; if(gi>=gen.length){render();window.llmStep();return;}}\n",
    "  render();\n",
    "};\n",
    "window.llmPlay=function(){if(timer){clearInterval(timer);timer=null;return;} const ms=[0,1600,1200,900,600,350][+$('spd').value]; timer=setInterval(window.llmStep,ms); window.llmStep();};\n",
    "window.llmReset();\n",
    "window.addEventListener('resize',drawArcs);\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0f4f61",
   "metadata": {},
   "source": [
    "### Setup (still Part 0)\n",
    "On Colab, go to **Runtime > Change runtime type > T4 GPU** first. Then run the next two cells.\n",
    "The first one fetches the toolboxes we need, one time only (on Colab most are already there, so\n",
    "it is quick). The second picks the fastest counting table available and fixes the dice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "478042b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install only the toolboxes that are missing. (On a machine where everything is already installed, this does nothing.)\n",
    "import importlib.util, subprocess, sys\n",
    "missing = [p for p in [\"torch\", \"transformers\", \"datasets\", \"tokenizers\", \"huggingface_hub\"] if importlib.util.find_spec(p) is None]\n",
    "if missing:\n",
    "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *missing])\n",
    "print(\"toolboxes ready\" + (f\" (installed {', '.join(missing)})\" if missing else \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00b89e95",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math, time, json, os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Pick the best device available: NVIDIA GPU (Colab), Apple GPU (Mac), or CPU.\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "torch.manual_seed(42)   # fix the dice: every run starts with the same random dials. Change the number for a different Momo.\n",
    "print(\"PyTorch\", torch.__version__, \"| device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33dd01e2",
   "metadata": {},
   "source": [
    "## 1. Tensors and gradients: the two building blocks\n",
    "\n",
    "A **tensor** is a grid of numbers with a shape. Picture the baskets in Momo's workshop:\n",
    "\n",
    "- **one row of baskets** is 1-D: shape `(4,)`. Grown-ups call it a **vector**.\n",
    "- **a tray of baskets**, rows and columns, is 2-D: shape `(3, 4)`. A **matrix**.\n",
    "- **a stack of trays**, one per story, is 3-D: shape `(2, 3, 4)`.\n",
    "\n",
    "Everything in Momo's head is one of these: the words coming in (one row per token), the\n",
    "dials (big trays), the guesses going out (one row per position). Click through the\n",
    "dimensions below and hover any basket to see its address."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b5c0de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: 1-D, 2-D and 3-D tensors. Run this cell, then click and hover.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vt\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>\n",
    "#vt button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}\n",
    "#vt button.on{background:#6a4a1a;border-color:#ffb347}\n",
    "#vt .cell{display:inline-block;width:38px;height:30px;line-height:30px;text-align:center;font-size:12px;border:1px solid #5a8ab8;background:#2f4f6f;margin:1px;border-radius:3px;cursor:crosshair}\n",
    "#vt .cell:hover{background:#ffb347;color:#000}\n",
    "#vt .grid{display:inline-block;padding:6px;background:#262626;border-radius:6px;margin:6px;vertical-align:top}\n",
    "#vt .cap{font-size:13px;line-height:1.5;margin-top:10px;background:#262626;border-radius:8px;padding:10px}\n",
    "</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">A tensor is a grid of numbers with a shape. Click a dimension, then hover or tap a basket.</div>\n",
    "<div><button id=\"vtB1\" class=\"on\" onclick=\"vtShow(1)\">1-D: vector</button><button id=\"vtB2\" onclick=\"vtShow(2)\">2-D: matrix</button><button id=\"vtB3\" onclick=\"vtShow(3)\">3-D: batch</button></div>\n",
    "<div id=\"vtArea\" style=\"margin-top:10px;min-height:120px\"></div>\n",
    "<div id=\"vtHover\" style=\"font-size:12px;color:#ffb347;height:18px\"></div>\n",
    "<div class=\"cap\" id=\"vtCap\"></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const caps={\n",
    "1:\"<b>Shape (4,)</b>: one row of numbers. Picture: one row of Momo's baskets, one number in each. Or one word's meaning as 4 numbers. In our model each token becomes a vector of <b>256</b> numbers.\",\n",
    "2:\"<b>Shape (3, 4)</b>: rows and columns, like Momo's baskets laid out on a tray. Picture: one sentence, one row per token, one column per feature. Our model holds a 256-token story as a (256, 256) matrix.\",\n",
    "3:\"<b>Shape (2, 3, 4)</b>: a stack of trays. Picture: a batch of sentences processed at the same time, which is how GPUs stay busy. Our training batches are 32 windows of 256 tokens: the token ids are shape (32, 256), and after the meaning map each token has 256 numbers, so the stack becomes (32, 256, 256).\"};\n",
    "function grid(r,c,prefix,seed){seed=seed||0;let h='<div class=\\\"grid\\\">';for(let i=0;i<r;i++){h+='<div>';for(let j=0;j<c;j++){const v=(Math.round((Math.sin(i*3+j*7+seed*5)*5)*10)/10).toFixed(1);h+=`<span class=\\\"cell\\\" onmouseover=\\\"vtH('${prefix}${i},${j}', ${v})\\\" onclick=\\\"vtH('${prefix}${i},${j}', ${v})\\\" onmouseout=\\\"vtH('')\\\">${v}</span>`;}h+='</div>';}return h+'</div>';}\n",
    "window.vtH=function(idx,v){document.getElementById('vtHover').textContent=idx?`index [${idx}]  value ${v}`:'';};\n",
    "window.vtShow=function(d){[1,2,3].forEach(k=>document.getElementById('vtB'+k).classList.toggle('on',k===d));\n",
    " const a=document.getElementById('vtArea');\n",
    " if(d===1)a.innerHTML=grid(1,4,'');\n",
    " if(d===2)a.innerHTML=grid(3,4,'');\n",
    " if(d===3)a.innerHTML='<span style=\"font-size:12px;color:#aaa\">sentence 0</span>'+grid(3,4,'0,',1)+'<span style=\"font-size:12px;color:#aaa\">sentence 1</span>'+grid(3,4,'1,',2);\n",
    " document.getElementById('vtCap').innerHTML=caps[d];};\n",
    "window.vtShow(1);\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dfe6dad",
   "metadata": {},
   "source": [
    "### How matrix multiplication actually happens\n",
    "\n",
    "Almost all the work inside Momo is one operation: multiplying two trays of numbers together.\n",
    "Textbooks show the formula first, so nobody understands it. Here is the story instead.\n",
    "\n",
    "Momo and Ella each ate some fruit today. Momo's row says what he ate: 2 bananas, 1 mango,\n",
    "3 coconuts. Ella's row says what she ate. Ella also knows two facts about every fruit, kept as\n",
    "two columns: how many bites it takes, and how many seeds it has. How many bites did Momo take\n",
    "in total? Take Momo's row and the \"bites\" column, multiply the pairs, add them up. That one\n",
    "number is one cell of the answer. Do it for every row and every column and you have the whole\n",
    "answer tray: bites and seeds, for Momo and for Ella. That is all matrix multiplication is.\n",
    "Press **Next cell** and watch.\n",
    "\n",
    "Why Momo cares: a word's embedding is a row like \"2 bananas, 1 mango, 3 coconuts\", and a\n",
    "layer's dials are a tray of columns like \"bites\" and \"seeds\". Multiplying them asks \"how much\n",
    "does this word have of each thing this layer is looking for?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94b0555",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: matrix multiplication, one output cell at a time.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vm\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:960px\">\n",
    "<style>\n",
    "#vm button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}\n",
    "#vm table{border-collapse:separate;border-spacing:3px;display:inline-table;vertical-align:middle;margin:0 6px}\n",
    "#vm td,#vm th{width:46px;height:34px;text-align:center;font-size:13px;border-radius:5px}\n",
    "#vm td{background:#2f4f6f;border:1px solid #5a8ab8;transition:all .25s}\n",
    "#vm th{color:#aaa;font-weight:normal;font-size:11px;background:transparent}\n",
    "#vm td.row{background:#6a4a1a;border-color:#ffb347}\n",
    "#vm td.col{background:#2f6f4f;border-color:#6fcf97}\n",
    "#vm td.out{background:#ffb347;color:#000;font-weight:600}\n",
    "#vm td.empty{background:#262626;border:1px dashed #555;color:#666}\n",
    "#vm .sym{font-size:22px;color:#aaa;vertical-align:middle;margin:0 4px}\n",
    "#vm .cap{font-size:13px;line-height:1.55;margin-top:10px;background:#262626;border-radius:8px;padding:10px;min-height:60px}\n",
    "#vm .work{font-family:monospace;font-size:14px;color:#ffb347;margin-top:6px}\n",
    "</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:4px\">Matrix multiplication, one cell at a time</div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin-bottom:10px\">Story: Momo and Ella each ate some fruit. Ella knows how many bites each fruit takes and how many seeds it has. How many bites, and how many seeds, did each of them have?</div>\n",
    "<div id=\"vmArea\" style=\"overflow-x:auto;white-space:nowrap\"></div>\n",
    "<div class=\"work\" id=\"vmWork\"></div>\n",
    "<div class=\"cap\" id=\"vmCap\"></div>\n",
    "<div style=\"margin-top:10px\"><button onclick=\"vmStep()\">Next cell</button><button onclick=\"vmAuto()\">Play all</button><button onclick=\"vmReset()\">Reset</button></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const A=[[2,1,3],[1,4,0]];           // animals x fruits   (bananas, mangoes, coconuts)\n",
    "const B=[[2,1],[3,2],[4,3]];         // fruits x facts      (bites per fruit, seeds per fruit)\n",
    "const fruits=[\"bananas\",\"mangoes\",\"coconuts\"], one=[\"banana\",\"mango\",\"coconut\"], custs=[\"Momo\",\"Ella\"], stalls=[\"bites\",\"seeds\"];\n",
    "const C=[[null,null],[null,null]]; let k=-1, timer=null;\n",
    "function tbl(M,rh,ch,cls){let h='<table><tr><th></th>'+ch.map(c=>`<th>${c}</th>`).join('')+'</tr>';M.forEach((r,i)=>{h+=`<tr><th>${rh[i]}</th>`+r.map((v,j)=>`<td class=\"${cls(i,j,v)}\">${v===null?'?':v}</td>`).join('')+'</tr>';});return h+'</table>';}\n",
    "function render(){const i=k<0?-1:Math.floor(k/2), j=k<0?-1:k%2;\n",
    " document.getElementById('vmArea').innerHTML=\n",
    "  tbl(A,custs,fruits,(r,c)=>r===i?'row':'')+'<span class=\"sym\">&times;</span>'+\n",
    "  tbl(B,fruits,stalls,(r,c)=>c===j?'col':'')+'<span class=\"sym\">=</span>'+\n",
    "  tbl(C,custs,stalls,(r,c,v)=>(r===i&&c===j)?'out':(v===null?'empty':''));\n",
    " if(k<0){document.getElementById('vmWork').textContent='';document.getElementById('vmCap').innerHTML='<b>Left table:</b> how many of each fruit each animal ate (rows = animals, columns = fruits).<br><b>Right table:</b> Ella\\'s two facts about each fruit: bites it takes, seeds it has (rows = fruits, columns = facts).<br>Press <b>Next cell</b>. Each answer cell = one <span style=\"color:#ffb347\">row</span> from the left paired with one <span style=\"color:#6fcf97\">column</span> from the right: multiply matching pairs, add them up.';return;}\n",
    " const terms=A[i].map((a,t)=>`${a}&times;${B[t][j]}`), prods=A[i].map((a,t)=>a*B[t][j]), tot=prods.reduce((x,y)=>x+y,0); C[i][j]=tot;\n",
    " document.getElementById('vmArea').innerHTML=tbl(A,custs,fruits,(r,c)=>r===i?'row':'')+'<span class=\"sym\">&times;</span>'+tbl(B,fruits,stalls,(r,c)=>c===j?'col':'')+'<span class=\"sym\">=</span>'+tbl(C,custs,stalls,(r,c,v)=>(r===i&&c===j)?'out':(v===null?'empty':''));\n",
    " document.getElementById('vmWork').innerHTML=`${custs[i]}, ${stalls[j]}:  ${terms.join(' + ')}  =  ${prods.join(' + ')}  =  <b>${tot}</b>`;\n",
    " document.getElementById('vmCap').innerHTML=`Take <span style=\"color:#ffb347\"><b>${custs[i]}'s row</b></span> (what ${custs[i]} ate) and the <span style=\"color:#6fcf97\"><b>${stalls[j]} column</b></span> (${stalls[j]} per fruit). Walk along both at the same time: ${A[i].map((a,t)=>`${a} ${a===1?one[t]:fruits[t]} &times; ${B[t][j]} ${stalls[j]} each`).join(', ')}. Add up: <b>${tot} ${stalls[j]}</b> for ${custs[i]}. That is one cell of the answer.<br><br>Why it matters: a neural-network layer is exactly this. Left table = your tokens (one row per token, one column per feature). Right table = the layer's weights. The answer = the tokens' new features. GPUs exist to do this billions of times per second.`;}\n",
    "window.vmStep=function(){if(k>=3){return;}k++;render();};\n",
    "window.vmAuto=function(){if(timer)return;timer=setInterval(()=>{if(k>=3){clearInterval(timer);timer=null;return;}vmStep();},1400);};\n",
    "window.vmReset=function(){clearInterval(timer);timer=null;k=-1;C[0]=[null,null];C[1]=[null,null];render();};\n",
    "render();\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e39e048b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([[1., 2., 3.],\n",
    "                  [4., 5., 6.]])          # shape (2, 3): 2 rows, 3 columns\n",
    "print(\"shape:\", x.shape)\n",
    "print(\"x @ x.T =\\n\", x @ x.T)            # matrix multiply: the core operation of every layer\n",
    "print(\"x * 2 =\\n\", x * 2)                # elementwise ops just work"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3febc443",
   "metadata": {},
   "source": [
    "### The gradient: which way to turn each dial\n",
    "\n",
    "A **gradient** answers one question for one dial: *if I turn you up a tiny bit, does the\n",
    "mistake get bigger or smaller, and how fast?* Every dial gets its own answer. Together those\n",
    "answers are the slope under Momo's feet.\n",
    "\n",
    "PyTorch finds them with **autograd**: we mark a tensor with `requires_grad=True` (switch the\n",
    "tape recorder on), do some math, call `.backward()` (read the tape backwards), and each\n",
    "dial's answer appears in `.grad`.\n",
    "\n",
    "Training is only this, repeated thousands of times: measure the mistake, get the gradients,\n",
    "turn every dial a small step the other way.\n",
    "\n",
    "Start with one dial. Momo's coconut sits on a bowl-shaped hill where height is the mistake.\n",
    "Drag the dial or press **Roll** and watch it find the bottom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e72fa2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: one dial. Momo's coconut rolls to the bottom of the bowl.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vg\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>#vg button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}#vg .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Momo's coconut rolling to the bottom of the bowl. Drag the dial, or let it roll.</div>\n",
    "<svg id=\"vgSvg\" width=\"560\" height=\"260\" style=\"background:#262626;border-radius:8px\"></svg>\n",
    "<div style=\"margin-top:8px\"><label style=\"font-size:13px\">weight w <input id=\"vgW\" type=\"range\" min=\"-4\" max=\"8\" step=\"any\" value=\"7\" style=\"width:260px;vertical-align:middle\" oninput=\"vgDraw(+this.value)\"></label>\n",
    "<button onclick=\"vgStep()\">Take one step</button><button onclick=\"vgAuto()\">Roll (auto)</button></div>\n",
    "<div class=\"cap\" id=\"vgCap\"></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const svg=document.getElementById('vgSvg'),W=560,H=260;const lossOf=w=>0.5*(w-2)*(w-2)+1, gradOf=w=>(w-2);\n",
    "const X=w=>40+(w+4)/12*(W-80), Y=l=>H-30-(l/19)*(H-60);\n",
    "let timer=null;\n",
    "window.vgDraw=function(w){document.getElementById('vgW').value=w;let d='';for(let x=-4;x<=8;x+=0.1){d+=(d?'L':'M')+X(x).toFixed(1)+','+Y(lossOf(x)).toFixed(1);}\n",
    " const l=lossOf(w),g=gradOf(w),lr=0.3,wn=w-lr*g;\n",
    " svg.innerHTML=`<path d=\"${d}\" fill=\"none\" stroke=\"#5a8ab8\" stroke-width=\"3\"/>\n",
    " <text x=\"${W/2}\" y=\"${H-6}\" fill=\"#aaa\" font-size=\"12\" text-anchor=\"middle\">dial setting w</text>\n",
    " <text x=\"14\" y=\"${H/2}\" fill=\"#aaa\" font-size=\"12\" transform=\"rotate(-90 14,${H/2})\" text-anchor=\"middle\">loss (how wrong)</text>\n",
    " ${Math.abs(w-2)<0.05?'':`<line x1=\"${X(w)}\" y1=\"${Y(l)}\" x2=\"${X(wn)}\" y2=\"${Y(l)}\" stroke=\"#ffb347\" stroke-width=\"3\" marker-end=\"url(#ar)\"/>`}\n",
    " <defs><marker id=\"ar\" markerWidth=\"8\" markerHeight=\"8\" refX=\"6\" refY=\"4\" orient=\"auto\"><path d=\"M0,0 L8,4 L0,8 z\" fill=\"#ffb347\"/></marker></defs>\n",
    " <circle cx=\"${X(w)}\" cy=\"${Y(l)}\" r=\"9\" fill=\"#ffb347\" stroke=\"#fff\" stroke-width=\"2\"/>\n",
    " <text x=\"${X(2)}\" y=\"${Y(1)+22}\" fill=\"#6fcf97\" font-size=\"12\" text-anchor=\"middle\">best setting (smallest mistake)</text>`;\n",
    " document.getElementById('vgCap').innerHTML=`w = <b>${w.toFixed(2)}</b> &nbsp; loss = <b>${l.toFixed(2)}</b> &nbsp; gradient = <b>${g.toFixed(2)}</b><br>The gradient is the slope under the coconut. Positive slope means \"loss goes up if w grows\", so we move w <i>down</i> by learning rate &times; gradient: ${w.toFixed(2)} &minus; 0.3 &times; ${g.toFixed(2)} = <b>${wn.toFixed(2)}</b>.<br>${Math.abs(w-2)<0.05?'<b>At the bottom.</b> The slope is zero here, so there is nothing left to fix: this is the best setting for this one dial. ':''}Picture Momo in fog feeling which way the ground slopes and stepping downhill. Training does this for all 5.9 million dials at once, thousands of times.`;\n",
    " return wn;};\n",
    "window.vgStep=function(){const w=+document.getElementById('vgW').value;const wn=vgDraw(w);vgDraw(Math.abs(wn-2)<0.03?2:wn);};\n",
    "window.vgAuto=function(){if(timer){clearInterval(timer);timer=null;return;}timer=setInterval(()=>{const w=+document.getElementById('vgW').value;if(Math.abs(w-2)<0.03){clearInterval(timer);timer=null;return;}vgStep();},350);};\n",
    "vgDraw(7);\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8c2b93f",
   "metadata": {},
   "source": [
    "### Momo in the foggy jungle: thorn bushes, marks, and the slope\n",
    "\n",
    "The coconut was a toy version. Now Momo himself is the walker, because his position *is*\n",
    "his dials. With millions of dials the hill is not a bowl but a whole jungle valley, and Momo\n",
    "is in fog. Here is the same idea as a story, because this is the part people picture wrongly.\n",
    "\n",
    "**The hill is the mistake.** Height is \"how wrong Momo is right now\". The grove at the\n",
    "bottom is the set of dial settings that makes the fewest mistakes. Momo's position *is* his\n",
    "dials. One step is one dial turn.\n",
    "\n",
    "**A thorn bush is a step that made things worse.** Momo tried a direction and the mistake\n",
    "went up. In the widget he steps back and marks the spot with a red X. Real training is blunter: it\n",
    "never steps back. Momo stays where he landed, reads the slope again, and the slope now points\n",
    "back downhill. The correction comes from the next reading, not from undoing the step.\n",
    "\n",
    "**What the gradient adds.** A real monkey learns a thorn bush is there by walking into it.\n",
    "The gradient is better than that. It is the slope under the feet, felt *before* stepping.\n",
    "Momo gets the slope for all 5.9 million dials in one go, without trying any of them (the\n",
    "blame game, right after this widget, shows how). Most of the time he never touches a bush.\n",
    "Guessing instead would mean testing 5.9 million directions one at a time.\n",
    "\n",
    "**So why does training still sometimes go uphill?**\n",
    "- The stride was too long. That is the **learning rate**. The slope was right, but Momo\n",
    "  overshot the grove and landed on the far wall. The next slope reading points back.\n",
    "- The reading was noisy. Each step uses one **batch** (a handful of story pieces, not the\n",
    "  whole library), so the slope is an estimate.\n",
    "\n",
    "**The marks on the map are real, in two forms.**\n",
    "- *The dials themselves.* Every step is written into the dials. There is no separate map;\n",
    "  after training Momo *is* the map. (Later lessons can partly overwrite earlier ones, which\n",
    "  is why the stories are shuffled and why the last steps are tiny.)\n",
    "- *The optimizer's memory.* The dial-turning rule we use, **AdamW**, keeps a running average\n",
    "  of recent slopes. That is Momo remembering \"the last ten readings all said left, so one odd\n",
    "  reading saying right is probably a bump; keep going left\". One bush does not fool him.\n",
    "\n",
    "**Fastest and most accurate.** Short careful steps at the start (**warm-up**, the ground is\n",
    "unknown), full stride in the middle, tiny steps near the grove (**cosine decay**) so Momo\n",
    "stops on the floor instead of bouncing between the walls.\n",
    "\n",
    "Try both ways of walking below and compare the step and bush counts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c8b20dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: Momo in the foggy jungle. Guess and hit thorn bushes, or read the slope first.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vh\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:960px\">\n",
    "<style>\n",
    "#vh button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin:0 6px 6px 0;cursor:pointer;font-size:13px}\n",
    "#vh button:hover{background:#444}\n",
    "#vh button.g{border-color:#6fcf97}#vh button.r{border-color:#ff6b6b}\n",
    "#vh .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}\n",
    "#vh .log{font-size:12px;background:#262626;border-radius:8px;padding:8px 10px;height:150px;overflow:auto;line-height:1.5}\n",
    "#vh .stat{display:inline-block;padding:4px 10px;border-radius:6px;background:#2a2a2a;border:1px solid #555;margin:2px 6px 2px 0;font-size:12px}\n",
    "</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:4px\">Momo in the foggy jungle: thorn bushes, marks, and the slope under his feet</div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin-bottom:8px\">The map is the mistake score (the loss). Dark blue is low (good), yellow is high (bad). The flag is the banana grove: the best set of dials. Momo cannot see the map; he only feels the ground under his feet.</div>\n",
    "<div style=\"display:flex;gap:16px;flex-wrap:wrap\">\n",
    "  <div><canvas id=\"vhC\" width=\"520\" height=\"360\" style=\"border-radius:8px;background:#262626\"></canvas></div>\n",
    "  <div style=\"flex:1;min-width:300px\">\n",
    "    <div style=\"font-size:12px;color:#aaa;margin-bottom:4px\">TWO WAYS TO WALK</div>\n",
    "    <button class=\"r\" onclick=\"vhBlind()\">Guess a step (may hit a thorn bush)</button>\n",
    "    <button class=\"g\" onclick=\"vhGrad()\">Read the slope, then step</button><br>\n",
    "    <button onclick=\"vhAuto('blind')\">Auto: guessing</button>\n",
    "    <button onclick=\"vhAuto('grad')\">Auto: slope</button>\n",
    "    <button onclick=\"vhReset()\">Reset</button>\n",
    "    <div style=\"margin:6px 0\">\n",
    "      <label style=\"font-size:12px\">Stride (learning rate) <input id=\"vhLr\" type=\"range\" min=\"1\" max=\"20\" value=\"4\" style=\"vertical-align:middle;width:120px\"></label>\n",
    "      <label style=\"font-size:12px;margin-left:8px\"><input id=\"vhMem\" type=\"checkbox\" checked> memory of recent slopes (momentum)</label>\n",
    "    </div>\n",
    "    <div><span class=\"stat\" id=\"vhSteps\">steps: 0</span><span class=\"stat\" id=\"vhTrees\">thorn bushes: 0</span><span class=\"stat\" id=\"vhLoss\">loss: 0</span></div>\n",
    "    <div style=\"font-size:12px;color:#aaa;margin:8px 0 4px\">MOMO'S NOTEBOOK (marks on the map)</div>\n",
    "    <div class=\"log\" id=\"vhLog\">Nothing tried yet.</div>\n",
    "  </div>\n",
    "</div>\n",
    "<div class=\"cap\" id=\"vhCap\"></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const C=document.getElementById('vhC'),ctx=C.getContext('2d'),W=C.width,H=C.height;\n",
    "// loss landscape in map coordinates (0..1, 0..1): a tilted bowl with two bumps (the \"trees\" of the terrain)\n",
    "const goal=[0.72,0.68];\n",
    "const loss=(x,y)=>{const dx=x-goal[0],dy=y-goal[1];let l=6*(dx*dx+dy*dy);\n",
    "  l+=1.1*Math.exp(-((x-0.40)**2+(y-0.55)**2)/0.012);l+=0.9*Math.exp(-((x-0.58)**2+(y-0.30)**2)/0.010);return l;};\n",
    "const grad=(x,y)=>{const h=1e-3;return [(loss(x+h,y)-loss(x-h,y))/(2*h),(loss(x,y+h)-loss(x,y-h))/(2*h)];};\n",
    "const px=x=>x*W, py=y=>(1-y)*H;\n",
    "let bg=null,pos,trail,xs,steps,trees,mem,timer=null,mode='';\n",
    "let minL=Infinity;for(let i=0;i<=100;i++)for(let j=0;j<=100;j++){const v=loss(i/100,j/100);if(v<minL)minL=v;}\n",
    "function paintBg(){const img=ctx.createImageData(W,H);let mx=0;const vals=new Float32Array(W*H);\n",
    "  for(let j=0;j<H;j++)for(let i=0;i<W;i++){const v=loss(i/W,1-j/H);vals[j*W+i]=v;if(v>mx)mx=v;}\n",
    "  for(let k=0;k<W*H;k++){const t=Math.min(1,vals[k]/mx);const r=Math.round(30+225*t),g=Math.round(40+170*t*t),b=Math.round(110*(1-t)+40);\n",
    "    img.data[k*4]=r;img.data[k*4+1]=g;img.data[k*4+2]=b;img.data[k*4+3]=255;}\n",
    "  ctx.putImageData(img,0,0);\n",
    "  // contour-ish rings\n",
    "  ctx.strokeStyle='rgba(255,255,255,0.12)';ctx.lineWidth=1;\n",
    "  for(let lv=0.3;lv<4;lv+=0.35){ctx.beginPath();let first=true;for(let a=0;a<=Math.PI*2+0.05;a+=0.05){\n",
    "    // crude: walk outward from goal until loss>=lv\n",
    "    let r=0;for(;r<1.2;r+=0.004){if(loss(goal[0]+r*Math.cos(a),goal[1]+r*Math.sin(a))>=lv)break;}\n",
    "    const x=px(goal[0]+r*Math.cos(a)),y=py(goal[1]+r*Math.sin(a));if(first){ctx.moveTo(x,y);first=false;}else ctx.lineTo(x,y);}ctx.stroke();}\n",
    "  bg=ctx.getImageData(0,0,W,H);}\n",
    "function draw(){ctx.putImageData(bg,0,0);\n",
    "  // flag\n",
    "  ctx.fillStyle='#fff';ctx.fillRect(px(goal[0])-1,py(goal[1])-22,2,22);ctx.fillStyle='#6fcf97';ctx.beginPath();ctx.moveTo(px(goal[0]),py(goal[1])-22);ctx.lineTo(px(goal[0])+14,py(goal[1])-17);ctx.lineTo(px(goal[0]),py(goal[1])-12);ctx.fill();\n",
    "  ctx.font='11px sans-serif';ctx.fillStyle='#fff';ctx.fillText('banana grove (best dials)',px(goal[0])-128,py(goal[1])-26);\n",
    "  // trail\n",
    "  ctx.strokeStyle='#6fcf97';ctx.lineWidth=2;ctx.beginPath();trail.forEach((p,i)=>{i?ctx.lineTo(px(p[0]),py(p[1])):ctx.moveTo(px(p[0]),py(p[1]));});ctx.stroke();\n",
    "  trail.forEach(p=>{ctx.fillStyle='#6fcf97';ctx.beginPath();ctx.arc(px(p[0]),py(p[1]),3,0,7);ctx.fill();});\n",
    "  // trees (bad tries)\n",
    "  xs.forEach(p=>{ctx.strokeStyle='#ff6b6b';ctx.lineWidth=2.5;const x=px(p[0]),y=py(p[1]);ctx.beginPath();ctx.moveTo(x-6,y-6);ctx.lineTo(x+6,y+6);ctx.moveTo(x+6,y-6);ctx.lineTo(x-6,y+6);ctx.stroke();});\n",
    "  // slope arrow (what the gradient says)\n",
    "  const arrived=loss(pos[0],pos[1])<minL+0.015;\n",
    "  const g=grad(pos[0],pos[1]);const n=Math.hypot(g[0],g[1])||1;const ax=-g[0]/n*34,ay=g[1]/n*34;\n",
    "  if(!arrived){ctx.strokeStyle='#ffb347';ctx.lineWidth=3;ctx.beginPath();ctx.moveTo(px(pos[0]),py(pos[1]));ctx.lineTo(px(pos[0])+ax,py(pos[1])+ay);ctx.stroke();\n",
    "  ctx.fillStyle='#ffb347';ctx.beginPath();const ang=Math.atan2(ay,ax);ctx.moveTo(px(pos[0])+ax,py(pos[1])+ay);ctx.lineTo(px(pos[0])+ax-9*Math.cos(ang-0.4),py(pos[1])+ay-9*Math.sin(ang-0.4));ctx.lineTo(px(pos[0])+ax-9*Math.cos(ang+0.4),py(pos[1])+ay-9*Math.sin(ang+0.4));ctx.fill();}\n",
    "  // hiker\n",
    "  ctx.fillStyle='#fff';ctx.beginPath();ctx.arc(px(pos[0]),py(pos[1]),8,0,7);ctx.fill();ctx.fillStyle='#1e1e1e';ctx.font='11px sans-serif';ctx.fillText('\\u{1F435}',px(pos[0])-6,py(pos[1])+4);\n",
    "  const l=loss(pos[0],pos[1]);\n",
    "  document.getElementById('vhSteps').textContent='steps: '+steps;document.getElementById('vhTrees').textContent='thorn bushes: '+trees;document.getElementById('vhLoss').textContent='loss: '+l.toFixed(2);\n",
    "  if(l<minL+0.015){document.getElementById('vhCap').innerHTML='<b>Arrived.</b> Momo is in the banana grove: these dial settings make the fewest mistakes. Steps: <b>'+steps+'</b>, thorn bushes: <b>'+trees+'</b>. Press Reset and try the other way of walking.';if(timer){clearInterval(timer);timer=null;}}\n",
    "}\n",
    "function logLine(html){const L=document.getElementById('vhLog');if(L.textContent.startsWith('Nothing'))L.innerHTML='';L.innerHTML='<div>'+html+'</div>'+L.innerHTML;}\n",
    "const lr=()=>0.01*(+document.getElementById('vhLr').value);\n",
    "window.vhBlind=function(){steps++;const a=Math.random()*Math.PI*2;const d=lr()*2.5;const nx=Math.min(0.98,Math.max(0.02,pos[0]+Math.cos(a)*d)),ny=Math.min(0.98,Math.max(0.02,pos[1]+Math.sin(a)*d));\n",
    "  const before=loss(pos[0],pos[1]),after=loss(nx,ny);\n",
    "  if(after>=before){trees++;xs.push([nx,ny]);logLine('&#10060; tried '+dir(a)+': loss went <b>up</b> '+before.toFixed(2)+' &rarr; '+after.toFixed(2)+'. Thorn bush. Marked it, stepped back.');\n",
    "    document.getElementById('vhCap').innerHTML='<b>Thorn bush.</b> The guess made the mistake worse, so Momo steps back and marks the spot with a red X. Guessing finds the grove eventually, but it pays for every bush. Real training cannot afford this: Momo has 5.9 million dials, and guessing would mean trying each one, one direction at a time.';}\n",
    "  else{pos=[nx,ny];trail.push(pos.slice());logLine('&#10004; tried '+dir(a)+': loss went down '+before.toFixed(2)+' &rarr; '+after.toFixed(2)+'. Kept it.');\n",
    "    document.getElementById('vhCap').innerHTML='<b>Lucky guess.</b> The mistake went down, so the step is kept and the green trail records it.';}\n",
    "  draw();};\n",
    "window.vhGrad=function(){steps++;let g=grad(pos[0],pos[1]);\n",
    "  // momentum: remember recent slopes\n",
    "  if(document.getElementById('vhMem').checked){mem=[0.8*mem[0]+0.2*g[0],0.8*mem[1]+0.2*g[1]];g=mem;}\n",
    "  const n=Math.hypot(g[0],g[1])||1;const d=Math.min(lr()*2.5,lr()*n);const nx=Math.min(0.98,Math.max(0.02,pos[0]-g[0]/n*d)),ny=Math.min(0.98,Math.max(0.02,pos[1]-g[1]/n*d));\n",
    "  const before=loss(pos[0],pos[1]),after=loss(nx,ny);pos=[nx,ny];trail.push(pos.slice());\n",
    "  if(after>=before){trees++;xs.push([nx,ny]);logLine('&#9888; slope said '+dir(Math.atan2(-g[1],-g[0]))+' but loss went <b>up</b> '+before.toFixed(2)+' &rarr; '+after.toFixed(2)+' (stride too long, or noisy reading). Next reading will point back.');\n",
    "    document.getElementById('vhCap').innerHTML='<b>Overshot.</b> The slope was right but the stride was too long, so Momo landed on the far wall. Nothing is lost: the next slope reading points straight back. Shorten the stride (a lower learning rate) and this stops happening. That is why training ends with tiny steps.';}\n",
    "  else{logLine('&#10004; slope said '+dir(Math.atan2(-g[1],-g[0]))+': loss '+before.toFixed(2)+' &rarr; '+after.toFixed(2)+'.');\n",
    "    document.getElementById('vhCap').innerHTML='<b>Read the slope first.</b> The orange arrow is the gradient: which way is downhill, felt through the feet before moving. No bush needed. The blame game (next widget) is how Momo gets this arrow for all 5.9 million dials at once. With memory on, the arrow is the average of recent readings, so one odd bump does not fool Momo.';}\n",
    "  draw();};\n",
    "window.vhAuto=function(m){if(timer){clearInterval(timer);timer=null;if(mode===m)return;}if(mode!==m)mem=[0,0];mode=m;timer=setInterval(()=>{m==='blind'?vhBlind():vhGrad();if(steps>400){clearInterval(timer);timer=null;}},m==='blind'?120:260);};\n",
    "window.vhReset=function(){if(timer){clearInterval(timer);timer=null;}pos=[0.12,0.86];trail=[pos.slice()];xs=[];steps=0;trees=0;mem=[0,0];document.getElementById('vhLog').textContent='Nothing tried yet.';\n",
    "  document.getElementById('vhCap').innerHTML='Momo starts top-left, in fog. Press the <b>red</b> button to guess a step. Press the <b>green</b> button to feel the slope first. Which way reaches the banana grove with fewer steps and fewer thorn bushes? Then push the stride slider to the right and watch what a too-long stride does.';draw();};\n",
    "function dir(a){const d=['east','north-east','north','north-west','west','south-west','south','south-east'];return d[Math.round(((a%(2*Math.PI))+2*Math.PI)%(2*Math.PI)/(Math.PI/4))%8];}\n",
    "paintBg();vhReset();\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea803726",
   "metadata": {},
   "source": [
    "### How does Momo feel the slope without walking? The blame game\n",
    "\n",
    "This is the part most explanations skip, and it is the real trick inside every AI.\n",
    "\n",
    "Momo has a job: pick bananas for Ella. Ella wants 10. Momo's rule has two dials:\n",
    "*bananas = dial A × bunch size + dial B*. Today the bunch size is 2. His dials say A = 1 and\n",
    "B = 3, so he picks 1 × 2 + 3 = 5 bananas. Ella wanted 10. The gap is 5. We score the mistake\n",
    "as the gap times itself, 5 × 5 = 25 (that is what \"squared\" means), so a big gap hurts much\n",
    "more than a small one.\n",
    "\n",
    "Which way should Momo turn dial A? And dial B? He could try: turn A a little, pick again,\n",
    "see if the mistake shrank. That works, but he must try every dial separately. The real Momo\n",
    "has 5.9 million dials. Far too slow.\n",
    "\n",
    "Instead Momo plays the **blame game**. While picking he wrote every step on a tape:\n",
    "\"multiplied A by 2\", \"added B\", \"squared the gap\". Now he reads the tape **backwards** and\n",
    "asks each step: *\"if you went up by one tiny bit, what would the mistake do?\"* Each step\n",
    "answers by taking the answer from the step after it and multiplying by how hard it pushes.\n",
    "At the start of the tape every dial has a number: \"turn me this way, this much\". Those\n",
    "numbers are the gradient, found for all dials in one backward trip.\n",
    "\n",
    "Grown-up words: the tape is **autograd**, reading it backwards is **backpropagation**, and\n",
    "\"multiply by how hard I push\" at each step is the **chain rule**. One detail the game leaves\n",
    "out: when one step feeds *two* later steps (a trick called a residual connection, which we meet\n",
    "in Part 6), the two blames are added together. And the game's mistake is a squared gap; the\n",
    "real Momo scores his mistake differently (cross-entropy, explained in Part 7). Same game,\n",
    "different first step. Press each stage below. The slow-way check lets you see that the\n",
    "backward trip tells the truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc8e8406",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: the blame game. Pick, pass the blame back, turn the dials.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"va\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:960px\">\n",
    "<style>\n",
    "#va button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin:0 6px 6px 0;cursor:pointer;font-size:13px}\n",
    "#va button:hover{background:#444}\n",
    "#va button.on{border-color:#ffb347;background:#3a2e1a}\n",
    "#va .box{display:inline-block;vertical-align:top;text-align:center;margin:0 2px;min-width:78px}\n",
    "#va .node{border:2px solid #555;border-radius:10px;background:#2a2a2a;padding:8px 6px;font-size:12px;min-height:54px}\n",
    "#va .node.dial{border-color:#5a8ab8}\n",
    "#va .node.lit{border-color:#5a8ab8;box-shadow:0 0 10px #5a8ab888}\n",
    "#va .node.blame{border-color:#ff6b6b;box-shadow:0 0 10px #ff6b6b88}\n",
    "#va .fw{color:#8ec6ff;font-weight:700;font-size:15px;min-height:20px}\n",
    "#va .bw{color:#ff8a8a;font-weight:700;font-size:13px;min-height:18px}\n",
    "#va .op{display:inline-block;vertical-align:top;margin-top:38px;color:#aaa;font-size:18px;width:20px;text-align:center}\n",
    "#va .cap{font-size:13px;line-height:1.55;margin-top:10px;background:#262626;border-radius:8px;padding:10px}\n",
    "#va .knob{width:36px;height:36px;border-radius:50%;border:2px solid #8ec6ff;margin:4px auto;position:relative;background:#1e1e1e}\n",
    "#va .knob i{position:absolute;left:50%;top:2px;width:2px;height:14px;background:#ffb347;margin-left:-1px;transform-origin:1px 16px}\n",
    "#va .lbl{font-size:11px;color:#aaa;text-transform:uppercase;letter-spacing:.05em}\n",
    "</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:4px\">The blame game: how Momo feels the slope without walking (this is autograd)</div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin-bottom:10px\">Momo picks bananas for Ella using two dials in his head. Ella wants <b>10</b> bananas. Blue numbers flow forward while Momo picks. Red numbers are the blame flowing backward.</div>\n",
    "\n",
    "<div id=\"vaRow\" style=\"display:flex;flex-wrap:wrap;align-items:flex-start;padding:4px 0\">\n",
    "  <div class=\"box\"><div class=\"lbl\">bunch size</div><div class=\"node\" id=\"nx\"><div style=\"font-size:11px;color:#aaa\">today</div><div class=\"fw\" id=\"fx\">2</div></div><div class=\"bw\"></div></div>\n",
    "  <div class=\"op\">&times;</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">dial A</div><div class=\"node dial\" id=\"nA\"><div class=\"knob\"><i id=\"kA\"></i></div><div class=\"fw\" id=\"fA\"></div></div><div class=\"bw\" id=\"bA\"></div></div>\n",
    "  <div class=\"op\">=</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">A &times; bunch</div><div class=\"node\" id=\"nm\"><div style=\"font-size:11px;color:#aaa\">step 1</div><div class=\"fw\" id=\"fm\"></div></div><div class=\"bw\" id=\"bm\"></div></div>\n",
    "  <div class=\"op\">+</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">dial B</div><div class=\"node dial\" id=\"nB\"><div class=\"knob\"><i id=\"kB\"></i></div><div class=\"fw\" id=\"fB\"></div></div><div class=\"bw\" id=\"bB\"></div></div>\n",
    "  <div class=\"op\">=</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">Momo's guess</div><div class=\"node\" id=\"ng\"><div style=\"font-size:11px;color:#aaa\">bananas picked</div><div class=\"fw\" id=\"fg\"></div></div><div class=\"bw\" id=\"bg\"></div></div>\n",
    "  <div class=\"op\">vs</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">Ella wants</div><div class=\"node\" id=\"nt\"><div style=\"font-size:11px;color:#aaa\">target</div><div class=\"fw\">10</div></div><div class=\"bw\"></div></div>\n",
    "  <div class=\"op\">&rarr;</div>\n",
    "  <div class=\"box\"><div class=\"lbl\">mistake (loss)</div><div class=\"node\" id=\"nL\"><div style=\"font-size:11px;color:#aaa\">(guess &minus; 10)&sup2;</div><div class=\"fw\" id=\"fL\"></div></div><div class=\"bw\" id=\"bL\"></div></div>\n",
    "</div>\n",
    "\n",
    "<div style=\"margin-top:8px\">\n",
    "  <button id=\"vaB1\" onclick=\"vaCook()\">1. Pick (forward)</button>\n",
    "  <button id=\"vaB2\" onclick=\"vaBlame()\">2. Pass the blame back (backward)</button>\n",
    "  <button id=\"vaB3\" onclick=\"vaTurn()\">3. Turn the dials (one step)</button>\n",
    "  <button onclick=\"vaWiggle()\">Check it the slow way (wiggle dial A)</button>\n",
    "  <button onclick=\"vaAuto()\">Auto: pick, blame, turn</button>\n",
    "  <button onclick=\"vaReset()\">Reset</button>\n",
    "  <span style=\"font-size:12px;color:#aaa;margin-left:6px\">round <b id=\"vaRound\">0</b></span>\n",
    "</div>\n",
    "<div class=\"cap\" id=\"vaCap\"></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const x=2,T=10,lr=0.02;let A,B,m,g,L,gA,gB,gg,phase,round,timer=null,bwT=[];\n",
    "const $=id=>document.getElementById(id);\n",
    "const f2=v=>(Math.round(v*100)/100).toFixed(2);\n",
    "function knobs(){$('kA').style.transform='rotate('+(A*60-60)+'deg)';$('kB').style.transform='rotate('+(B*40-120)+'deg)';$('fA').textContent=f2(A);$('fB').textContent=f2(B);}\n",
    "function clearFw(){['fm','fg','fL'].forEach(i=>$(i).textContent='');}\n",
    "function clearBw(){bwT.forEach(clearTimeout);bwT=[];['bA','bm','bB','bg','bL'].forEach(i=>$(i).innerHTML='');document.querySelectorAll('#va .node').forEach(n=>n.classList.remove('blame','lit'));}\n",
    "function lit(ids){document.querySelectorAll('#va .node').forEach(n=>n.classList.remove('lit'));ids.forEach(i=>$(i).classList.add('lit'));}\n",
    "function btn(n){[1,2,3].forEach(k=>$('vaB'+k).classList.toggle('on',k===n));}\n",
    "window.vaCook=function(){m=A*x;g=m+B;L=(g-T)*(g-T);clearBw();$('fm').textContent=f2(m);$('fg').textContent=f2(g);$('fL').textContent=f2(L);lit(['nx','nA','nm','nB','ng','nt','nL']);phase=1;btn(2);\n",
    "  const d=g-T;\n",
    "  $('vaCap').innerHTML=`<b>Pick.</b> Momo follows his rule: dial A &times; bunch size + dial B = ${f2(A)} &times; 2 + ${f2(B)} = <b>${f2(g)}</b> bananas. Ella wanted 10, so the guess is ${Math.abs(d).toFixed(2)} too ${d<0?'low':'high'}. Mistake = (${f2(g)} &minus; 10)&sup2; = <b>${f2(L)}</b>.<br>While picking, Momo wrote every step on a tape: \"I multiplied A by 2\", \"I added B\", \"I squared the gap\". PyTorch does exactly this; the tape is what <code>requires_grad=True</code> switches on.`;};\n",
    "window.vaBlame=function(){if(phase<1){vaCook();}\n",
    "  const d=g-T;gg=2*d;gB=gg*1;gA=gg*x;phase=2;btn(3);\n",
    "  document.querySelectorAll('#va .node').forEach(n=>n.classList.remove('lit'));\n",
    "  const steps=[\n",
    "   ()=>{$('nL').classList.add('blame');$('bL').innerHTML='start: 1';},\n",
    "   ()=>{$('ng').classList.add('blame');$('bg').innerHTML='blame '+f2(gg);},\n",
    "   ()=>{$('nB').classList.add('blame');$('bB').innerHTML='blame '+f2(gB);$('nm').classList.add('blame');$('bm').innerHTML='blame '+f2(gg);},\n",
    "   ()=>{$('nA').classList.add('blame');$('bA').innerHTML='blame '+f2(gA);},\n",
    "  ];\n",
    "  const gap=timer?150:450;bwT=steps.map((s,i)=>setTimeout(s,i*gap));\n",
    "  $('vaCap').innerHTML=`<b>Pass the blame back.</b> Momo reads the tape <i>backwards</i> and asks every step the same question: <i>\"if you went up by one tiny bit, what would the mistake do?\"</i><br>\n",
    "  &bull; <b>Mistake &larr; guess.</b> The gap is ${f2(d)}. Raising the guess by 1 changes (gap)&sup2; by about 2 &times; gap = <b>${f2(gg)}</b>. ${d<0?'Minus means: a higher guess makes the mistake <i>smaller</i>.':'Plus means: a higher guess makes the mistake <i>bigger</i>.'}<br>\n",
    "  &bull; <b>Guess &larr; dial B.</b> Raise B by 1 and the guess rises by exactly 1. So B's blame = ${f2(gg)} &times; 1 = <b>${f2(gB)}</b>.<br>\n",
    "  &bull; <b>Guess &larr; dial A.</b> Raise A by 1 and the guess rises by 2, because A is multiplied by the bunch size. So A's blame = ${f2(gg)} &times; 2 = <b>${f2(gA)}</b>.<br>\n",
    "  Every answer is <i>the answer from the step after me</i> &times; <i>how hard I push</i>. Multiplying along the chain like this is called the <b>chain rule</b>, and reading the tape backwards is called <b>backpropagation</b>. The two red numbers under the dials, ${f2(gA)} and ${f2(gB)}, are the <b>gradient</b>: the slope under Momo's feet in the foggy jungle, for both dials, found in <b>one</b> backward trip and with no walking.`;};\n",
    "window.vaTurn=function(){if(phase<2){vaBlame();bwT.push(setTimeout(vaTurn,timer?700:1900));return;}\n",
    "  const A0=A,B0=B;A=A-lr*gA;B=B-lr*gB;round++;$('vaRound').textContent=round;knobs();clearBw();clearFw();phase=0;btn(1);\n",
    "  $('vaCap').innerHTML=`<b>Turn the dials against their blame.</b> Stride (learning rate) = 0.02.<br>A = ${f2(A0)} &minus; 0.02 &times; (${f2(gA)}) = <b>${f2(A)}</b><br>B = ${f2(B0)} &minus; 0.02 &times; (${f2(gB)}) = <b>${f2(B)}</b><br>Press <b>Pick</b> again and watch the mistake shrink. This is one step of the foggy jungle walk above: the red numbers are the orange arrow. Our real model has 5.9 million dials, and one backward trip gives the blame for all of them at once.`;};\n",
    "window.vaWiggle=function(){if(phase<1)vaCook();const eps=0.1;const L2=((A+eps)*x+B-T)**2;const dA=L2-L;const pred=(phase>=2?gA:2*(g-T)*x)*eps;\n",
    "  $('vaCap').innerHTML=`<b>The slow way, to check.</b> Turn dial A up by ${eps} and pick again: A = ${f2(A+eps)} gives mistake ${f2(L2)} instead of ${f2(L)}, a change of <b>${f2(dA)}</b>.<br>The blame game predicted ${f2(pred/eps)} &times; ${eps} = <b>${f2(pred)}</b>. Same answer (tiny difference because ${eps} is not <i>truly</i> tiny). But the slow way needed one extra picking round for <i>one</i> dial. With 5.9 million dials that is 5.9 million extra rounds per step. The blame game gets all of them from one backward trip. That is why autograd made modern AI possible.`;};\n",
    "window.vaAuto=function(){if(timer){clearTimeout(timer);timer=null;return;}\n",
    "  const tick=()=>{if(phase===0){vaCook();}else if(phase===1){vaBlame();}else{vaTurn();if(L<0.1){timer=null;$('vaCap').innerHTML+='<br><b>Done:</b> Momo now picks '+f2(g)+' bananas, close enough to 10. Dials A = '+f2(A)+', B = '+f2(B)+'. Press Reset to play again.';return;}}\n",
    "    timer=setTimeout(tick,round<2?1600:600);};\n",
    "  timer=setTimeout(tick,0);};\n",
    "window.vaReset=function(){if(timer){clearTimeout(timer);timer=null;}A=1;B=3;round=0;phase=0;$('vaRound').textContent=0;knobs();clearFw();clearBw();btn(1);\n",
    "  $('vaCap').innerHTML='Momo starts with dial A = 1 and dial B = 3. He has no idea which way to turn them. Press <b>Pick</b>, then <b>Pass the blame back</b>, then <b>Turn the dials</b>. Or press <b>Auto</b>.';};\n",
    "vaReset();\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbbc8f63",
   "metadata": {},
   "source": [
    "Now the blame game in three lines of PyTorch. Three dials, one mistake, `backward()`, and\n",
    "each dial gets its blame. Then we turn every dial a small step against its blame. Walking\n",
    "downhill like this, one stride at a time, has a name: **gradient descent**. `torch.no_grad()`\n",
    "means \"switch the tape recorder off while we turn the dials; this part is not a step to be\n",
    "blamed later\". Below you will see three blames, one per dial, and then each dial has moved\n",
    "a little the other way. That is one lesson. Training is thousands of these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db819628",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = torch.tensor([1.0, -1.0, 0.5], requires_grad=True)   # three dials; requires_grad switches the tape on\n",
    "loss = ((x @ w) ** 2).sum()                              # the mistake, which depends on the dials\n",
    "loss.backward()                                          # the blame game: read the tape backwards, fill w.grad\n",
    "print(\"loss:\", loss.item())\n",
    "print(\"d(loss)/d(w):\", w.grad)\n",
    "\n",
    "# One step of gradient descent: turn every dial a little against its blame (stride 0.01).\n",
    "with torch.no_grad():\n",
    "    w -= 0.01 * w.grad\n",
    "print(\"w after one step:\", w)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99aa0e23",
   "metadata": {},
   "source": [
    "## 2. Data: Ella's bedtime stories (TinyStories)\n",
    "\n",
    "Ella needs stories to read. [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)\n",
    "is a public **dataset** (a collection of examples to learn from) of short children's stories\n",
    "written in very simple English. It was created to show that *small* models can learn to write\n",
    "real sentences if the language is simple enough.\n",
    "\n",
    "We stream a slice of it so the download stays small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc46895d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "N_STORIES = 20_000   # enough for a tiny model; raise it later if you want a better one\n",
    "\n",
    "stream = load_dataset(\"roneneldan/TinyStories\", split=\"train\", streaming=True)\n",
    "stories = [row[\"text\"].strip() for row in stream.take(N_STORIES)]\n",
    "\n",
    "print(f\"{len(stories):,} stories, {sum(len(s) for s in stories)/1e6:.1f}M characters\")\n",
    "print(\"-\" * 60)\n",
    "print(stories[0][:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1eca3fe",
   "metadata": {},
   "source": [
    "### Ella's own stories about Momo\n",
    "\n",
    "TinyStories is full of Lily, Tom and their mums. If we want Momo and Ella to show up when the\n",
    "model writes, the model has to *read* about them. So we add about a thousand short stories\n",
    "about our two characters, built from sentence pools in the same simple style (every story has\n",
    "an opening, Ella's introduction, a small problem, Ella helping, a resolution and an ending).\n",
    "We repeat them three times so they are not drowned out by the 20,000 other stories. This is\n",
    "one way a model is taught new material: mix it into the training text. (Grown-ups more often\n",
    "*fine-tune*, which is the same mixing done to a model that has already been trained.)\n",
    "\n",
    "One more thing happens here. Ella hides the last 200 TinyStories and never reads them to\n",
    "Momo. They are her **quiz**, the **validation set**. After each lesson she asks Momo about\n",
    "*those*, so we learn whether he understands the language or has only memorised the stories he\n",
    "heard. The Momo stories go only into the lessons, never into the quiz, so a repeated story can\n",
    "never make the quiz look better than it is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "215e5fa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Momo and Ella stories, built from sentence pools. Same code as scripts/momo_stories.py in the repo.\n",
    "import random\n",
    "\n",
    "PLACES = [\"by the river\", \"under the big mango tree\", \"in the green jungle\", \"at the waterfall\",\n",
    "          \"on the sunny hill\", \"near the little pond\", \"in the tall grass\", \"by the old rock\"]\n",
    "THINGS = [\"a shiny red ball\", \"a yellow banana\", \"a big green leaf\", \"a round grey stone\",\n",
    "          \"a blue feather\", \"a little drum\", \"a brown coconut\", \"a soft pink flower\",\n",
    "          \"a long stick\", \"a red kite\"]\n",
    "FRIENDS = [\"Pip the parrot\", \"Tilly the turtle\", \"Bo the bear cub\", \"Lulu the little lion\"]\n",
    "FOODS = [\"bananas\", \"mangoes\", \"sweet berries\", \"coconuts\", \"crunchy leaves\"]\n",
    "\n",
    "OPENINGS = [\n",
    "    \"Once upon a time, there was a little monkey named Momo.\",\n",
    "    \"Momo was a little monkey who lived {place}.\",\n",
    "    \"One day, Momo the monkey woke up very early.\",\n",
    "    \"Momo was a small monkey with a big smile.\",\n",
    "    \"In the jungle there lived a little monkey called Momo.\",\n",
    "    \"Every morning, Momo the monkey jumped out of his tree.\",\n",
    "]\n",
    "ELLA_INTRO = [\n",
    "    \"His best friend was Ella, a big grey elephant.\",\n",
    "    \"Ella the elephant was his best friend. Ella never forgot anything.\",\n",
    "    \"He loved his friend Ella. Ella was an elephant, and she remembered everything.\",\n",
    "    \"Ella the elephant was big and kind and very wise.\",\n",
    "    \"His friend Ella was an elephant with a long trunk and a good memory.\",\n",
    "]\n",
    "# Each scene: (problem, Ella helps, resolution). They always go together.\n",
    "SCENES = [\n",
    "    (\"Momo wanted to play with {thing}, but he could not find it. He looked {place}. It was not there. He felt sad.\",\n",
    "     \"Ella came slowly. \\\"Do not worry, Momo,\\\" she said. \\\"I remember. You left it {place2}.\\\"\",\n",
    "     \"Momo ran and found {thing} right where Ella said. He hugged her trunk.\"),\n",
    "    (\"Momo climbed very high to get {thing}. Then he looked down and felt scared. \\\"I cannot get down!\\\" he cried.\",\n",
    "     \"Ella walked under the tree and lifted her trunk. \\\"I am here,\\\" she said. \\\"You are safe. Come down slowly.\\\"\",\n",
    "     \"Momo climbed down step by step and jumped into Ella's arms.\"),\n",
    "    (\"It started to rain. Momo got wet and cold. He hid {place} and shivered.\",\n",
    "     \"Ella held a big leaf over Momo like an umbrella. \\\"Stay dry, my friend,\\\" she said.\",\n",
    "     \"The rain stopped. The sun came out. Momo and Ella walked home together.\"),\n",
    "    (\"Momo met {friend} {place}. {friend_short} had {thing} and did not want to share. Momo felt sad.\",\n",
    "     \"Ella walked over and said, \\\"Let us all play together. Sharing makes everyone happy.\\\"\",\n",
    "     \"Momo and {friend_short} became friends and played with {thing} all day.\"),\n",
    "    (\"Momo tried to count his {food}. \\\"One, two, three...\\\" Then he forgot the next number.\",\n",
    "     \"Ella said, \\\"After three comes four. I will count with you.\\\" They counted together.\",\n",
    "     \"Momo counted all the way to ten. Ella was very proud.\"),\n",
    "    (\"Momo had a bad dream. He woke up scared and ran to find Ella.\",\n",
    "     \"Ella sat with Momo until he felt better. \\\"Being scared is okay,\\\" she said. \\\"I am right here.\\\"\",\n",
    "     \"Momo fell asleep next to Ella and had happy dreams.\"),\n",
    "    (\"Momo wanted to tell a story, but he did not know how to start. He sat down {place} and thought and thought.\",\n",
    "     \"Ella said, \\\"Start with 'Once upon a time'. Then say what happens next. One word at a time.\\\"\",\n",
    "     \"Momo told his first story. It was about {thing}. Ella clapped her big ears.\"),\n",
    "    (\"Momo saw a big shadow {place}. He thought it was a monster. He hid behind {thing}.\",\n",
    "     \"Ella laughed softly. \\\"That shadow is only me,\\\" she said. Momo peeked out and laughed too.\",\n",
    "     \"They played shadow games until the sun went down.\"),\n",
    "    (\"Momo was sad because he could not swim like {friend}. He sat {place} with his feet in the water.\",\n",
    "     \"Ella said, \\\"Try again, little one. Every time you try, you learn a little more.\\\" She held him up with her trunk.\",\n",
    "     \"Momo kicked and kicked. Soon he could swim all by himself.\"),\n",
    "    (\"Momo and Ella went {place} to look for {food}. They looked and looked, but they found none.\",\n",
    "     \"Ella thought for a moment. Elephants never forget. \\\"I know a secret place,\\\" she said.\",\n",
    "     \"They found lots of {food} and ate them together under the stars.\"),\n",
    "    (\"Momo broke {thing} that belonged to {friend}. He hid it and felt bad all day.\",\n",
    "     \"Ella said, \\\"A true friend says sorry. Go and tell {friend_short}.\\\"\",\n",
    "     \"Momo said sorry. {friend_short} smiled and said, \\\"We can fix it together.\\\"\"),\n",
    "    (\"Momo wanted to be big like Ella. He stood on {thing} and stretched, but he was still small.\",\n",
    "     \"Ella smiled. \\\"You are small, but you can climb and jump and I cannot,\\\" she said. \\\"We are different, and that is good.\\\"\",\n",
    "     \"Momo climbed to the top of the tree and threw {food} down to Ella. They both laughed.\"),\n",
    "]\n",
    "ENDINGS = [\n",
    "    \"The end.\",\n",
    "    \"Momo was happy. Ella was happy. The end.\",\n",
    "    \"And they were best friends forever. The end.\",\n",
    "    \"\\\"Thank you, Ella,\\\" said Momo. \\\"You always remember.\\\" The end.\",\n",
    "    \"That night, Momo dreamed about his friend Ella. The end.\",\n",
    "]\n",
    "\n",
    "\n",
    "def make_story(rng: random.Random) -> str:\n",
    "    friend = rng.choice(FRIENDS)\n",
    "    place, place2 = rng.sample(PLACES, 2)\n",
    "    slots = dict(place=place, place2=place2, thing=rng.choice(THINGS), friend=friend,\n",
    "                 friend_short=friend.split(\" the \")[0], food=rng.choice(FOODS))\n",
    "    problem, helps, resolution = rng.choice(SCENES)\n",
    "    parts = [rng.choice(OPENINGS), rng.choice(ELLA_INTRO), problem, helps, resolution, rng.choice(ENDINGS)]\n",
    "    return \" \".join(p.format(**slots) for p in parts)\n",
    "\n",
    "\n",
    "def momo_stories(n: int = 1000, seed: int = 0) -> list[str]:\n",
    "    # n distinct stories: repeats from the pools are skipped, so \"unique\" really is unique.\n",
    "    rng = random.Random(seed)\n",
    "    out, seen = [], set()\n",
    "    while len(out) < n:\n",
    "        s = make_story(rng)\n",
    "        if s not in seen:\n",
    "            seen.add(s)\n",
    "            out.append(s)\n",
    "    return out\n",
    "\n",
    "momo = momo_stories(1000, seed=0)\n",
    "val_stories = stories[-200:]                    # Ella's hidden quiz: TinyStories only, never trained on\n",
    "train_stories = stories[:-200] + momo * 3       # 19,800 TinyStories + 3,000 Momo stories (1,000 unique, each 3 times)\n",
    "random.Random(42).shuffle(train_stories)        # shuffle so Momo stories are spread through the lessons\n",
    "stories = train_stories                         # from here on, \"stories\" means the lesson stories\n",
    "print(f\"{len(train_stories):,} lesson stories ({len(momo)*3:,} about Momo, {len(set(momo)):,} unique) | {len(val_stories)} quiz stories\")\n",
    "print(\"-\" * 60)\n",
    "print(momo[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "550278c3",
   "metadata": {},
   "source": [
    "## 3. Momo's first trick: the bigram\n",
    "\n",
    "A language model answers one question: **given what came before, what comes next?**\n",
    "\n",
    "Momo's very first trick is to remember only the *last letter*. He keeps a tray: \"after the\n",
    "letter t, how often does h come next? how often does e?\" That tray is a **bigram model**\n",
    "(bi = two: it looks at pairs of letters). In code it is one `nn.Embedding` table of size\n",
    "*vocab × vocab*. An **embedding table** is just a tray where row number *i* belongs to token *i*.\n",
    "In the introduction the embedding was \"a spot on the meaning map\": that map is this kind of\n",
    "tray, one row per token. Here we borrow the same kind of tray for a different job: row *i*\n",
    "holds \"what tends to follow token *i*\". Same box, different contents.\n",
    "\n",
    "First Momo needs a vocabulary. For now, one token = one letter. **Encode** means turning\n",
    "letters into their numbers; **decode** means turning the numbers back into letters. Momo only\n",
    "ever sees the numbers.\n",
    "\n",
    "Play with the idea first. The widget counts, from a few sentences about Momo and Ella, which\n",
    "letter tends to follow which. Click a letter. That tray *is* the whole model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b4022c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: the bigram table, counted live from a few sentences.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vb\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>#vb .k{display:inline-block;width:26px;height:26px;line-height:26px;text-align:center;margin:2px;border-radius:5px;background:#2f4f6f;border:1px solid #5a8ab8;cursor:pointer;font-size:13px}#vb .k.on{background:#ffb347;color:#000}#vb .bar{height:14px;background:#6fcf97;border-radius:3px}#vb .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Bigram model: click a letter to see what it thinks comes next.</div>\n",
    "<div id=\"vbKeys\"></div>\n",
    "<div id=\"vbOut\" style=\"margin-top:10px;min-height:120px\"></div>\n",
    "<div class=\"cap\">A bigram model is nothing but this table: for every character, how often each other character followed it in the training text. It is counted here live from a few sample sentences. It cannot remember anything further back than one letter, which is why its stories are gibberish. The GPT fixes that with attention.</div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const corpus=\"once upon a time there was a little monkey called momo. momo liked to play in the jungle with his friend ella the elephant. one day they found a shiny red ball. they were very happy and played all day. the end. ella had a big red balloon. she gave it to momo. momo smiled and said thank you.\";\n",
    "const keys=\"abcdefghijklmnopqrstuvwxyz. \".split('');const counts={};\n",
    "for(let i=0;i<corpus.length-1;i++){const a=corpus[i],b=corpus[i+1];counts[a]=counts[a]||{};counts[a][b]=(counts[a][b]||0)+1;}\n",
    "const kd=document.getElementById('vbKeys');keys.forEach(k=>{const s=document.createElement('span');s.className='k';s.textContent=k===' '?'␣':k;s.onclick=()=>vbShow(k);kd.appendChild(s);});\n",
    "window.vbShow=function(k){[...kd.children].forEach((s,i)=>s.classList.toggle('on',keys[i]===k));const c=counts[k]||{};const tot=Object.values(c).reduce((a,b)=>a+b,0);\n",
    " const all=Object.entries(c).sort((a,b)=>b[1]-a[1]);const rows=all.slice(0,6);const rest=all.slice(6).reduce((s,[,n])=>s+n,0);\n",
    " document.getElementById('vbOut').innerHTML=tot?`<div style=\"font-size:13px;margin-bottom:6px\">After <b>\"${k===' '?'space':k}\"</b> (seen ${tot} times), the next character was:</div>`+rows.map(([ch,n])=>`<div style=\"display:flex;align-items:center;gap:8px;font-size:13px\"><span style=\"width:40px;text-align:right\">${ch===' '?'␣':ch}</span><div class=\"bar\" style=\"width:${Math.round(n/tot*300)}px\"></div><span style=\"color:#aaa\">${Math.round(n/tot*100)}%</span></div>`).join('')+(rest?`<div style=\"display:flex;align-items:center;gap:8px;font-size:13px;color:#888\"><span style=\"width:40px;text-align:right\">other</span><div class=\"bar\" style=\"width:${Math.round(rest/tot*300)}px;background:#555\"></div><span>${Math.round(rest/tot*100)}% (${all.length-6} other letters)</span></div>`:''):`<div style=\"font-size:13px;color:#aaa\">\"${k}\" never appears in the sample text, so the model has no idea.</div>`;};\n",
    "vbShow('t');\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc2334a9",
   "metadata": {},
   "source": [
    "Now the same thing in code: first the letter vocabulary and the encode/decode pair."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ebe181d",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\\n\".join(stories)\n",
    "chars = sorted(set(text))\n",
    "char_vocab = len(chars)\n",
    "stoi = {c: i for i, c in enumerate(chars)}   # string -> int\n",
    "itos = {i: c for c, i in stoi.items()}       # int -> string\n",
    "encode_c = lambda s: [stoi[c] for c in s]\n",
    "decode_c = lambda ids: \"\".join(itos[i] for i in ids)\n",
    "\n",
    "print(\"character vocab size:\", char_vocab)\n",
    "print(\"encode('hello') ->\", encode_c(\"hello\"))\n",
    "print(\"decode back     ->\", decode_c(encode_c(\"hello\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5cc8528",
   "metadata": {},
   "source": [
    "Next, Momo's first lessons. Every lesson in this notebook, here and for the big model in\n",
    "Part 7, is the same **five moves**:\n",
    "\n",
    "1. Ella picks a **batch**: a handful of short pieces of text (here 64 pieces of 8 letters),\n",
    "   and for each piece the same text shifted by one letter, which is the list of right answers.\n",
    "2. Momo guesses the next letter at every position.\n",
    "3. Ella scores the guesses with the **loss**.\n",
    "4. The blame game runs backwards: every dial gets its gradient.\n",
    "5. The optimizer turns the dials.\n",
    "\n",
    "Then Momo writes: start from one letter, turn his scores into percentages, spin a wheel whose\n",
    "slices are those percentages to pick the next letter, and repeat. We use only the first two\n",
    "million letters so this toy runs in seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1c9d084",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bigram(nn.Module):\n",
    "    def __init__(self, vocab):\n",
    "        super().__init__()\n",
    "        self.table = nn.Embedding(vocab, vocab)   # row i = scores for \"what follows token i\"\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        logits = self.table(idx)                  # (B, T) -> (B, T, vocab)\n",
    "        loss = None\n",
    "        if targets is not None:\n",
    "            # cross-entropy compares predicted scores with the true next token\n",
    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        for _ in range(max_new_tokens):\n",
    "            logits, _ = self(idx[:, -1:])                       # only the last token matters\n",
    "            probs = F.softmax(logits[:, -1, :], dim=-1)             # scores -> percentages\n",
    "            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)  # spin a wheel whose slices are the percentages\n",
    "        return idx\n",
    "\n",
    "char_data = torch.tensor(encode_c(text[:2_000_000]), dtype=torch.long)   # first 2 million letters are plenty for a toy\n",
    "\n",
    "def get_char_batch(batch_size=64, block_size=8):\n",
    "    ix = torch.randint(len(char_data) - block_size - 1, (batch_size,))\n",
    "    x = torch.stack([char_data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([char_data[i+1:i+block_size+1] for i in ix])   # the \"next character\" of each position\n",
    "    return x.to(device), y.to(device)\n",
    "\n",
    "bigram = Bigram(char_vocab).to(device)\n",
    "opt = torch.optim.AdamW(bigram.parameters(), lr=1e-2)   # the dial-turning rule, stride 0.01\n",
    "\n",
    "for step in range(2000):\n",
    "    xb, yb = get_char_batch()                 # move 1: a batch and its right answers\n",
    "    _, loss = bigram(xb, yb)                  # moves 2 and 3: guess, and score with the loss\n",
    "    opt.zero_grad(); loss.backward(); opt.step()   # move 4: blame game; move 5: turn the dials (zero_grad wipes last lesson's blames first)\n",
    "    if step % 500 == 0:\n",
    "        print(f\"step {step:4d}  loss {loss.item():.3f}\")\n",
    "\n",
    "start = torch.tensor([encode_c(\"T\")], device=device)\n",
    "print(\"\\nBigram writes:\\n\", decode_c(bigram.generate(start, 300)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be9bfbb8",
   "metadata": {},
   "source": [
    "Momo learned *something* (the loss dropped, and he produces word-ish shapes with spaces), but\n",
    "it is gibberish. He only looks one letter back. To tell a story he needs two upgrades:\n",
    "**better bites** and a way to **look further back**. That is the rest of the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f160514",
   "metadata": {},
   "source": [
    "## 4. A real tokenizer: how Momo decides where to bite\n",
    "\n",
    "Momo bites a sentence into pieces before he can think about it. How big should a bite be?\n",
    "\n",
    "- Bite every letter on its own, and Momo wastes all his effort learning to spell.\n",
    "- Bite whole words, and he needs a piece for every word in the world, and a word he has\n",
    "  never seen breaks him.\n",
    "- **BPE** (byte-pair encoding) is the middle way. Start with single letters. Find the two\n",
    "  pieces that sit next to each other most often (\"t\" and \"h\") and glue them into one piece\n",
    "  (\"th\"). Repeat thousands of times. Common words like \" the\" become one bite. A rare word\n",
    "  like \" hippopotamus\" becomes a few bites (\"hipp\", \"opot\", \"amus\"). Every big model (GPT,\n",
    "  Llama, Qwen) bites this way.\n",
    "\n",
    "We teach Momo his own bites from the stories, 4,096 pieces in total. Each piece gets a\n",
    "number, and from then on Momo only ever sees the numbers.\n",
    "\n",
    "Type a sentence below and watch it get bitten into pieces and turned into numbers. (This\n",
    "widget only pretends, to show the idea. The real biter, trained from our stories, comes in the\n",
    "next cell and may bite differently.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026dd481",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: type a sentence, see tokens and ids (illustrative splitter).\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vk\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>#vk .chip{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f4f6f;border:1px solid #5a8ab8;font-size:13px}#vk .id{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f6f4f;border:1px solid #6fcf97;font-size:12px;font-family:monospace}#vk input{width:95%;padding:8px;border-radius:6px;border:1px solid #777;background:#262626;color:#eee;font-size:14px}#vk .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Tokenizer: type anything and watch it get bitten into tokens, then numbers.</div>\n",
    "<input id=\"vkIn\" value=\"Once upon a time, Momo the monkey found a shiny red ball.\" oninput=\"vkRun()\">\n",
    "<div style=\"font-size:12px;color:#aaa;margin:8px 0 2px\">tokens (&#9251; marks a leading space)</div><div id=\"vkTok\"></div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin:8px 0 2px\">pretend token ids (the real ones come from the trained biter in the next cell)</div><div id=\"vkIds\"></div>\n",
    "<div class=\"cap\">A tokenizer splits text into pieces it has seen often. Common words stay whole (\"little\"), rare or long words break into parts (\"unbeli\" + \"evable\"), and punctuation is its own token. Each piece has a fixed number in the dictionary. <i>This demo mimics the idea in JavaScript; the real one is trained from data in Part 4.</i></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "function hash(s){let h=7;for(const ch of s)h=(h*31+ch.codePointAt(0))%4096;return h;}\n",
    "function split(t){const out=[];const re=/\\s?\\p{L}+|\\s?\\d+|\\s?[^\\s\\p{L}\\d]/gu;let m;while((m=re.exec(t))){let p=m[0];if(p.replace(/^\\s/,'').length>7){const lead=p.startsWith(' ')?' ':'';const w=p.trim();out.push(lead+w.slice(0,Math.ceil(w.length/2)));out.push(w.slice(Math.ceil(w.length/2)));}else out.push(p);}return out;}\n",
    "window.vkRun=function(){const t=document.getElementById('vkIn').value;const toks=split(t);\n",
    " document.getElementById('vkTok').innerHTML=toks.map(p=>`<span class=\"chip\">${p.replace(/^ /,'␣')}</span>`).join('');\n",
    " document.getElementById('vkIds').innerHTML=toks.map(p=>`<span class=\"id\">${hash(p)}</span>`).join('')+`<span style=\"font-size:12px;color:#aaa;margin-left:8px\">${t.length} characters &#8594; ${toks.length} tokens</span>`;};\n",
    "vkRun();\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f80df23",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import ByteLevelBPETokenizer\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "\n",
    "VOCAB_SIZE = 4096\n",
    "EOS = \"<|endoftext|>\"   # special token that marks where one story ends\n",
    "\n",
    "bpe = ByteLevelBPETokenizer()\n",
    "bpe.train_from_iterator(stories, vocab_size=VOCAB_SIZE, min_frequency=2, special_tokens=[EOS])   # min_frequency: a pair must appear at least twice to be glued\n",
    "\n",
    "# Wrap it in the standard Hugging Face tokenizer class. We reuse the one special marker for\n",
    "# \"story starts\" (bos), \"story ends\" (eos) and \"empty space\" (pad); a tiny model needs no more.\n",
    "tokenizer = PreTrainedTokenizerFast(tokenizer_object=bpe, bos_token=EOS, eos_token=EOS, pad_token=EOS)\n",
    "vocab_size = len(tokenizer)\n",
    "\n",
    "sample = \"Once upon a time, Momo the monkey found a shiny red ball.\"\n",
    "ids = tokenizer(sample)[\"input_ids\"]\n",
    "print(\"vocab size:\", vocab_size)\n",
    "print(f\"{len(sample)} characters -> {len(ids)} tokens\")\n",
    "print(tokenizer.convert_ids_to_tokens(ids))   # 'Ġ' means 'there is a space before this piece'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a8b12f",
   "metadata": {},
   "source": [
    "Now every story becomes numbers. We glue all the lesson stories into one long stream of token\n",
    "ids with the \"story is over\" marker between them, and do the same, separately, for Ella's\n",
    "hidden quiz stories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d96b25f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_ids(story_list):\n",
    "    return torch.tensor(tokenizer(EOS.join(story_list) + EOS)[\"input_ids\"], dtype=torch.long)\n",
    "\n",
    "train_data = to_ids(train_stories)   # Momo's lessons\n",
    "val_data = to_ids(val_stories)       # Ella's quiz, never trained on\n",
    "print(f\"train {len(train_data)/1e6:.2f}M tokens | quiz {len(val_data)/1e3:.0f}K tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7ab127",
   "metadata": {},
   "source": [
    "## 5. Attention, by hand: the story circle in numbers\n",
    "\n",
    "This is the idea that made modern LLMs possible. Each token gets to **look at the tokens\n",
    "before it** and decide how much to listen to each one. It is the story circle from the\n",
    "introduction, now in code.\n",
    "\n",
    "Every token makes three vectors, its three cards:\n",
    "- a **query** (\"what am I looking for?\"),\n",
    "- a **key** (\"what am I?\"),\n",
    "- a **value** (\"what do I tell you if you pick me?\").\n",
    "\n",
    "A token's output is a weighted blend of the *values* of earlier tokens (and its own). The\n",
    "weights come from how well its *query* matches their *keys*, squashed into percentages by\n",
    "**softmax**. A triangle-shaped **mask** blocks the future, because the future is what Momo is\n",
    "trying to guess.\n",
    "\n",
    "In the code, `nn.Linear` is one tray of dials: it multiplies whatever comes in by that tray,\n",
    "exactly like the fruit widget in Part 1. Three trays make the three cards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69abace2",
   "metadata": {},
   "outputs": [],
   "source": [
    "B, T, C = 1, 4, 8                      # batch, time (tokens), numbers per token\n",
    "x_demo = torch.randn(B, T, C)\n",
    "head_size = 4                          # how long each card is\n",
    "q = nn.Linear(C, head_size, bias=False)(x_demo)\n",
    "k = nn.Linear(C, head_size, bias=False)(x_demo)\n",
    "v = nn.Linear(C, head_size, bias=False)(x_demo)\n",
    "\n",
    "scores = q @ k.transpose(-2, -1) / math.sqrt(head_size)  # (B, T, T): how much token i matches token j.\n",
    "                                                          # Dividing by the square root of the card length keeps scores from getting huge when cards get long.\n",
    "mask = torch.tril(torch.ones(T, T))                       # 1 = allowed to look, 0 = future\n",
    "scores = scores.masked_fill(mask == 0, float(\"-inf\"))     # future positions get -inf -> 0 after softmax\n",
    "weights = F.softmax(scores, dim=-1)\n",
    "out = weights @ v                                         # (B, T, head_size)\n",
    "\n",
    "torch.set_printoptions(precision=2, sci_mode=False)\n",
    "print(\"attention weights (row = token, column = who it looks at):\\n\", weights[0])\n",
    "print(\"\\nRow 1 only looks at tokens 0-1, row 3 can see all four. Each row sums to 1.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c24308af",
   "metadata": {},
   "source": [
    "### Ella's shelf (the KV cache), in about a dozen lines\n",
    "\n",
    "When Momo writes, a new token arrives at every step. The keys and values of the tokens\n",
    "already written **never change**, so Ella keeps them on her shelf (the cache) and Momo makes\n",
    "cards only for the newcomer. Below: four tokens are already on the shelf, a fifth arrives,\n",
    "and we check that the shelf gives exactly the same answer as redoing all five from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84949030",
   "metadata": {},
   "outputs": [],
   "source": [
    "Wq, Wk, Wv = nn.Linear(C, head_size, bias=False), nn.Linear(C, head_size, bias=False), nn.Linear(C, head_size, bias=False)\n",
    "\n",
    "x_old = torch.randn(1, 4, C)                  # 4 tokens already processed earlier\n",
    "K_cache, V_cache = Wk(x_old), Wv(x_old)       # their key/value cards, stored once: Ella's shelf\n",
    "\n",
    "x_new = torch.randn(1, 1, C)                  # the 5th token arrives\n",
    "q_new, k_new, v_new = Wq(x_new), Wk(x_new), Wv(x_new)   # only ITS cards are computed\n",
    "K_cache = torch.cat([K_cache, k_new], dim=1)  # the shelf grows by one card\n",
    "V_cache = torch.cat([V_cache, v_new], dim=1)\n",
    "\n",
    "att = F.softmax(q_new @ K_cache.transpose(-2, -1) / math.sqrt(head_size), dim=-1)  # 1 query vs 5 stored keys\n",
    "out_cached = att @ V_cache\n",
    "\n",
    "# Check against the slow way: recompute attention for all 5 tokens from scratch, take the last row.\n",
    "x_all = torch.cat([x_old, x_new], dim=1)\n",
    "scores = Wq(x_all) @ Wk(x_all).transpose(-2, -1) / math.sqrt(head_size)\n",
    "scores = scores.masked_fill(torch.tril(torch.ones(5, 5)) == 0, float(\"-inf\"))\n",
    "out_full = (F.softmax(scores, dim=-1) @ Wv(x_all))[:, -1:]\n",
    "\n",
    "print(\"tokens stored in cache:\", K_cache.shape[1])\n",
    "print(\"cached shortcut == full recompute:\", torch.allclose(out_cached, out_full, atol=1e-6))\n",
    "print(\"work done this step: 1 token's cards instead of 5. For a 1,000-word chat, that is 1 instead of 1,000.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "626f6cd8",
   "metadata": {},
   "source": [
    "## 6. A tiny GPT: the whole of Momo\n",
    "\n",
    "**GPT** stands for Generative Pre-trained Transformer. *Generative* because it makes new\n",
    "text. *Pre-trained* because it learns from stories before anyone uses it. *Transformer*\n",
    "because it is built from attention layers stacked on top of each other. Here is the whole of\n",
    "Momo, bottom to top, like the tall jungle tree from the introduction. The rungs come in\n",
    "order, and every token climbs every rung:\n",
    "\n",
    "1. **Token embedding**: each token number becomes a row of `n_embd` numbers, its spot on the\n",
    "   meaning map.\n",
    "2. **Position embedding**: each seat number (1st, 2nd, 3rd...) also gets a row, so Momo knows\n",
    "   word order.\n",
    "3. **Blocks**, repeated `n_layer` times (the 6 rungs of the tree). Each block has:\n",
    "   - **multi-head attention**: the story circle, run 4 times at once with different cards;\n",
    "   - a **feed-forward network**: after listening to everyone, each token sits on its own\n",
    "     twig and thinks alone (a small two-layer network; research suggests many facts end up\n",
    "     stored here);\n",
    "   - **LayerNorm** (Ella keeping the numbers at a steady volume) and **residual connections**\n",
    "     (each block *adds* its thoughts to what came in instead of replacing it, so nothing is\n",
    "     lost on the way up the tree).\n",
    "4. **Final LayerNorm and the scoring tray**: turn each row into a score for every token in\n",
    "   the vocabulary (the logits). Programmers call this tray the \"output head\", a different use\n",
    "   of the word head from the attention heads.\n",
    "\n",
    "We also write `generate`: guess the next token, add it, repeat.\n",
    "\n",
    "Click any layer in the tree below to read what it does and why it is there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af4363f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: the model's layers, bottom to top. Click a layer.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vs\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>#vs .box{display:block;width:300px;margin:4px auto;padding:8px;text-align:center;border-radius:8px;border:2px solid #5a8ab8;background:#2f4f6f;cursor:pointer;font-size:13px;transition:all .2s}#vs .box:hover,#vs .box.on{border-color:#ffb347;background:#3a2e1a}#vs .blk{border:2px dashed #6fcf97;border-radius:10px;padding:8px;width:340px;margin:6px auto;position:relative}#vs .tag{position:absolute;right:-70px;top:40%;color:#6fcf97;font-size:13px;font-weight:600}#vs .ar{text-align:center;color:#888;font-size:16px;line-height:14px}#vs .cap{font-size:13px;line-height:1.5;background:#262626;border-radius:8px;padding:10px;min-height:70px}@media (max-width:760px){#vs .wrap{flex-direction:column-reverse}}</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">The tiny GPT, bottom to top. Click any layer.</div>\n",
    "<div class=\"wrap\" style=\"display:flex;gap:16px;flex-wrap:wrap\"><div style=\"flex:0 0 420px\">\n",
    "<span class=\"box\" onclick=\"vsX(this,'out')\">next-word scores (4,096 numbers)</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'lnf')\">final LayerNorm + scoring tray (output head)</span><div class=\"ar\">&#8593;</div>\n",
    "<div class=\"blk\"><span class=\"tag\">&times; 6 blocks</span>\n",
    "<span class=\"box\" onclick=\"vsX(this,'res2')\">+ add back (residual)</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'ff')\">feed-forward network</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'ln')\">LayerNorm</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'res1')\">+ add back (residual)</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'attn')\">multi-head attention (4 heads)</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'ln')\">LayerNorm</span></div>\n",
    "<div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'emb')\">token embedding + position embedding</span><div class=\"ar\">&#8593;</div>\n",
    "<span class=\"box\" onclick=\"vsX(this,'tok')\">token ids (up to 256 of them)</span>\n",
    "</div><div style=\"flex:1;min-width:260px\"><div class=\"cap\" id=\"vsCap\">Click a layer to read what it does and why it is there.</div></div></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){const t={\n",
    "tok:\"<b>Token ids.</b> The input: a list of up to 256 integers from the tokenizer. Picture: the number painted on each of Momo's bite-sized pieces.\",\n",
    "emb:\"<b>Embeddings.</b> Each id is swapped for a 256-number vector (its meaning), and a second vector for its position (1st, 2nd, 3rd...) is added, so the model knows word order. Picture: each animal in the story circle gets a name tag and a seat number.\",\n",
    "ln:\"<b>LayerNorm.</b> Rescales the numbers so they stay in a sensible range. It does no 'thinking'; it keeps training stable. Picture: Ella turning the volume to a steady level before each round.\",\n",
    "attn:\"<b>Multi-head attention.</b> The story circle: each token looks back at earlier tokens (and itself) and blends in what matters. 4 heads run in parallel, each looking for something different (grammar, names, position...). This is the only place tokens talk to each other.\",\n",
    "res1:\"<b>Residual (add back).</b> The attention output is <i>added</i> to the original, not replacing it. The token keeps its own meaning plus what it learned. It also gives gradients a direct path, which is what lets deep networks train at all.\",\n",
    "ff:\"<b>Feed-forward network.</b> A small two-layer network applied to every token on its own, 256 &#8594; 1024 &#8594; 256. Research suggests many facts and patterns end up stored here. Picture: after listening to everyone in the circle (attention), each animal sits on its own twig and thinks alone.\",\n",
    "res2:\"<b>Residual again.</b> Same trick: add the feed-forward result back onto the stream.\",\n",
    "lnf:\"<b>Final LayerNorm + scoring tray.</b> Programmers call this the output head, a different use of the word from the attention heads. One last rescale, then each token's 256-number vector is multiplied against the embedding table (shared weights, 'weight tying') to produce a score for all 4,096 dictionary entries.\",\n",
    "out:\"<b>Scores.</b> 4,096 numbers per position. Softmax turns them into probabilities, and the next word is sampled from them. During training these scores are compared with the real next word to compute the loss.\"};\n",
    "window.vsX=function(el,k){document.querySelectorAll('#vs .box').forEach(b=>b.classList.remove('on'));el.classList.add('on');document.getElementById('vsCap').innerHTML=t[k];};\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54e5e5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config:\n",
    "    vocab_size = vocab_size\n",
    "    block_size = 256     # the context window from the introduction: how many tokens Momo can see at once\n",
    "    n_embd = 256         # size of each token's vector\n",
    "    n_head = 4           # attention heads per block\n",
    "    n_layer = 6          # number of blocks\n",
    "    dropout = 0.0        # tiny model + lots of data -> no dropout needed\n",
    "\n",
    "cfg = Config()\n",
    "\n",
    "class CausalSelfAttention(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.n_head = cfg.n_head\n",
    "        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd)      # query, key, value for all heads at once\n",
    "        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd)\n",
    "        self.drop = nn.Dropout(cfg.dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape\n",
    "        q, k, v = self.qkv(x).split(C, dim=2)\n",
    "        # reshape so each head works on its own slice: (B, n_head, T, head_size)\n",
    "        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
    "        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
    "        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
    "        # exactly the masked softmax(QK^T/sqrt(d))V from Part 5, done by a fast built-in kernel\n",
    "        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)\n",
    "        y = y.transpose(1, 2).contiguous().view(B, T, C)      # merge the heads back together\n",
    "        return self.drop(self.proj(y))\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(cfg.n_embd, 4 * cfg.n_embd),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(4 * cfg.n_embd, cfg.n_embd),\n",
    "            nn.Dropout(cfg.dropout),\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.ln1 = nn.LayerNorm(cfg.n_embd)\n",
    "        self.attn = CausalSelfAttention(cfg)\n",
    "        self.ln2 = nn.LayerNorm(cfg.n_embd)\n",
    "        self.ff = FeedForward(cfg)\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))   # \"x +\" is the residual connection\n",
    "        x = x + self.ff(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)\n",
    "        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)\n",
    "        self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(cfg.n_embd)\n",
    "        self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)\n",
    "        self.head.weight = self.tok_emb.weight    # weight tying: same matrix in and out, saves parameters\n",
    "        self.apply(self._init)\n",
    "\n",
    "    def _init(self, m):\n",
    "        if isinstance(m, (nn.Linear, nn.Embedding)):\n",
    "            nn.init.normal_(m.weight, mean=0.0, std=0.02)\n",
    "        if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "            nn.init.zeros_(m.bias)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        pos = torch.arange(T, device=idx.device)\n",
    "        x = self.tok_emb(idx) + self.pos_emb(pos)          # (B, T, n_embd)\n",
    "        x = self.ln_f(self.blocks(x))\n",
    "        logits = self.head(x)                              # (B, T, vocab_size)\n",
    "        loss = None\n",
    "        if targets is not None:\n",
    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40, eos_id=None):\n",
    "        for _ in range(max_new_tokens):\n",
    "            idx_cond = idx[:, -self.cfg.block_size:]       # never feed more than block_size tokens\n",
    "            logits, _ = self(idx_cond)\n",
    "            logits = logits[:, -1, :] / temperature         # lower temperature = safer, more boring\n",
    "            if top_k is not None:\n",
    "                v, _ = torch.topk(logits, top_k)\n",
    "                logits[logits < v[:, [-1]]] = float(\"-inf\")  # keep only the k most likely tokens\n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            next_id = torch.multinomial(probs, 1)\n",
    "            idx = torch.cat([idx, next_id], dim=1)\n",
    "            if eos_id is not None and (next_id == eos_id).all():   # the model decided the story is over\n",
    "                break\n",
    "        return idx\n",
    "\n",
    "model = GPT(cfg).to(device)\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"GPT with {cfg.n_layer} blocks, {cfg.n_head} heads, {cfg.n_embd}-dim vectors -> {n_params/1e6:.1f}M parameters\")\n",
    "print(\"(for scale: GPT-2 small is 124M, Llama-3 8B is 8,000M)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "246a4e91",
   "metadata": {},
   "source": [
    "**Four things in that code we have not talked about.**\n",
    "\n",
    "- `dropout`: during lessons, randomly ignore a few dials so Momo does not lean on any one of\n",
    "  them. Ours is 0, switched off: a tiny model with lots of stories does not need it.\n",
    "- `nn.GELU()`: a gentle bend between the two trays of the feed-forward network. Without a\n",
    "  bend, two trays in a row would be just one bigger tray, and the twig could only multiply.\n",
    "- *Weight tying*, `self.head.weight = self.tok_emb.weight`: the meaning map and the scoring\n",
    "  tray are the same tray used twice (once to go from token to meaning, once to go back), so\n",
    "  Momo needs fewer dials.\n",
    "- `std=0.02`: at the start every dial is set to a small random number. Small, so the first\n",
    "  guesses are gentle; random, so the dials do not all learn the same thing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31bb615e",
   "metadata": {},
   "source": [
    "## 7. Training: Ella's bedtime lessons\n",
    "\n",
    "Every lesson is the same five moves Momo used for the bigram, just bigger:\n",
    "1. Ella picks a **batch** of story pieces (32 windows of 256 tokens) and, for each, the same\n",
    "   window shifted by one token: the right answers.\n",
    "2. Momo guesses the next token at every position.\n",
    "3. Ella scores the guesses with the **loss** (ours is **cross-entropy**: \"how surprised were\n",
    "   you by the real next word?\").\n",
    "4. The blame game runs backwards: every dial gets its gradient.\n",
    "5. The **optimizer** turns the dials. Ours is **AdamW**, which also keeps a memory of recent\n",
    "   slopes, like Momo remembering his last few readings on the foggy hill. The W is for\n",
    "   **weight decay**: every lesson Ella also nudges every dial a hair back toward zero, so no\n",
    "   dial grows wild.\n",
    "\n",
    "Two words you will see in the code: `model.train()` is lesson mode and `model.eval()` is quiz\n",
    "mode (on our model they behave the same, since dropout is off, but it is good manners).\n",
    "\n",
    "Two extras every real training run uses:\n",
    "- a **learning-rate schedule**: tiny steps first (**warm-up**, the ground is unknown), full\n",
    "  stride in the middle, shrinking steps at the end (**cosine decay**) so Momo lands in the\n",
    "  grove instead of bouncing;\n",
    "- **mixed precision** on NVIDIA GPUs: do most sums with shorter numbers, about 2x faster.\n",
    "\n",
    "About 1,500 lessons takes about two minutes on a T4 or an M-series Mac. Watch the loss fall,\n",
    "on the lessons and on Ella's hidden quiz.\n",
    "\n",
    "Step through the loop below first to see the five moves and a loss curve drawing itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0edd20df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive: the training loop, stage by stage, with a loss curve.\n",
    "from IPython.display import HTML, display\n",
    "display(HTML(r\"\"\"<div id=\"vr\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
    "<style>#vr button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}#vr .st{display:inline-block;padding:8px 10px;margin:3px;border-radius:8px;border:2px solid #555;background:#2a2a2a;font-size:13px;min-width:120px;text-align:center;transition:all .2s}#vr .st.on{border-color:#ffb347;background:#3a2e1a}#vr .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px;min-height:48px}</style>\n",
    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">The training loop. Step through it and watch the loss fall.</div>\n",
    "<div id=\"vrStages\">\n",
    "<span class=\"st\" data-i=\"0\">1. Ella picks a batch<br><small>32 random story windows</small></span><span style=\"color:#888\">&#8594;</span>\n",
    "<span class=\"st\" data-i=\"1\">2. Momo guesses<br><small>the next token, everywhere</small></span><span style=\"color:#888\">&#8594;</span>\n",
    "<span class=\"st\" data-i=\"2\">3. Ella scores<br><small>the loss: how wrong, on average</small></span><span style=\"color:#888\">&#8594;</span>\n",
    "<span class=\"st\" data-i=\"3\">4. blame game<br><small>autograd: a gradient per dial</small></span><span style=\"color:#888\">&#8594;</span>\n",
    "<span class=\"st\" data-i=\"4\">5. optimizer step<br><small>turn all 5.9M dials</small></span><span style=\"color:#888\">&#8617;</span>\n",
    "</div>\n",
    "<div style=\"display:flex;gap:16px;flex-wrap:wrap;margin-top:10px\"><svg id=\"vrSvg\" width=\"460\" height=\"200\" style=\"background:#262626;border-radius:8px\"></svg>\n",
    "<div style=\"flex:1;min-width:240px\"><div class=\"cap\" id=\"vrCap\">Press Step, or Run 20 steps.</div>\n",
    "<div style=\"margin-top:8px\"><button onclick=\"vrStep()\">Step</button><button onclick=\"vrRun()\">Run 20 steps</button><button onclick=\"vrReset()\">Reset</button></div>\n",
    "<div style=\"font-size:12px;color:#aaa;margin-top:6px\">step <b id=\"vrN\">0</b> &nbsp; loss <b id=\"vrL\">-</b></div></div></div>\n",
    "</div>\n",
    "<script>\n",
    "(function(){\n",
    "const txt=[\"<b>Ella picks a batch.</b> 32 random 256-token windows from the stories, plus the same windows shifted one token right: the right answers.\",\"<b>Momo guesses.</b> The batch runs through Momo. Out come scores for every position's next token.\",\"<b>Ella scores.</b> Cross-entropy compares the scores with the true next tokens. One number: average surprise. Lower is better.\",\"<b>The blame game.</b> Autograd reads the tape backwards and gives every one of the 5.9 million dials its gradient: 'turn me this way to lower the loss'.\",\"<b>Turn the dials.</b> AdamW moves every dial a tiny bit against its gradient. Then the loop repeats, 1,500 times.\"];\n",
    "let stage=-1,n=0,pts=[];const svg=document.getElementById('vrSvg');let timer=null;\n",
    "function loss(k){return 1.9+6.5*Math.exp(-k/6)+ (Math.sin(k*1.7)*0.12);}\n",
    "function draw(){const show=pts.slice(-60);let d='';show.forEach((l,i)=>{d+=(i?'L':'M')+(30+i*(410/60)).toFixed(1)+','+(180-(l-1.5)/7.5*160).toFixed(1);});\n",
    " const ticks=[2,4,6,8].map(v=>`<text x=\"26\" y=\"${(180-(v-1.5)/7.5*160+4).toFixed(1)}\" fill=\"#777\" font-size=\"10\" text-anchor=\"end\">${v}</text><line x1=\"30\" x2=\"440\" y1=\"${(180-(v-1.5)/7.5*160).toFixed(1)}\" y2=\"${(180-(v-1.5)/7.5*160).toFixed(1)}\" stroke=\"#333\"/>`).join('');\n",
    " svg.innerHTML=ticks+`<text x=\"240\" y=\"196\" fill=\"#aaa\" font-size=\"11\" text-anchor=\"middle\">training steps${pts.length>60?' (last 60 shown)':''}</text><text x=\"12\" y=\"100\" fill=\"#aaa\" font-size=\"11\" transform=\"rotate(-90 12,100)\" text-anchor=\"middle\">loss</text><path d=\"${d}\" fill=\"none\" stroke=\"#6fcf97\" stroke-width=\"2.5\"/>`+(show.length?`<circle cx=\"${30+(show.length-1)*(410/60)}\" cy=\"${180-(show[show.length-1]-1.5)/7.5*160}\" r=\"4\" fill=\"#ffb347\"/>`:'');\n",
    " document.querySelectorAll('#vr .st').forEach(s=>s.classList.toggle('on',+s.dataset.i===stage));\n",
    " document.getElementById('vrCap').innerHTML=stage<0?'Press Step, or Run 20 steps.':txt[stage];document.getElementById('vrN').textContent=n;document.getElementById('vrL').textContent=pts.length?pts[pts.length-1].toFixed(2):'-';}\n",
    "window.vrStep=function(){stage=(stage+1)%5;if(stage===2){n++;pts.push(loss(n));}draw();};\n",
    "window.vrRun=function(){if(timer)return;let c=0;timer=setInterval(()=>{vrStep();if(++c>=100){clearInterval(timer);timer=null;}},120);};\n",
    "window.vrReset=function(){clearInterval(timer);timer=null;stage=-1;n=0;pts=[];draw();};\n",
    "draw();\n",
    "})();\n",
    "</script>\n",
    "\"\"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed19704",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "MAX_STEPS = 1500\n",
    "EVAL_EVERY = 250\n",
    "LR = 1e-3\n",
    "WARMUP = 100\n",
    "\n",
    "def get_batch(split):\n",
    "    d = train_data if split == \"train\" else val_data\n",
    "    ix = torch.randint(len(d) - cfg.block_size - 1, (BATCH_SIZE,))\n",
    "    x = torch.stack([d[i:i+cfg.block_size] for i in ix])\n",
    "    y = torch.stack([d[i+1:i+cfg.block_size+1] for i in ix])\n",
    "    return x.to(device), y.to(device)\n",
    "\n",
    "@torch.no_grad()\n",
    "def estimate_loss(iters=20):   # Ella's quiz: average loss over 20 batches of lesson text and 20 of hidden text\n",
    "    model.eval()\n",
    "    out = {}\n",
    "    for split in (\"train\", \"val\"):\n",
    "        losses = torch.zeros(iters)\n",
    "        for i in range(iters):\n",
    "            xb, yb = get_batch(split)\n",
    "            _, loss = model(xb, yb)\n",
    "            losses[i] = loss.item()\n",
    "        out[split] = losses.mean().item()\n",
    "    model.train()\n",
    "    return out\n",
    "\n",
    "def lr_at(step):\n",
    "    if step < WARMUP:\n",
    "        return LR * (step + 1) / WARMUP\n",
    "    progress = (step - WARMUP) / max(1, MAX_STEPS - WARMUP)\n",
    "    return 0.1 * LR + 0.9 * LR * 0.5 * (1 + math.cos(math.pi * progress))\n",
    "\n",
    "# betas: how long AdamW's memory of recent slopes is. weight_decay: the W. (Real runs exempt LayerNorm and biases from decay; we keep it simple.)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1, betas=(0.9, 0.95))\n",
    "use_amp = device == \"cuda\"\n",
    "scaler = torch.amp.GradScaler(enabled=use_amp)\n",
    "\n",
    "history = []\n",
    "t0 = time.time()\n",
    "model.train()\n",
    "for step in range(MAX_STEPS + 1):\n",
    "    for g in optimizer.param_groups:\n",
    "        g[\"lr\"] = lr_at(step)\n",
    "    if step % EVAL_EVERY == 0 or step == MAX_STEPS:\n",
    "        losses = estimate_loss()\n",
    "        history.append((step, losses[\"train\"], losses[\"val\"]))\n",
    "        print(f\"step {step:5d} | lesson loss {losses['train']:.3f} | quiz loss {losses['val']:.3f} \"\n",
    "              f\"| perplexity {math.exp(losses['val']):6.1f} | {(time.time()-t0)/60:4.1f} min\")\n",
    "    if step == MAX_STEPS:\n",
    "        break\n",
    "    xb, yb = get_batch(\"train\")\n",
    "    with torch.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=use_amp):\n",
    "        _, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    scaler.scale(loss).backward()\n",
    "    scaler.unscale_(optimizer)\n",
    "    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # stops rare huge gradients from wrecking training\n",
    "    scaler.step(optimizer)\n",
    "    scaler.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce14f79b",
   "metadata": {},
   "source": [
    "**How to read the numbers.** The loss is Ella's average \"how surprised were you\" per token.\n",
    "**Perplexity** is `exp(loss)`: roughly, \"between how many tokens is Momo still choosing?\"\n",
    "Random guessing over 4,096 tokens would be perplexity 4,096. Around 10 or below on this data\n",
    "gives real sentences, and our run lands right about there. The quiz loss comes from the 200\n",
    "stories Momo never saw, with no Momo stories in them, so it measures real learning rather\n",
    "than memorised repeats. If the lesson loss kept falling while the quiz loss rose, Momo would\n",
    "be memorising, not learning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13ba6e8b",
   "metadata": {},
   "source": [
    "## 8. Momo tells a story: the use case\n",
    "\n",
    "Give Momo the first few words and let him continue, one token at a time, as in the live flow\n",
    "diagram at the top (except that our short `generate` redoes the whole journey each step\n",
    "instead of using Ella's shelf). Two knobs control how he picks each token:\n",
    "- **temperature**: very low means he nearly always takes the safest token (dull but\n",
    "  sensible), high means he gambles on unlikely ones (creative but risky);\n",
    "- **top-k**: he only ever picks from the *k* most likely tokens, so the wildest ones are\n",
    "  never chosen.\n",
    "\n",
    "Run it a few times; each run is a fresh story. Try your own prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b00a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tell_story(prompt, max_new_tokens=200, temperature=0.8):\n",
    "    model.eval()\n",
    "    ids = [tokenizer.eos_token_id] + tokenizer(prompt)[\"input_ids\"]   # start with the story marker, as every training story did\n",
    "    idx = torch.tensor([ids], device=device)\n",
    "    out = model.generate(idx, max_new_tokens, temperature=temperature, eos_id=tokenizer.eos_token_id)\n",
    "    return tokenizer.decode(out[0].tolist(), skip_special_tokens=True)\n",
    "\n",
    "for prompt in [\"Once upon a time, a little monkey\", \"Ella had a red balloon.\", \"One day, Momo went to the river\"]:\n",
    "    print(tell_story(prompt))\n",
    "    print(\"-\" * 60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc77f782",
   "metadata": {},
   "source": [
    "**Look closely at what we saw.** The Momo stories are almost word for word the sentences we\n",
    "fed him three times each, sometimes with the slots mixed up (he looks for a coconut and finds\n",
    "a feather). That is memorising, not inventing: a few thousand repeats of the same sentences\n",
    "are easy for 5.9 million dials to store. The stories without Momo wander, forget who is\n",
    "speaking, and sometimes change the character's name halfway. Both are honest signs of a tiny\n",
    "model after 1,500 lessons. More stories and more lessons (see the ideas at the end) make him\n",
    "invent more and recite less."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24859a73",
   "metadata": {},
   "source": [
    "## 9. Put Momo's book in the library: the Hugging Face Hub\n",
    "\n",
    "The **Hugging Face Hub** is the public library of models and datasets. A model there is just\n",
    "three things: its **weights** (Momo's dials, the tensors we trained), its **config** (the\n",
    "shape of his head: how many layers, heads and so on), and its **tokenizer** (how he bites).\n",
    "We save all three to a folder, write a short **model card** (the book's cover page: what it\n",
    "is, how it was made, how to use it), and upload the folder as a **repo** (a folder on the\n",
    "Hub) under your account.\n",
    "\n",
    "You need a Hugging Face **access token** with *write* permission (Settings > Access Tokens).\n",
    "This \"token\" has nothing to do with text tokens: it is a secret key, like a library card. On\n",
    "Colab the easiest way is the **Secrets** panel (the key icon on the left): add a secret named\n",
    "`HF_TOKEN`, switch on notebook access, and the cell finds it by itself. Otherwise `login()`\n",
    "asks for it. Locally, put `HF_TOKEN=...` in your shell environment.\n",
    "\n",
    "We create the repo **private** first, so only you can see it. Flip `private=False` in the\n",
    "code when you are ready to share."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e55d046",
   "metadata": {},
   "outputs": [],
   "source": [
    "REPO_NAME = \"tiny-storyteller\"       # the model will live at huggingface.co/<you>/tiny-storyteller\n",
    "PUSH_TO_HUB = True                   # set False to only save locally\n",
    "\n",
    "save_dir = REPO_NAME\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "torch.save(model.state_dict(), f\"{save_dir}/model.pt\")\n",
    "with open(f\"{save_dir}/config.json\", \"w\") as f:\n",
    "    json.dump({k: getattr(cfg, k) for k in [\"vocab_size\", \"block_size\", \"n_embd\", \"n_head\", \"n_layer\", \"dropout\"]}, f, indent=2)\n",
    "tokenizer.save_pretrained(save_dir)\n",
    "\n",
    "final_val = history[-1][2]\n",
    "card = f'''---\n",
    "license: mit\n",
    "language: en\n",
    "datasets: [roneneldan/TinyStories]\n",
    "tags: [tiny, from-scratch, pytorch, educational, text-generation]\n",
    "---\n",
    "# Tiny Storyteller\n",
    "\n",
    "A {n_params/1e6:.1f}M-parameter GPT-style language model trained **from scratch** (random init, no pretrained weights)\n",
    "on {len(train_stories) - 3*len(momo):,} TinyStories plus {3*len(momo):,} short synthetic stories about two characters, Momo the monkey and\n",
    "Ella the elephant ({len(momo):,} unique stories, each repeated 3 times, generated by `momo_stories.py`), as a learning project.\n",
    "Use case: a bedtime story starter for young kids.\n",
    "\n",
    "- Architecture: {cfg.n_layer} transformer blocks, {cfg.n_head} heads, {cfg.n_embd}-dim embeddings, {cfg.block_size}-token context\n",
    "- Tokenizer: byte-level BPE, {vocab_size:,} tokens, trained on the same data\n",
    "- Training: {MAX_STEPS} steps, batch {BATCH_SIZE}, AdamW, cosine schedule\n",
    "- Final validation loss {final_val:.3f} (perplexity {math.exp(final_val):.1f}) on {len(val_stories)} held-out TinyStories (no Momo stories in the validation set)\n",
    "\n",
    "The full build, explained step by step, is in `tiny_llm_from_scratch.ipynb` in this repo.\n",
    "Load with plain PyTorch: `GPT(cfg).load_state_dict(torch.load(\"model.pt\"))`.\n",
    "'''\n",
    "with open(f\"{save_dir}/README.md\", \"w\") as f:\n",
    "    f.write(card)\n",
    "print(\"saved:\", sorted(os.listdir(save_dir)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cacb9a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "if PUSH_TO_HUB:\n",
    "    from huggingface_hub import HfApi, login, whoami, get_token\n",
    "    if get_token() is None:                        # no HF_TOKEN env var, no Colab secret, no saved login\n",
    "        login()                                    # ask for the access token (nothing is printed)\n",
    "    api = HfApi()\n",
    "    user = whoami()[\"name\"]\n",
    "    repo_id = f\"{user}/{REPO_NAME}\"\n",
    "    api.create_repo(repo_id, repo_type=\"model\", private=True, exist_ok=True)   # private by default; flip to False to share publicly\n",
    "    api.upload_folder(folder_path=save_dir, repo_id=repo_id, repo_type=\"model\")\n",
    "    nb = \"tiny_llm_from_scratch.ipynb\"\n",
    "    if os.path.exists(nb):\n",
    "        api.upload_file(path_or_fileobj=nb, path_in_repo=nb, repo_id=repo_id)\n",
    "    print(\"published at https://huggingface.co/\" + repo_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a4c8d7",
   "metadata": {},
   "source": [
    "## 10. Loading it back later\n",
    "You, on another machine, can rebuild Momo from the Hub with the `GPT` class from Part 6 and\n",
    "these few lines. The biter (tokenizer) comes along too. While the repo is private, only your\n",
    "own access token can open it; once you set `private=False`, anyone can."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59455da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "def load_from_hub(repo_id):\n",
    "    with open(hf_hub_download(repo_id, \"config.json\")) as f:\n",
    "        c = json.load(f)\n",
    "    class C: pass\n",
    "    for k, v in c.items(): setattr(C, k, v)\n",
    "    m = GPT(C()).to(device)\n",
    "    m.load_state_dict(torch.load(hf_hub_download(repo_id, \"model.pt\"), map_location=device))\n",
    "    tok = PreTrainedTokenizerFast.from_pretrained(repo_id)   # the biter travels with the dials\n",
    "    return m.eval(), tok\n",
    "\n",
    "# Example (uncomment after publishing):\n",
    "# model, tokenizer = load_from_hub(\"singhpratech/tiny-storyteller\")\n",
    "# print(tell_story(\"Once upon a time, a little monkey\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bef94f79",
   "metadata": {},
   "source": [
    "## What we built, in one paragraph\n",
    "Ella's stories became numbers through a tokenizer Momo learned himself. A tiny transformer\n",
    "turned those numbers into spots on a meaning map, let each token look around the story\n",
    "circle at the ones before it, and guessed the next token. Ella measured how wrong he was with\n",
    "cross-entropy, the blame game (autograd) found every dial's gradient, and AdamW turned the\n",
    "dials, thousands of times. That is the entire recipe. Big models use the same recipe with more\n",
    "stories, more branches, and more GPUs.\n",
    "\n",
    "## Ideas to try next\n",
    "- Raise `N_STORIES` and `MAX_STEPS` for a noticeably better storyteller.\n",
    "- Change `n_layer` or `n_embd` and watch how parameters, speed and loss trade off.\n",
    "- Lower `temperature` for safer stories, raise it for stranger ones.\n",
    "- Replace TinyStories with your own text (poems, letters, your journal) and the same code\n",
    "  learns to write in that style.\n",
    "- Make Momo invent instead of recite: more unique Momo stories (add scenes to the pools) and\n",
    "  fewer repeats of each."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}