{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "3d5936ca",
      "metadata": {},
      "source": [
        "# Let's build GPT — self-attention to a full transformer\n",
        "\n",
        "Give every token the power to **look back** — **concept -> code -> Your turn** each step.\n",
        "\n",
        "This is notebook **06**, the payoff. Notebook `01` built a bigram that only sees one character\n",
        "back. Here we add **self-attention** so each position can gather information from all earlier\n",
        "positions, then stack it into a real (small) **GPT** and train it on tiny Shakespeare.\n",
        "Follows Karpathy's *\"Let's build GPT from scratch\"*."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8b6c4fb3",
      "metadata": {},
      "source": [
        "## Prologue — in plain English\n",
        "\n",
        "The bigram's weakness: predicting the next letter from only the previous one. WaveNet (`05`)\n",
        "widened the view but fused neighbours in a fixed pattern. **Attention** is smarter: every\n",
        "token looks at all earlier tokens and decides, on its own, *which* ones matter right now.\n",
        "\n",
        "Real-life picture: you are reading a sentence and hit the word \"it.\" To know what \"it\" refers\n",
        "to, your eyes jump back to the relevant earlier noun — not to a fixed position, but to whichever\n",
        "word is relevant. Self-attention lets each token make that kind of targeted look-back.\n",
        "\n",
        "We build it up in small pieces: the averaging trick, then queries/keys/values, then multiple\n",
        "heads, then the full transformer block (attention + feed-forward + residual + LayerNorm)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "028dce18",
      "metadata": {},
      "source": [
        "### 0.1 Setup and data (standalone)\n",
        "\n",
        "Self-contained: load `input.txt`, build the character vocabulary, and make `get_batch`. Uses\n",
        "GPU if available, else CPU. The model is intentionally **small** so it trains on a laptop CPU;\n",
        "the comments show how to scale it up if you have a GPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "36f4f6e6",
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(\"device:\", device, \"| torch\", torch.__version__)\n",
        "\n",
        "# --- hyperparameters (small for CPU; bump these up on a GPU) ---\n",
        "batch_size = 16\n",
        "block_size = 32       # max context length\n",
        "n_embd = 64\n",
        "n_head = 4\n",
        "n_layer = 3\n",
        "dropout = 0.0\n",
        "learning_rate = 1e-3\n",
        "max_iters = 3000\n",
        "eval_interval = 1000\n",
        "eval_iters = 50\n",
        "\n",
        "# --- data ---\n",
        "text = open(\"input.txt\", \"r\", encoding=\"utf-8\").read()\n",
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "stoi = {ch: i for i, ch in enumerate(chars)}\n",
        "itos = {i: ch for i, ch in enumerate(chars)}\n",
        "encode = lambda s: [stoi[c] for c in s]\n",
        "decode = lambda l: \"\".join(itos[i] for i in l)\n",
        "\n",
        "data = torch.tensor(encode(text), dtype=torch.long)\n",
        "n = int(0.9 * len(data))\n",
        "train_data, val_data = data[:n], data[n:]\n",
        "\n",
        "def get_batch(split):\n",
        "    d = train_data if split == \"train\" else val_data\n",
        "    ix = torch.randint(len(d) - block_size, (batch_size,))\n",
        "    x = torch.stack([d[i:i + block_size] for i in ix])\n",
        "    y = torch.stack([d[i + 1:i + block_size + 1] for i in ix])\n",
        "    return x.to(device), y.to(device)\n",
        "\n",
        "print(\"vocab_size:\", vocab_size, \"| chars in dataset:\", len(text))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d2b12520",
      "metadata": {},
      "source": [
        "### 1.1 The math trick: a token can average its past\n",
        "\n",
        "Before real attention, here is the key move in its simplest form: let each position become the\n",
        "**average of itself and all positions before it** (never the future — that would be cheating).\n",
        "We show three ways that give the *same* result, from slow to elegant."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "54c69615",
      "metadata": {},
      "outputs": [],
      "source": [
        "torch.manual_seed(1337)\n",
        "B, T, C = 4, 8, 2\n",
        "x = torch.randn(B, T, C)\n",
        "\n",
        "# v1: explicit loop — average of all previous tokens\n",
        "xbow = torch.zeros((B, T, C))\n",
        "for b in range(B):\n",
        "    for t in range(T):\n",
        "        xbow[b, t] = x[b, :t + 1].mean(0)\n",
        "\n",
        "# v2: same thing as a matrix multiply with a lower-triangular average matrix\n",
        "wei = torch.tril(torch.ones(T, T))\n",
        "wei = wei / wei.sum(1, keepdim=True)\n",
        "xbow2 = wei @ x\n",
        "\n",
        "# v3: same again, but the weights come from softmax over a masked matrix\n",
        "tril = torch.tril(torch.ones(T, T))\n",
        "wei = torch.zeros((T, T)).masked_fill(tril == 0, float(\"-inf\"))\n",
        "wei = F.softmax(wei, dim=-1)\n",
        "xbow3 = wei @ x\n",
        "\n",
        "print(\"v1 == v2 :\", torch.allclose(xbow, xbow2, atol=1e-6), \" (max diff %.1e - just float rounding)\" % (xbow - xbow2).abs().max())\n",
        "print(\"v1 == v3 :\", torch.allclose(xbow, xbow3, atol=1e-6), \" (max diff %.1e - just float rounding)\" % (xbow - xbow3).abs().max())\n",
        "print(\"\\nthe softmax weight matrix (row t can only attend to columns <= t):\")\n",
        "print(wei)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aa46672a",
      "metadata": {},
      "source": [
        "Why the `-inf` + softmax version matters: those weights do **not** have to be a flat\n",
        "average. If the numbers in `wei` are *learned*, each token can decide how much to listen to each\n",
        "earlier token. That is exactly what attention does next.\n",
        "\n",
        "Analogy: a meeting where you may only hear people who already spoke (the mask), and you choose\n",
        "how much weight to give each speaker (the learned weights)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b522ef98",
      "metadata": {},
      "source": [
        "### 2.1 A single self-attention head — query, key, value\n",
        "\n",
        "Each token emits three vectors:\n",
        "\n",
        "- **query**: \"what am I looking for?\"\n",
        "- **key**: \"what do I offer?\"\n",
        "- **value**: \"what will I actually pass on if chosen?\"\n",
        "\n",
        "A token's attention to another = how well its **query** matches the other's **key**\n",
        "(dot product). We scale by `1/sqrt(head_size)` to keep the numbers tame, mask the future, softmax\n",
        "into weights, then take the weighted sum of **values**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9adfb9ea",
      "metadata": {},
      "outputs": [],
      "source": [
        "torch.manual_seed(1337)\n",
        "B, T, C = 4, 8, 32\n",
        "x = torch.randn(B, T, C)\n",
        "head_size = 16\n",
        "\n",
        "key   = nn.Linear(C, head_size, bias=False)\n",
        "query = nn.Linear(C, head_size, bias=False)\n",
        "value = nn.Linear(C, head_size, bias=False)\n",
        "\n",
        "k = key(x)        # (B, T, head_size)\n",
        "q = query(x)      # (B, T, head_size)\n",
        "wei = q @ k.transpose(-2, -1) * head_size ** -0.5   # (B, T, T) affinities\n",
        "\n",
        "tril = torch.tril(torch.ones(T, T))\n",
        "wei = wei.masked_fill(tril == 0, float(\"-inf\"))\n",
        "wei = F.softmax(wei, dim=-1)\n",
        "\n",
        "v = value(x)\n",
        "out = wei @ v     # (B, T, head_size)\n",
        "print(\"attention output shape:\", tuple(out.shape))\n",
        "print(\"\\nrow 0 of the learned attention weights (note they are NOT uniform):\")\n",
        "print(wei[0, -1])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ff4eea4c",
      "metadata": {},
      "source": [
        "### 3.1 The full GPT, assembled from clean modules\n",
        "\n",
        "Now we package it. The pieces:\n",
        "\n",
        "- **Head** — one self-attention head (above), with dropout and a causal mask buffer.\n",
        "- **MultiHeadAttention** — several heads in parallel, concatenated, then projected. Analogy:\n",
        "  several readers each look back for a different kind of clue, then pool notes.\n",
        "- **FeedForward** — a small per-token MLP: \"time to think about what I gathered.\"\n",
        "- **Block** — attention + feed-forward, each wrapped with a **residual** connection and\n",
        "  **LayerNorm**. Residual = keep the original and add the new insight (so deep stacks still\n",
        "  train). LayerNorm = re-level each token's vector (the `03` lesson, per-token).\n",
        "- **GPTLanguageModel** — token + position embeddings -> a stack of Blocks -> final LayerNorm ->\n",
        "  a linear head to vocab logits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b0a4b08d",
      "metadata": {},
      "outputs": [],
      "source": [
        "class Head(nn.Module):\n",
        "    def __init__(self, head_size):\n",
        "        super().__init__()\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.register_buffer(\"tril\", torch.tril(torch.ones(block_size, block_size)))\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.shape\n",
        "        k = self.key(x)\n",
        "        q = self.query(x)\n",
        "        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5\n",
        "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float(\"-inf\"))\n",
        "        wei = F.softmax(wei, dim=-1)\n",
        "        wei = self.dropout(wei)\n",
        "        v = self.value(x)\n",
        "        return wei @ v\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, num_heads, head_size):\n",
        "        super().__init__()\n",
        "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
        "        self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
        "        return self.dropout(self.proj(out))\n",
        "\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "    def __init__(self, n_embd):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(dropout),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.sa = MultiHeadAttention(n_head, head_size)\n",
        "        self.ffwd = FeedForward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.sa(self.ln1(x))     # residual around attention\n",
        "        x = x + self.ffwd(self.ln2(x))   # residual around feed-forward\n",
        "        return x\n",
        "\n",
        "\n",
        "class GPTLanguageModel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
        "        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(n_embd)\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "        tok_emb = self.token_embedding_table(idx)                       # (B,T,C)\n",
        "        pos_emb = self.position_embedding_table(torch.arange(T, device=device))\n",
        "        x = tok_emb + pos_emb\n",
        "        x = self.blocks(x)\n",
        "        x = self.ln_f(x)\n",
        "        logits = self.lm_head(x)                                        # (B,T,vocab)\n",
        "        if targets is None:\n",
        "            return logits, None\n",
        "        B, T, Cc = logits.shape\n",
        "        loss = F.cross_entropy(logits.view(B * T, Cc), targets.view(B * T))\n",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        for _ in range(max_new_tokens):\n",
        "            idx_cond = idx[:, -block_size:]          # never exceed the context window\n",
        "            logits, _ = self(idx_cond)\n",
        "            logits = logits[:, -1, :]                # last position only\n",
        "            probs = F.softmax(logits, dim=-1)\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)\n",
        "            idx = torch.cat((idx, idx_next), dim=1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "model = GPTLanguageModel().to(device)\n",
        "print(\"parameters (millions):\", sum(p.numel() for p in model.parameters()) / 1e6)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dbb781c5",
      "metadata": {},
      "source": [
        "### 3.2 Positional embeddings — why we need them\n",
        "\n",
        "Attention by itself is **order-blind**: it treats the context as a bag of tokens. But \"dog\n",
        "bites man\" and \"man bites dog\" must differ. So besides each token's meaning\n",
        "(`token_embedding_table`), we add a vector for its **position**\n",
        "(`position_embedding_table`). Now a token knows both *what* it is and *where* it is.\n",
        "\n",
        "(That addition already happens inside `forward`: `x = tok_emb + pos_emb`.)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0385a0ef",
      "metadata": {},
      "source": [
        "### 4.1 Train the GPT\n",
        "\n",
        "Same training shape as always, with the AdamW optimizer. Watch train and val loss fall well\n",
        "below the bigram's ~2.5. (Small model + CPU, so this is a few minutes; scale up on a GPU for\n",
        "real Shakespeare.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "395e6922",
      "metadata": {},
      "outputs": [],
      "source": [
        "@torch.no_grad()\n",
        "def estimate_loss():\n",
        "    out = {}\n",
        "    model.eval()\n",
        "    for split in [\"train\", \"val\"]:\n",
        "        losses = torch.zeros(eval_iters)\n",
        "        for k in range(eval_iters):\n",
        "            X, Y = get_batch(split)\n",
        "            _, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean().item()\n",
        "    model.train()\n",
        "    return out\n",
        "\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for it in range(max_iters):\n",
        "    if it % eval_interval == 0 or it == max_iters - 1:\n",
        "        losses = estimate_loss()\n",
        "        print(f\"step {it:4d}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "    xb, yb = get_batch(\"train\")\n",
        "    _, loss = model(xb, yb)\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "print(\"done training\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0ae5a608",
      "metadata": {},
      "source": [
        "### 5.1 Generate Shakespeare-ish text\n",
        "\n",
        "Start from a single newline and let the model write. It will not be perfect Shakespeare (tiny\n",
        "model, short training), but compared to the bigram it now produces real-looking words, spacing,\n",
        "and line breaks — because each character can attend to the ones before it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "61db808c",
      "metadata": {},
      "outputs": [],
      "source": [
        "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
        "print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6b43261f",
      "metadata": {},
      "source": [
        "**Your turn 6** — If you have a GPU, scale up (e.g. `n_embd=384`, `n_head=6`,\n",
        "`n_layer=6`, `block_size=256`, `dropout=0.2`, `max_iters=5000`) and retrain — the samples get\n",
        "dramatically better. On CPU, try just raising `n_layer` to 4 or `max_iters` to 5000 and compare\n",
        "the val loss."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1e1142ab",
      "metadata": {},
      "source": [
        "## What's next — and the connection to ChatGPT\n",
        "\n",
        "You built a real transformer: token + position embeddings, multi-head causal self-attention,\n",
        "feed-forward layers, residuals, and LayerNorm — the exact architecture behind GPT.\n",
        "\n",
        "What you trained here is a **pretrained** model in miniature: it only learns to *continue text*\n",
        "(predict the next token) from a big pile of text. ChatGPT starts the same way, then adds steps\n",
        "this series does not cover:\n",
        "\n",
        "- **Supervised finetuning** on example conversations (teach it to answer, not just continue).\n",
        "- **Reinforcement learning from human feedback (RLHF)** to align responses with what people\n",
        "  prefer.\n",
        "\n",
        "The remaining notebook:\n",
        "\n",
        "- `07_gpt_tokenizer.ipynb` — real models do not use one-character-per-token; they use **Byte\n",
        "  Pair Encoding**. That is the final piece of the pipeline.\n",
        "\n",
        "**Checklist**\n",
        "- [ ] The averaging trick (loop = tril matmul = masked softmax)\n",
        "- [ ] Query / key / value self-attention head\n",
        "- [ ] Multi-head attention + projection\n",
        "- [ ] Feed-forward, residuals, LayerNorm -> Block\n",
        "- [ ] Token + position embeddings -> full GPT\n",
        "- [ ] Trained below the bigram baseline; generated text"
      ]
    }
  ],
  "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
}