{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "286c5a48",
   "metadata": {},
   "source": [
    "# Topic 2: The same bigram model in PyTorch\n",
    "\n",
    "We rebuild the Topic 1 bigram name generator almost cell-for-cell, but in **PyTorch**.\n",
    "Keep the Topic 1 notebook open side by side: the forward pass looks nearly identical, but the hand-derived backward pass disappears."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82a108b5",
   "metadata": {},
   "source": [
    "## 1. Setup, random seed and device\n",
    "We import PyTorch and its functional API (`F`), fix the seed (1337), and pick a **device**: a GPU (`\"cuda\"`) if one is available, otherwise the CPU.\n",
    "A separate `torch.Generator` gives us a reproducible random stream for sampling names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d26c437c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:42.759636Z",
     "iopub.status.busy": "2026-10-01T17:09:42.759493Z",
     "iopub.status.idle": "2026-10-01T17:09:44.476942Z",
     "shell.execute_reply": "2026-10-01T17:09:44.476310Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch 2.13.0+cu132 | device: cpu\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "g = torch.Generator().manual_seed(1337)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"torch\", torch.__version__, \"| device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45ab8352",
   "metadata": {},
   "source": [
    "## 2. Load the dataset\n",
    "Reading the names and building the character vocabulary is plain Python, exactly as in Topic 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d82d85c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.479510Z",
     "iopub.status.busy": "2026-10-01T17:09:44.479272Z",
     "iopub.status.idle": "2026-10-01T17:09:44.487260Z",
     "shell.execute_reply": "2026-10-01T17:09:44.486738Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32033 names, 27 tokens\n"
     ]
    }
   ],
   "source": [
    "words = open(\"names.txt\").read().splitlines()\n",
    "chars = sorted(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(stoi)\n",
    "print(len(words), \"names,\", vocab_size, \"tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b11d987",
   "metadata": {},
   "source": [
    "The (current, next) pairs are now stored in a `torch.tensor` instead of a NumPy array.\n",
    "A tensor is PyTorch's n-dimensional array: like a NumPy array it has a `shape` and a `dtype` (here 64-bit integers), but it can also live on a GPU and track gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abd124a8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.489129Z",
     "iopub.status.busy": "2026-10-01T17:09:44.488990Z",
     "iopub.status.idle": "2026-10-01T17:09:44.635662Z",
     "shell.execute_reply": "2026-10-01T17:09:44.635110Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "228146 training pairs | shape: torch.Size([228146]) | dtype: torch.int64\n"
     ]
    }
   ],
   "source": [
    "xs, ys = [], []\n",
    "for w in words:\n",
    "    chs = [\".\"] + list(w) + [\".\"]\n",
    "    for ch1, ch2 in zip(chs, chs[1:]):\n",
    "        xs.append(stoi[ch1])\n",
    "        ys.append(stoi[ch2])\n",
    "xs = torch.tensor(xs)\n",
    "ys = torch.tensor(ys)\n",
    "num = xs.numel()\n",
    "print(num, \"training pairs | shape:\", xs.shape, \"| dtype:\", xs.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c3a9cb",
   "metadata": {},
   "source": [
    "`F.one_hot` does the one-hot encoding in one call. It returns integers, so we convert with `.float()` because matrix multiplication with `W` needs floating-point numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "967a6afc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.637776Z",
     "iopub.status.busy": "2026-10-01T17:09:44.637579Z",
     "iopub.status.idle": "2026-10-01T17:09:44.696263Z",
     "shell.execute_reply": "2026-10-01T17:09:44.695685Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([228146, 27]) torch.float32\n",
      "tensor([1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0.])\n"
     ]
    }
   ],
   "source": [
    "xenc = F.one_hot(xs, num_classes=vocab_size).float()\n",
    "print(xenc.shape, xenc.dtype)\n",
    "print(xenc[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d0741c",
   "metadata": {},
   "source": [
    "## 3. Initialize the model\n",
    "`W` is again a random 27x27 matrix. The one new ingredient is `requires_grad=True`, which tells PyTorch to record every operation involving `W` so it can compute gradients for us later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "090fd9b8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.698151Z",
     "iopub.status.busy": "2026-10-01T17:09:44.698005Z",
     "iopub.status.idle": "2026-10-01T17:09:44.700784Z",
     "shell.execute_reply": "2026-10-01T17:09:44.700398Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([27, 27]) True\n"
     ]
    }
   ],
   "source": [
    "W = torch.randn((vocab_size, vocab_size), requires_grad=True)\n",
    "print(W.shape, W.requires_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f44e8ea",
   "metadata": {},
   "source": [
    "## 4. Forward pass and cross-entropy loss\n",
    "The forward pass is written out exactly as in Topic 1: logits, a numerically stable softmax, and the average negative log-probability of the correct next character.\n",
    "Only the names change slightly (`axis` → `dim`, `keepdims` → `keepdim`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c8db2dc8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.707804Z",
     "iopub.status.busy": "2026-10-01T17:09:44.707661Z",
     "iopub.status.idle": "2026-10-01T17:09:44.775463Z",
     "shell.execute_reply": "2026-10-01T17:09:44.774855Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 3.8462\n"
     ]
    }
   ],
   "source": [
    "logits = xenc @ W                                     # (num, 27)\n",
    "counts = (logits - logits.max(dim=1, keepdim=True).values).exp()\n",
    "probs = counts / counts.sum(dim=1, keepdim=True)      # each row sums to 1\n",
    "loss = -probs[torch.arange(num), ys].log().mean()\n",
    "print(f\"loss: {loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04d107c6",
   "metadata": {},
   "source": [
    "## 5. Backward pass: autograd\n",
    "Instead of deriving the gradient by hand, we call `loss.backward()`. PyTorch walks back through the recorded operations (chain rule, automatically) and stores d(loss)/dW in `W.grad`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "33537801",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.777392Z",
     "iopub.status.busy": "2026-10-01T17:09:44.777243Z",
     "iopub.status.idle": "2026-10-01T17:09:44.885210Z",
     "shell.execute_reply": "2026-10-01T17:09:44.884585Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([27, 27])\n"
     ]
    }
   ],
   "source": [
    "loss.backward()\n",
    "print(W.grad.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "230186b6",
   "metadata": {},
   "source": [
    "Is it right? Let's compute the Topic 1 hand-derived formula `dW = xenc.T @ (probs - onehot(ys)) / num` and compare. `.detach()` gives us `probs` without gradient tracking, and `torch.allclose` checks the two agree up to floating-point rounding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e7bdafac",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.887260Z",
     "iopub.status.busy": "2026-10-01T17:09:44.887115Z",
     "iopub.status.idle": "2026-10-01T17:09:44.938730Z",
     "shell.execute_reply": "2026-10-01T17:09:44.938185Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "autograd matches hand-derived gradient: True\n",
      "max abs difference: 2.7008354663848877e-08\n"
     ]
    }
   ],
   "source": [
    "dW = xenc.T @ ((probs.detach() - F.one_hot(ys, vocab_size).float()) / num)\n",
    "print(\"autograd matches hand-derived gradient:\", torch.allclose(W.grad, dW, atol=1e-6))\n",
    "print(\"max abs difference:\", (W.grad - dW).abs().max().item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b87454",
   "metadata": {},
   "source": [
    "## 6. Gradient descent step (SGD)\n",
    "We update `W` inside `torch.no_grad()`, because the update itself is not part of the model and should not be recorded for autograd.\n",
    "We then **reset `W.grad`**: PyTorch *adds* new gradients to whatever is already in `.grad`, so forgetting this would mix gradients from different steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "42dde5db",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.940673Z",
     "iopub.status.busy": "2026-10-01T17:09:44.940528Z",
     "iopub.status.idle": "2026-10-01T17:09:44.988725Z",
     "shell.execute_reply": "2026-10-01T17:09:44.988154Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss after one step: 3.4023\n"
     ]
    }
   ],
   "source": [
    "lr = 50.0\n",
    "with torch.no_grad():\n",
    "    W -= lr * W.grad\n",
    "W.grad = None                                         # clear the accumulated gradient\n",
    "\n",
    "logits = xenc @ W\n",
    "counts = (logits - logits.max(dim=1, keepdim=True).values).exp()\n",
    "probs = counts / counts.sum(dim=1, keepdim=True)\n",
    "loss = -probs[torch.arange(num), ys].log().mean()\n",
    "print(f\"loss after one step: {loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "224c6fb6",
   "metadata": {},
   "source": [
    "## 7. Generate names with the barely-trained model\n",
    "Sampling works as before; `torch.multinomial` draws the next index from a probability vector, using our seeded generator `g`. Expect gibberish after just one step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cccefaa0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.990751Z",
     "iopub.status.busy": "2026-10-01T17:09:44.990605Z",
     "iopub.status.idle": "2026-10-01T17:09:44.997029Z",
     "shell.execute_reply": "2026-10-01T17:09:44.996573Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opn\n",
      "klxeawv\n",
      "dm\n",
      "peudalloea\n",
      "lnahmpexus\n"
     ]
    }
   ],
   "source": [
    "def generate(W, n=5):\n",
    "    for _ in range(n):\n",
    "        out, ix = [], 0\n",
    "        while True:\n",
    "            p = F.softmax(W[ix].detach().cpu(), dim=0)  # same as onehot(ix) @ W, then softmax\n",
    "            ix = torch.multinomial(p, num_samples=1, generator=g).item()\n",
    "            if ix == 0:\n",
    "                break\n",
    "            out.append(itos[ix])\n",
    "        print(\"\".join(out))\n",
    "\n",
    "generate(W)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e75c61",
   "metadata": {},
   "source": [
    "## 8. The idiomatic PyTorch version\n",
    "In practice we let PyTorch supply the building blocks too: `nn.Linear(27, 27, bias=False)` holds the weight matrix, `F.cross_entropy` computes softmax + cross-entropy from the logits in one (more stable) step, and `torch.optim.SGD` performs the update.\n",
    "We also move the model and the data to `device`; on a GPU this is the only change needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bdb4bec2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:44.998768Z",
     "iopub.status.busy": "2026-10-01T17:09:44.998601Z",
     "iopub.status.idle": "2026-10-01T17:09:46.307063Z",
     "shell.execute_reply": "2026-10-01T17:09:46.306471Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linear(in_features=27, out_features=27, bias=False)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "model = nn.Linear(vocab_size, vocab_size, bias=False).to(device)\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
    "xenc, ys = xenc.to(device), ys.to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98744f09",
   "metadata": {},
   "source": [
    "Each epoch is now forward → `zero_grad` / `backward` / `step`. These three lines are the core training step for essentially every PyTorch model, from this bigram to a large language model.\n",
    "`nn.Linear` starts from smaller random weights than our `torch.randn`, so the loss begins near the uniform 3.30, and it should reach about 2.47, just like Topic 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "dcd18fc6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:46.309202Z",
     "iopub.status.busy": "2026-10-01T17:09:46.309005Z",
     "iopub.status.idle": "2026-10-01T17:09:51.783660Z",
     "shell.execute_reply": "2026-10-01T17:09:51.782851Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch   0  loss 3.2992\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  10  loss 2.6052\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  20  loss 2.5362\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  30  loss 2.5100\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  40  loss 2.4962\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  50  loss 2.4877\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  60  loss 2.4819\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  70  loss 2.4777\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  80  loss 2.4746\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  90  loss 2.4722\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  99  loss 2.4704\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(100):\n",
    "    logits = model(xenc)                              # same as xenc @ W\n",
    "    loss = F.cross_entropy(logits, ys)\n",
    "\n",
    "    optimizer.zero_grad()                             # clear old gradients\n",
    "    loss.backward()                                   # autograd computes gradients\n",
    "    optimizer.step()                                  # W <- W - lr * W.grad\n",
    "\n",
    "    if epoch % 10 == 0 or epoch == 99:\n",
    "        print(f\"epoch {epoch:3d}  loss {loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "003c5c55",
   "metadata": {},
   "source": [
    "`nn.Linear` stores its weight as (out_features, in_features), i.e. the transpose of our `W`, so we pass `model.weight.T` to the same `generate` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3d134a60",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-01T17:09:51.785905Z",
     "iopub.status.busy": "2026-10-01T17:09:51.785699Z",
     "iopub.status.idle": "2026-10-01T17:09:51.793640Z",
     "shell.execute_reply": "2026-10-01T17:09:51.793069Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mzdarrinam\n",
      "rn\n",
      "vybeartoqxy\n",
      "kugfwanee\n",
      "vayioklian\n"
     ]
    }
   ],
   "source": [
    "generate(model.weight.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba2e952f",
   "metadata": {},
   "source": [
    "## Takeaways\n",
    "- **Autograd**: `loss.backward()` computed exactly the gradient we derived by hand in Topic 1, with no calculus on our side.\n",
    "- **Building blocks**: `nn.Linear`, `F.cross_entropy` and `torch.optim.SGD` replace our hand-written layer, loss and update.\n",
    "- **GPU for free**: `.to(device)` moves the model and data; the rest of the code is unchanged.\n",
    "- The `zero_grad` / `backward` / `step` loop scales to any model. In Topic 3 we use it to train a deeper MLP that looks at more than one previous character."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}