{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "d24a5a8b",
      "metadata": {},
      "source": [
        "# makemore — a WaveNet\n",
        "\n",
        "Go **deeper and hierarchical** — **concept -> code -> Your turn** each step.\n",
        "\n",
        "This is notebook **05**. The MLP in `02`/`03` squashed the whole context into **one** layer.\n",
        "Here we widen the context to 8 letters and fuse them **gradually**, two at a time, in a\n",
        "tree — the structure behind DeepMind's WaveNet (2016). We also wrap everything in tidy\n",
        "`torch.nn`-style building blocks. Follows Karpathy's *\"makemore part 5: WaveNet\"*."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "187db6ed",
      "metadata": {},
      "source": [
        "## Prologue — in plain English\n",
        "\n",
        "Reading 8 letters by mashing them into one layer throws away structure. WaveNet's idea: fuse\n",
        "information **slowly**.\n",
        "\n",
        "Real-life picture: a knockout tournament. 8 players (letters) play in pairs -> 4 winners ->\n",
        "play in pairs -> 2 -> 1 champion. Each round combines neighbours into a richer summary. Our\n",
        "network does the same: combine letters in pairs, then pairs-of-pairs, then again, until a\n",
        "single vector represents the whole context.\n",
        "\n",
        "We also build small reusable layers (`Linear`, `BatchNorm1d`, `Tanh`, `Embedding`,\n",
        "`FlattenConsecutive`, `Sequential`) that behave like PyTorch's own — so the model reads like a\n",
        "clean list of steps."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6bf209d0",
      "metadata": {},
      "source": [
        "### 0.1 Imports and data (context grows to 8)\n",
        "\n",
        "The only data change from earlier notebooks: `block_size = 8`, so the model sees up to 8\n",
        "previous letters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "88825091",
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
        "\n",
        "import random\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "try:\n",
        "    import matplotlib.pyplot as plt\n",
        "    HAS_PLT = True\n",
        "except Exception:\n",
        "    HAS_PLT = False\n",
        "\n",
        "words = open(\"names.txt\").read().splitlines()\n",
        "chars = sorted(list(set(\"\".join(words))))\n",
        "stoi = {ch: i + 1 for i, ch in enumerate(chars)}\n",
        "stoi[\".\"] = 0\n",
        "itos = {i: ch for ch, i in stoi.items()}\n",
        "vocab_size = len(itos)\n",
        "block_size = 8\n",
        "\n",
        "def build_dataset(word_list):\n",
        "    X, Y = [], []\n",
        "    for w in word_list:\n",
        "        context = [0] * block_size\n",
        "        for ch in w + \".\":\n",
        "            ix = stoi[ch]\n",
        "            X.append(context); Y.append(ix)\n",
        "            context = context[1:] + [ix]\n",
        "    return torch.tensor(X), torch.tensor(Y)\n",
        "\n",
        "random.seed(42)\n",
        "random.shuffle(words)\n",
        "n1, n2 = int(0.8 * len(words)), int(0.9 * len(words))\n",
        "Xtr, Ytr = build_dataset(words[:n1])\n",
        "Xdev, Ydev = build_dataset(words[n1:n2])\n",
        "Xte, Yte = build_dataset(words[n2:])\n",
        "print(\"train:\", tuple(Xtr.shape), \" (context length =\", block_size, \")\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4c8200dd",
      "metadata": {},
      "source": [
        "### 1.1 PyTorch-style building blocks\n",
        "\n",
        "Each layer is a tiny class with a `__call__` (forward) and `parameters()`. The interesting one\n",
        "is **`FlattenConsecutive(n)`**: it groups every `n` neighbouring positions together — that is\n",
        "the \"combine in pairs\" tournament step. `BatchNorm1d` is written to also handle the 3-D tensors\n",
        "that appear between fusion steps (a subtle but important detail)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7695c863",
      "metadata": {},
      "outputs": [],
      "source": [
        "class Linear:\n",
        "    def __init__(self, fan_in, fan_out, bias=True):\n",
        "        self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in ** 0.5\n",
        "        self.bias = torch.zeros(fan_out) if bias else None\n",
        "    def __call__(self, x):\n",
        "        self.out = x @ self.weight\n",
        "        if self.bias is not None:\n",
        "            self.out = self.out + self.bias\n",
        "        return self.out\n",
        "    def parameters(self):\n",
        "        return [self.weight] + ([] if self.bias is None else [self.bias])\n",
        "\n",
        "\n",
        "class BatchNorm1d:\n",
        "    def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
        "        self.eps = eps; self.momentum = momentum; self.training = True\n",
        "        self.gamma = torch.ones(dim); self.beta = torch.zeros(dim)\n",
        "        self.running_mean = torch.zeros(dim); self.running_var = torch.ones(dim)\n",
        "    def __call__(self, x):\n",
        "        if self.training:\n",
        "            dim = 0 if x.ndim == 2 else (0, 1)     # 3-D inputs: average over batch AND time\n",
        "            xmean = x.mean(dim, keepdim=True)\n",
        "            xvar = x.var(dim, keepdim=True)\n",
        "        else:\n",
        "            xmean = self.running_mean\n",
        "            xvar = self.running_var\n",
        "        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)\n",
        "        self.out = self.gamma * xhat + self.beta\n",
        "        if self.training:\n",
        "            with torch.no_grad():\n",
        "                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean\n",
        "                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar\n",
        "        return self.out\n",
        "    def parameters(self):\n",
        "        return [self.gamma, self.beta]\n",
        "\n",
        "\n",
        "class Tanh:\n",
        "    def __call__(self, x):\n",
        "        self.out = torch.tanh(x); return self.out\n",
        "    def parameters(self):\n",
        "        return []\n",
        "\n",
        "\n",
        "class Embedding:\n",
        "    def __init__(self, num_embeddings, embedding_dim):\n",
        "        self.weight = torch.randn((num_embeddings, embedding_dim), generator=g)\n",
        "    def __call__(self, IX):\n",
        "        self.out = self.weight[IX]; return self.out\n",
        "    def parameters(self):\n",
        "        return [self.weight]\n",
        "\n",
        "\n",
        "class FlattenConsecutive:\n",
        "    def __init__(self, n):\n",
        "        self.n = n\n",
        "    def __call__(self, x):\n",
        "        B, T, C = x.shape\n",
        "        x = x.view(B, T // self.n, C * self.n)   # fuse every n neighbours\n",
        "        if x.shape[1] == 1:\n",
        "            x = x.squeeze(1)\n",
        "        self.out = x; return self.out\n",
        "    def parameters(self):\n",
        "        return []\n",
        "\n",
        "\n",
        "class Sequential:\n",
        "    def __init__(self, layers):\n",
        "        self.layers = layers\n",
        "    def __call__(self, x):\n",
        "        for layer in self.layers:\n",
        "            x = layer(x)\n",
        "        self.out = x; return self.out\n",
        "    def parameters(self):\n",
        "        return [p for layer in self.layers for p in layer.parameters()]\n",
        "\n",
        "print(\"layers defined\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ab8cd26d",
      "metadata": {},
      "source": [
        "### 1.2 Build the hierarchical (WaveNet) model\n",
        "\n",
        "`block_size = 8` fuses in three rounds: 8 -> 4 -> 2 -> 1. Each round is\n",
        "`FlattenConsecutive(2) -> Linear -> BatchNorm1d -> Tanh`. Finally a `Linear` maps to the 27\n",
        "next-letter scores. We shrink the last layer's weights so the first loss starts near the ideal\n",
        "`ln(27)`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1dcd57ee",
      "metadata": {},
      "outputs": [],
      "source": [
        "g = torch.Generator().manual_seed(42)\n",
        "n_embd = 24\n",
        "n_hidden = 128\n",
        "\n",
        "model = Sequential([\n",
        "    Embedding(vocab_size, n_embd),\n",
        "    FlattenConsecutive(2), Linear(n_embd * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
        "    FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
        "    FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),\n",
        "    Linear(n_hidden, vocab_size),\n",
        "])\n",
        "\n",
        "with torch.no_grad():\n",
        "    model.layers[-1].weight *= 0.1   # humble output\n",
        "\n",
        "parameters = model.parameters()\n",
        "for p in parameters:\n",
        "    p.requires_grad = True\n",
        "print(\"total parameters:\", sum(p.nelement() for p in parameters))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dfe5653b",
      "metadata": {},
      "source": [
        "### 1.3 Watch the tournament shrink the sequence\n",
        "\n",
        "Run one batch and print each layer's output shape. See the time dimension halve 8 -> 4 -> 2 -> 1\n",
        "as neighbours fuse, while the feature dimension does the work."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5a37517c",
      "metadata": {},
      "outputs": [],
      "source": [
        "ix = torch.randint(0, Xtr.shape[0], (4,), generator=g)\n",
        "_ = model(Xtr[ix])\n",
        "for layer in model.layers:\n",
        "    print(f\"{layer.__class__.__name__:18s} : {tuple(layer.out.shape)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "43949ea9",
      "metadata": {},
      "source": [
        "### 2.1 Train\n",
        "\n",
        "Standard loop: minibatch -> forward -> `loss.backward()` -> nudge. (We use autograd again now;\n",
        "notebook `04` already proved we could do it by hand.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f0800ce0",
      "metadata": {},
      "outputs": [],
      "source": [
        "max_steps = 20000\n",
        "batch_size = 32\n",
        "lossi = []\n",
        "\n",
        "for i in range(max_steps):\n",
        "    ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n",
        "    Xb, Yb = Xtr[ix], Ytr[ix]\n",
        "\n",
        "    logits = model(Xb)\n",
        "    loss = F.cross_entropy(logits, Yb)\n",
        "\n",
        "    for p in parameters:\n",
        "        p.grad = None\n",
        "    loss.backward()\n",
        "    lr = 0.1 if i < 15000 else 0.01\n",
        "    for p in parameters:\n",
        "        p.data += -lr * p.grad\n",
        "\n",
        "    lossi.append(loss.log10().item())\n",
        "    if i % 4000 == 0 or i == max_steps - 1:\n",
        "        print(f\"step {i:5d}/{max_steps}   loss {loss.item():.4f}\")\n",
        "\n",
        "if HAS_PLT:\n",
        "    # smooth the noisy per-step loss by averaging in chunks of 200\n",
        "    lt = torch.tensor(lossi)\n",
        "    plt.figure(figsize=(5, 4))\n",
        "    plt.plot(lt.view(-1, 200).mean(1))\n",
        "    plt.xlabel(\"step / 200\"); plt.ylabel(\"log10(loss)\"); plt.title(\"WaveNet training loss\")\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "053ba8a7",
      "metadata": {},
      "source": [
        "### 2.2 Evaluate\n",
        "\n",
        "Switch BatchNorm layers to eval mode (use running stats), then score train and dev. The deeper\n",
        "hierarchical model should edge below the single-layer MLP's ~2.1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3a5d6228",
      "metadata": {},
      "outputs": [],
      "source": [
        "for layer in model.layers:\n",
        "    if isinstance(layer, BatchNorm1d):\n",
        "        layer.training = False\n",
        "\n",
        "@torch.no_grad()\n",
        "def split_loss(split):\n",
        "    x, y = {\"train\": (Xtr, Ytr), \"dev\": (Xdev, Ydev), \"test\": (Xte, Yte)}[split]\n",
        "    logits = model(x)\n",
        "    return F.cross_entropy(logits, y).item()\n",
        "\n",
        "print(\"train loss:\", round(split_loss(\"train\"), 4))\n",
        "print(\"dev   loss:\", round(split_loss(\"dev\"), 4))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b78eaf70",
      "metadata": {},
      "source": [
        "### 3.1 Generate names\n",
        "\n",
        "Sample letter by letter as before. With 8 letters of context and the hierarchical fusion, the\n",
        "names tend to look a little more name-like."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "33833066",
      "metadata": {},
      "outputs": [],
      "source": [
        "g2 = torch.Generator().manual_seed(2147483647 + 10)\n",
        "\n",
        "for _ in range(20):\n",
        "    out = []\n",
        "    context = [0] * block_size\n",
        "    while True:\n",
        "        logits = model(torch.tensor([context]))\n",
        "        probs = F.softmax(logits, dim=1)\n",
        "        ix = torch.multinomial(probs, num_samples=1, generator=g2).item()\n",
        "        context = context[1:] + [ix]\n",
        "        if ix == 0:\n",
        "            break\n",
        "        out.append(itos[ix])\n",
        "    print(\"\".join(out))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "31d52503",
      "metadata": {},
      "source": [
        "**Your turn 5** — The model is just a list. Add another fusion round (so it handles\n",
        "`block_size = 16`), or raise `n_hidden`, retrain, and compare dev loss. Because everything is a\n",
        "clean `Sequential`, experimenting is just editing the list."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "894cc92e",
      "metadata": {},
      "source": [
        "## What's next\n",
        "\n",
        "You built a deeper, hierarchical model out of reusable PyTorch-style blocks — the same shape of\n",
        "code real libraries use.\n",
        "\n",
        "Where this leads in the series:\n",
        "\n",
        "- `06_build_gpt_attention.ipynb` — the transformer. Instead of a fixed tournament that always\n",
        "  fuses neighbours, **attention** lets every position decide *which* earlier positions to listen\n",
        "  to. That flexibility is what makes GPT powerful.\n",
        "- `07_gpt_tokenizer.ipynb` — how real models turn text into tokens in the first place.\n",
        "\n",
        "**Checklist**\n",
        "- [ ] Context widened to 8 letters\n",
        "- [ ] PyTorch-style layers (Linear, BatchNorm1d, Tanh, Embedding, FlattenConsecutive, Sequential)\n",
        "- [ ] Hierarchical fusion 8 -> 4 -> 2 -> 1\n",
        "- [ ] BatchNorm1d handles 3-D tensors\n",
        "- [ ] Trained, evaluated, sampled names"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv (3.12.10)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}