Upload tiny_llm_from_scratch.ipynb with huggingface_hub

tiny_llm_from_scratch.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "9ff84951",
+   "id": "b2fe09d4",
    "metadata": {},
    "source": [
     "# Tiny Storyteller: building a language model from scratch\n",
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3fcb1915",
+   "id": "c4f48f79",
    "metadata": {},
    "source": [
     "# Before any code: Momo, Ella, and what we are about to build\n",
@@ -67,7 +67,8 @@
     "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",
-    "Each of these gets its own section below.\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.\n",
     "\n",
     "## 1. What is PyTorch?\n",
     "\n",
@@ -82,8 +83,8 @@
     "  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. No ordinary programming tool does this, and it is the reason\n",
-    "  modern AI works at all. Part 1 has a game that shows exactly how the tape is read.\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",
@@ -138,7 +139,7 @@
     "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 \"I'm done\" token. You see words appear one by one\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.\n",
     "\n",
     "## 5. How attention looks things up: the story circle\n",
@@ -150,7 +151,9 @@
     "- 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",
-    "The newest animal holds up its query card and compares it with everyone's key card. A good\n",
+    "Every animal holds up its query card and compares it with the key cards of the animals\n",
+    "before it (during training all of them do this at the same time; when writing, only the\n",
+    "newest one needs to). A good\n",
     "match scores high, a poor match scores low. The scores are squashed into percentages\n",
     "(softmax again) and the animal collects a blend of everyone's value cards, weighted by those\n",
     "percentages. \"capital\" might take 70 percent from \"France\", 20 percent from \"of\", and crumbs\n",
@@ -180,7 +183,10 @@
     "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 twelve lines of code.\n",
+    "Part 5 shows the shelf in twelve 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.\n",
     "\n",
     "## 7. The strange names, decoded\n",
     "\n",
@@ -233,7 +239,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a93c6c86",
+   "id": "4d7312e3",
    "metadata": {},
    "source": [
     "## Live flow diagram\n",
@@ -248,7 +254,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c19f30be",
+   "id": "cdae01ca",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -373,7 +379,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "754198a6",
+   "id": "94499493",
    "metadata": {},
    "source": [
     "## 0. Setup\n",
@@ -383,7 +389,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "68b59dbf",
+   "id": "0fe52b75",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -393,7 +399,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1d0052fe",
+   "id": "cf74d97b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -410,7 +416,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "511c964a",
+   "id": "30cd04b6",
    "metadata": {},
    "source": [
     "## 1. Tensors and gradients: the two building blocks\n",
@@ -429,7 +435,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "aa29df50",
+   "id": "7fd75d97",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -472,7 +478,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3717ba99",
+   "id": "6d9343cc",
    "metadata": {},
    "source": [
     "### How matrix multiplication actually happens\n",
@@ -494,7 +500,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ba02b9f9",
+   "id": "88849be7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -551,7 +557,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "00afbdb2",
+   "id": "c1522ad3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -564,7 +570,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9ced2269",
+   "id": "59188275",
    "metadata": {},
    "source": [
     "### The gradient: which way to turn each dial\n",
@@ -587,7 +593,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0a31a332",
+   "id": "4e4aef04",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -627,7 +633,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "69480f68",
+   "id": "b4ec1642",
    "metadata": {},
    "source": [
     "### Momo in the foggy jungle: thorn bushes, marks, and the slope\n",
@@ -640,9 +646,9 @@
     "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. He steps back, marks the spot with a red X, and corrects, say to the left. That\n",
-    "correction is the gradient doing its job: after a bad step the slope under his feet points\n",
-    "back the other way.\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",
@@ -656,8 +662,9 @@
     "- The reading was noisy. Each step uses one batch of stories, 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 and nothing is forgotten.\n",
-    "  Momo does not get a map; after training he *is* the map.\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",
@@ -672,7 +679,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f4cdba75",
+   "id": "73a8941e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -782,7 +789,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0fc03ef8",
+   "id": "a6a84861",
    "metadata": {},
    "source": [
     "### How does Momo feel the slope without walking? The blame game\n",
@@ -805,14 +812,17 @@
     "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**. Press each stage below. The\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 (our residual connections do this), the two blames\n",
+    "are added together. And the game's mistake is a squared gap; the real Momo uses cross-entropy.\n",
+    "Same game, different first step. Press each stage below. The\n",
     "slow-way check lets you see that the backward trip tells the truth."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6c9e4319",
+   "id": "acb78f06",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -913,7 +923,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ec8003d8",
+   "id": "39d72d64",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -931,7 +941,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d1f811c6",
+   "id": "24e1da8f",
    "metadata": {},
    "source": [
     "## 2. Data: Ella's bedtime stories (TinyStories)\n",
@@ -947,7 +957,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b3b653c1",
+   "id": "ebb5d3e6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -965,7 +975,125 @@
   },
   {
    "cell_type": "markdown",
-   "id": "080dcf9d",
+   "id": "b9bf532d",
+   "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 a few hundred short stories about\n",
+    "our two characters, built from sentence pools in the same simple style (every story has an\n",
+    "opening, Ella's introduction, a small problem, Ella helping, a resolution and an ending). We\n",
+    "repeat them three times so they are not drowned out by the 20,000 other stories. This is\n",
+    "exactly how a real model is taught a company's own documents: mix them into the training text."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40a20bc2",
+   "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",
+    "    rng = random.Random(seed)\n",
+    "    return [make_story(rng) for _ in range(n)]\n",
+    "\n",
+    "momo = momo_stories(1000, seed=0)\n",
+    "stories = stories + momo * 3          # 20,000 TinyStories + 3,000 Momo stories (1,000 unique)\n",
+    "import random; random.Random(42).shuffle(stories)\n",
+    "print(f\"{len(stories):,} stories in total, {len(momo):,} unique Momo stories\")\n",
+    "print(\"-\" * 60)\n",
+    "print(momo[0])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b3fc1b7",
    "metadata": {},
    "source": [
     "## 3. Momo's first trick: the bigram\n",
@@ -986,7 +1114,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0309bf44",
+   "id": "d7c45f72",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1017,7 +1145,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "aa0fa0fa",
+   "id": "3a50e68a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1037,7 +1165,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0be6fc04",
+   "id": "942d14ad",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1086,7 +1214,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e7ebebbc",
+   "id": "55fb61d0",
    "metadata": {},
    "source": [
     "Momo learned *something* (the loss dropped, and he produces word-ish shapes with spaces), but\n",
@@ -1096,7 +1224,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1c63ff73",
+   "id": "b287c96b",
    "metadata": {},
    "source": [
     "## 4. A real tokenizer: how Momo decides where to bite\n",
@@ -1121,7 +1249,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "609916e1",
+   "id": "6cc8a788",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1151,7 +1279,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4ca0837a",
+   "id": "9edae39b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1177,7 +1305,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d939e9d0",
+   "id": "07134162",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1191,7 +1319,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "537003ab",
+   "id": "435e5e35",
    "metadata": {},
    "source": [
     "## 5. Attention, by hand: the story circle in numbers\n",
@@ -1213,7 +1341,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dcca4697",
+   "id": "98f928a1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1237,7 +1365,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3d14ba96",
+   "id": "dd9059e3",
    "metadata": {},
    "source": [
     "### Ella's shelf (the KV cache), in twelve lines\n",
@@ -1251,7 +1379,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f3cef3da",
+   "id": "807c906b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1281,7 +1409,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e614ea56",
+   "id": "e7557457",
    "metadata": {},
    "source": [
     "## 6. A tiny GPT: Momo's whole head\n",
@@ -1298,7 +1426,8 @@
     "3. **Blocks**, repeated `n_layer` times (6 branches 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 goes off to its own\n",
-    "     branch and thinks alone (a small two-layer network, where most facts end up stored);\n",
+    "     branch 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",
@@ -1313,7 +1442,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d35b8f49",
+   "id": "40b8d767",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1357,7 +1486,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1fbd25da",
+   "id": "d79f1d74",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1468,7 +1597,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6bad119f",
+   "id": "4e8831e7",
    "metadata": {},
    "source": [
     "## 7. Training: Ella's bedtime lessons\n",
@@ -1497,7 +1626,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5e1db782",
+   "id": "882b3bdb",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1539,7 +1668,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "133e9333",
+   "id": "481c902a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1606,7 +1735,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a9476de0",
+   "id": "944ca70e",
    "metadata": {},
    "source": [
     "**How to read the numbers.** The loss is Ella's average \"how surprised were you\" per token.\n",
@@ -1617,7 +1746,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2e496a79",
+   "id": "387f64a3",
    "metadata": {},
    "source": [
     "## 8. Momo tells a story: the use case\n",
@@ -1635,13 +1764,14 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f2cc11bf",
+   "id": "b545d4c1",
    "metadata": {},
    "outputs": [],
    "source": [
     "def tell_story(prompt, max_new_tokens=200, temperature=0.8):\n",
     "    model.eval()\n",
-    "    idx = torch.tensor([tokenizer(prompt)[\"input_ids\"]], device=device)\n",
+    "    ids = tokenizer(prompt)[\"input_ids\"] or [tokenizer.eos_token_id]   # empty prompt: start from the story marker\n",
+    "    idx = torch.tensor([ids], device=device)\n",
     "    out = model.generate(idx, max_new_tokens, temperature=temperature)\n",
     "    return tokenizer.decode(out[0].tolist(), skip_special_tokens=True)\n",
     "\n",
@@ -1652,7 +1782,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ee01a9ea",
+   "id": "455a433e",
    "metadata": {},
    "source": [
     "## 9. Put Momo's book in the library: the Hugging Face Hub\n",
@@ -1671,7 +1801,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e392a0ac",
+   "id": "00878bab",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1712,7 +1842,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c5250a84",
+   "id": "45371319",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1733,7 +1863,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6db272e7",
+   "id": "9c85c1d9",
    "metadata": {},
    "source": [
     "## Loading it back later\n",
@@ -1744,7 +1874,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0970216e",
+   "id": "dd50896e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1765,7 +1895,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "334dcc17",
+   "id": "c107a885",
    "metadata": {},
    "source": [
     "## What we built, in one paragraph\n",