Upload tiny_llm_from_scratch.ipynb with huggingface_hub

tiny_llm_from_scratch.ipynb

@@ -0,0 +1,769 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e2785f32",
+   "metadata": {},
+   "source": [
+    "# Tiny Storyteller: building a language model from scratch\n",
+    "\n",
+    "**Author:** singhpratech, with Claude. **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, no fine-tuning, no magic. We start with raw PyTorch tensors,\n",
+    "build the simplest possible model, then add one idea at a time until we have a small\n",
+    "GPT-style transformer that we train ourselves.\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 fox\"*, and it continues with a short, simple story.\n",
+    "\n",
+    "## The road map\n",
+    "| Part | What we build | Why |\n",
+    "|---|---|---|\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 and share the model |\n",
+    "\n",
+    "Run every cell top to bottom. Each cell is small on purpose."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5a109db5",
+   "metadata": {},
+   "source": [
+    "## 0. Setup\n",
+    "On Colab, go to **Runtime > Change runtime type > T4 GPU** first. Then run this cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00b4393c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install -q torch transformers datasets tokenizers huggingface_hub"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e97e5f1c",
+   "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)\n",
+    "print(\"PyTorch\", torch.__version__, \"| device:\", device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5eecb952",
+   "metadata": {},
+   "source": [
+    "## 1. Tensors and gradients: the two building blocks\n",
+    "\n",
+    "A **tensor** is just a grid of numbers with a shape. A vector is 1-D, a matrix is 2-D,\n",
+    "a batch of sentences is 3-D. Everything in a neural network (the input text, the weights,\n",
+    "the output probabilities) is a tensor."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3f77d7e0",
+   "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": "2054c860",
+   "metadata": {},
+   "source": [
+    "A **gradient** tells us how much the loss would change if we nudged a weight.\n",
+    "PyTorch computes gradients for us with *autograd*: we mark a tensor with\n",
+    "`requires_grad=True`, do some math, call `.backward()`, and the gradient appears in `.grad`.\n",
+    "\n",
+    "Training a model is only this, repeated millions of times: compute a loss, get the gradients,\n",
+    "move every weight a tiny step against its gradient."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d5918ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "w = torch.tensor([1.0, -1.0, 0.5], requires_grad=True)   # a \"weight\" we want to learn\n",
+    "loss = ((x @ w) ** 2).sum()                              # some loss that depends on w\n",
+    "loss.backward()                                          # autograd fills in w.grad\n",
+    "print(\"loss:\", loss.item())\n",
+    "print(\"d(loss)/d(w):\", w.grad)\n",
+    "\n",
+    "# One step of gradient descent: move w a little bit against the gradient.\n",
+    "with torch.no_grad():\n",
+    "    w -= 0.01 * w.grad\n",
+    "print(\"w after one step:\", w)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4c3d59ca",
+   "metadata": {},
+   "source": [
+    "## 2. Data: TinyStories\n",
+    "\n",
+    "[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) is a public dataset of\n",
+    "short children's stories written in very simple English. It was created specifically to show\n",
+    "that *small* models can learn to write coherent text 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": "6af6bd64",
+   "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": "ea6277bf",
+   "metadata": {},
+   "source": [
+    "## 3. The simplest language model: a bigram\n",
+    "\n",
+    "A language model answers one question: **given what came before, what comes next?**\n",
+    "\n",
+    "The simplest version looks only at the *last character* and keeps a table of\n",
+    "\"after character A, how likely is character B\". That is a **bigram model**.\n",
+    "It is just one `nn.Embedding` table of size *vocab x vocab*.\n",
+    "\n",
+    "First we need a vocabulary. For now, one token = one character."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a29bbc9",
+   "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": "code",
+   "execution_count": null,
+   "id": "9990f156",
+   "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)\n",
+    "            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)\n",
+    "        return idx\n",
+    "\n",
+    "char_data = torch.tensor(encode_c(text[:2_000_000]), dtype=torch.long)\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)\n",
+    "\n",
+    "for step in range(2000):\n",
+    "    xb, yb = get_char_batch()\n",
+    "    _, loss = bigram(xb, yb)\n",
+    "    opt.zero_grad(); loss.backward(); opt.step()\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": "658f7e72",
+   "metadata": {},
+   "source": [
+    "It learned *something* (the loss dropped, and it produces word-ish shapes with spaces),\n",
+    "but it is gibberish. It only sees one character back. To write a story, the model needs\n",
+    "two upgrades: **better tokens** and a way to **look further back**. That is the rest of the notebook."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "763829c4",
+   "metadata": {},
+   "source": [
+    "## 4. A real tokenizer: byte-pair encoding (BPE)\n",
+    "\n",
+    "Characters are too small; the model wastes effort learning to spell. Whole words are too\n",
+    "many (and unknown words break). BPE is the middle ground: it starts from characters and\n",
+    "repeatedly merges the most common pairs, so frequent words become single tokens and rare\n",
+    "words split into pieces. Every big model (GPT, Llama, Qwen) uses some form of this.\n",
+    "\n",
+    "We train our own tokenizer on the stories, with a small vocabulary of 4,096 tokens."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8bcb2397",
+   "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])\n",
+    "\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, a little fox 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": "code",
+   "execution_count": null,
+   "id": "23cb929f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Turn all stories into one long stream of token ids, with EOS between stories.\n",
+    "ids = tokenizer(EOS.join(stories) + EOS)[\"input_ids\"]\n",
+    "data = torch.tensor(ids, dtype=torch.long)\n",
+    "n = int(0.99 * len(data))\n",
+    "train_data, val_data = data[:n], data[n:]\n",
+    "print(f\"total tokens: {len(data)/1e6:.2f}M | train {len(train_data)/1e6:.2f}M | val {len(val_data)/1e3:.0f}K\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2e9b068",
+   "metadata": {},
+   "source": [
+    "## 5. Attention, by hand\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 pay attention to each one.\n",
+    "\n",
+    "Every token produces three vectors:\n",
+    "- a **query** (\"what am I looking for?\"),\n",
+    "- a **key** (\"what do I contain?\"),\n",
+    "- a **value** (\"what do I pass on if someone attends to me?\").\n",
+    "\n",
+    "A token's output is a weighted average of the *values* of earlier tokens, where the weights\n",
+    "come from how well its *query* matches their *keys*. The weights are made with a softmax,\n",
+    "and a triangular mask stops tokens from peeking at the future (we are predicting it!)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c68c6f5c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "B, T, C = 1, 4, 8                      # batch, time (tokens), channels (features)\n",
+    "x_demo = torch.randn(B, T, C)\n",
+    "head = 4\n",
+    "q = nn.Linear(C, head, bias=False)(x_demo)\n",
+    "k = nn.Linear(C, head, bias=False)(x_demo)\n",
+    "v = nn.Linear(C, head, bias=False)(x_demo)\n",
+    "\n",
+    "scores = q @ k.transpose(-2, -1) / math.sqrt(head)       # (B, T, T): how much token i matches token j\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)\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": "0da003d3",
+   "metadata": {},
+   "source": [
+    "## 6. A tiny GPT\n",
+    "\n",
+    "Now the full model. It is the same attention from above, used many times:\n",
+    "\n",
+    "1. **Token embedding**: each token id becomes a vector of `n_embd` numbers.\n",
+    "2. **Position embedding**: each position (0, 1, 2, ...) also gets a vector, so the model knows word order.\n",
+    "3. **Blocks** (repeated `n_layer` times). Each block has:\n",
+    "   - multi-head attention (several attention \"heads\" in parallel, each looking for different things),\n",
+    "   - a feed-forward network (a small MLP applied to every token on its own),\n",
+    "   - layer norm and residual connections, which keep training stable.\n",
+    "4. **Final layer norm + output head**: turn each vector into scores over the vocabulary.\n",
+    "\n",
+    "We also write a `generate` method: predict the next token, append it, repeat."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b00876c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Config:\n",
+    "    vocab_size = vocab_size\n",
+    "    block_size = 256     # how many tokens the model 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):\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 next_id.item() == tokenizer.eos_token_id:   # 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": "9c446a37",
+   "metadata": {},
+   "source": [
+    "## 7. Training\n",
+    "\n",
+    "The loop is the same one we used for the bigram, just bigger:\n",
+    "1. grab a random batch of token windows, and the same windows shifted by one (the targets),\n",
+    "2. run the model, get the loss,\n",
+    "3. backward, step the optimizer.\n",
+    "\n",
+    "Two small extras that every real training run uses:\n",
+    "- a **learning-rate schedule** (warm up, then cosine decay) so the start is gentle and the end is precise,\n",
+    "- **mixed precision** on NVIDIA GPUs, which is roughly 2x faster.\n",
+    "\n",
+    "About 1,500 steps takes a few minutes on a T4 or an M-series Mac. Watch the loss fall."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "222e1c51",
+   "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):\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",
+    "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} | train loss {losses['train']:.3f} | val 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": "4dee51b4",
+   "metadata": {},
+   "source": [
+    "**How to read the numbers.** The loss is the average \"surprise\" per token. *Perplexity* is\n",
+    "`exp(loss)`: roughly, \"between how many tokens is the model guessing\". A random guess over\n",
+    "4,096 tokens would be perplexity 4,096. Anything under 10 on this data means real sentences."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "284708da",
+   "metadata": {},
+   "source": [
+    "## 8. Generate stories: the use case\n",
+    "\n",
+    "Give the model the first few words of a story and let it continue.\n",
+    "Run it a few times; each run is a fresh story. Try your own prompts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ce19e9d9",
+   "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",
+    "    out = model.generate(idx, max_new_tokens, temperature=temperature)\n",
+    "    return tokenizer.decode(out[0].tolist(), skip_special_tokens=True)\n",
+    "\n",
+    "for prompt in [\"Once upon a time, a little fox\", \"Lily had a red balloon.\", \"One day, Tom went to the park\"]:\n",
+    "    print(tell_story(prompt))\n",
+    "    print(\"-\" * 60)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3a6cfad",
+   "metadata": {},
+   "source": [
+    "## 9. Save and publish on the Hugging Face Hub\n",
+    "\n",
+    "A model is just its **weights** (the tensors we trained), its **config** (the shape), and its\n",
+    "**tokenizer**. We save all three to a folder, write a short model card, and upload the folder\n",
+    "as a repo under your account.\n",
+    "\n",
+    "You need a Hugging Face token with *write* permission (Settings > Access Tokens).\n",
+    "On Colab, `login()` will ask for it. Locally, put `HF_TOKEN=...` in your environment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5fc0838",
+   "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",
+    "json.dump({k: getattr(cfg, k) for k in [\"vocab_size\", \"block_size\", \"n_embd\", \"n_head\", \"n_layer\", \"dropout\"]},\n",
+    "          open(f\"{save_dir}/config.json\", \"w\"), 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 {N_STORIES:,} TinyStories, as a learning project. 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})\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",
+    "open(f\"{save_dir}/README.md\", \"w\").write(card)\n",
+    "print(\"saved:\", sorted(os.listdir(save_dir)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7e1ae430",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if PUSH_TO_HUB:\n",
+    "    from huggingface_hub import HfApi, login, whoami\n",
+    "    if not os.environ.get(\"HF_TOKEN\"):\n",
+    "        login()                                    # on Colab this opens a token prompt\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": "34f8cabf",
+   "metadata": {},
+   "source": [
+    "## Loading it back later\n",
+    "Anyone (including you, on another machine) can rebuild the model from the Hub with the\n",
+    "`GPT` class from Part 6 and these few lines:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5e86850",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import hf_hub_download\n",
+    "\n",
+    "def load_from_hub(repo_id):\n",
+    "    c = json.load(open(hf_hub_download(repo_id, \"config.json\")))\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",
+    "    return m.eval()\n",
+    "\n",
+    "# Example (uncomment after publishing):\n",
+    "# model = load_from_hub(\"singhpratech/tiny-storyteller\")\n",
+    "# print(tell_story(\"Once upon a time, a little fox\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0f44cd1",
+   "metadata": {},
+   "source": [
+    "## What we built, in one paragraph\n",
+    "Text became numbers through a tokenizer we trained ourselves. A tiny transformer turned\n",
+    "those numbers into vectors, let each token attend to the ones before it, and predicted the\n",
+    "next token. We measured how wrong it was with cross-entropy, used autograd to get gradients,\n",
+    "and nudged every weight with AdamW, thousands of times. That is the entire recipe. Big\n",
+    "models use the same recipe with more data, more layers, 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 (recipes, poems, your journal) and the same code\n",
+    "  learns to write in that style."
+   ]
+  }
+ ],
+ "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
+}