{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's build GPT \u2014 from scratch\n",
    "\n",
    "We are rebuilding a tiny GPT on the *tiny Shakespeare* dataset, following Karpathy's\n",
    "\"Let's build GPT\" but with extra explanation at each step.\n",
    "\n",
    "**Before running:** in VS Code, click the kernel picker (top-right) and choose the\n",
    "interpreter at `gpt_from_scratch/.venv`. That venv has CUDA PyTorch installed.\n",
    "\n",
    "The roadmap inside this notebook:\n",
    "1. Data + tokenization\n",
    "2. Batching (the input/target shapes)\n",
    "3. A Bigram baseline (simplest possible language model) + training loop\n",
    "4. *(next session)* Self-attention \u2014 the heart of the Transformer\n",
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Setup check\n",
    "\n",
    "First confirm PyTorch can actually see your GPU. If `CUDA available` is False,\n",
    "stop \u2014 we're accidentally on the CPU build again."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "print(\"torch:\", torch.__version__)\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"device:\", device)\n",
    "if device == \"cuda\":\n",
    "    print(\"GPU:\", torch.cuda.get_device_name(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. The data\n",
    "\n",
    "A language model learns to predict the *next* piece of text given what came before.\n",
    "So all we need is a big pile of text. Ours is ~1MB of Shakespeare."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()\n",
    "\n",
    "print(\"length in characters:\", len(text))\n",
    "print(\"----\")\n",
    "print(text[:250])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Tokenization\n",
    "\n",
    "A neural net only eats **numbers**, not letters. *Tokenization* = deciding how to\n",
    "chop text into pieces (\"tokens\") and mapping each piece to an integer.\n",
    "\n",
    "Real GPTs use *sub-word* tokens (e.g. \"ing\", \"tion\") via an algorithm called BPE.\n",
    "We'll start with the simplest possible scheme: **one token = one character**.\n",
    "\n",
    "- Pro: trivially simple, tiny vocabulary.\n",
    "- Con: sequences get long (every letter is a step), so the model has to work harder.\n",
    "\n",
    "> Analogy: tokenization is choosing your alphabet. Characters = small alphabet, long\n",
    "> words. Sub-words = bigger alphabet, shorter words. GPT picks the bigger alphabet."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# all the unique characters that occur in the text = our vocabulary\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(\"\".join(chars))\n",
    "print(\"vocab_size:\", vocab_size)\n",
    "\n",
    "# build the two lookup tables: char->int (stoi) and int->char (itos)\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]          # string -> list of ints\n",
    "decode = lambda l: \"\".join(itos[i] for i in l)   # list of ints -> string\n",
    "\n",
    "print(encode(\"hello there\"))\n",
    "print(decode(encode(\"hello there\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Encode the whole dataset into a tensor\n",
    "\n",
    "We turn the entire text into one long 1-D tensor of integers, then hold out the last\n",
    "10% as a validation set (to check the model isn't just memorising)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:50])   # the first 50 chars, now as numbers\n",
    "\n",
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "print(\"train:\", len(train_data), \" val:\", len(val_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Context: what does \"predict the next token\" actually mean?\n",
    "\n",
    "We never feed the whole 1M-character book at once. We feed **chunks** of a fixed\n",
    "length called `block_size` (the context window). Within one chunk of length 8 there\n",
    "are actually 8 training examples packed in \u2014 predict char 2 from char 1, char 3 from\n",
    "chars 1-2, and so on. This teaches the model to handle contexts of every length from\n",
    "1 up to `block_size`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "block_size = 8\n",
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size + 1]   # y is x shifted left by one\n",
    "for t in range(block_size):\n",
    "    context = x[:t + 1]\n",
    "    target = y[t]\n",
    "    print(f\"when input is {context.tolist()} -> target is {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Batching\n",
    "\n",
    "GPUs are happiest doing many things in parallel. So instead of one chunk, we grab a\n",
    "**batch** of several random chunks and process them at once. Two key numbers:\n",
    "\n",
    "- `batch_size` = how many independent chunks we stack (parallelism).\n",
    "- `block_size` = how long each chunk is (context length).\n",
    "\n",
    "`get_batch` picks random starting points and returns:\n",
    "- `x`: the inputs, shape `(batch_size, block_size)`\n",
    "- `y`: the targets (each input shifted by one), same shape.\n",
    "\n",
    "`.to(device)` moves the tensors onto the GPU."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4\n",
    "block_size = 8\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,))   # random start indices\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",
    "xb, yb = get_batch(\"train\")\n",
    "print(\"inputs \", xb.shape)\n",
    "print(xb)\n",
    "print(\"targets\", yb.shape)\n",
    "print(yb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. The Bigram model \u2014 our baseline\n",
    "\n",
    "Simplest possible language model: to predict the next character, look **only** at the\n",
    "current character. We do this with an `nn.Embedding` table of shape\n",
    "`(vocab_size, vocab_size)`. Row `i` is directly the scores (\"logits\") for what comes\n",
    "after token `i`.\n",
    "\n",
    "Key concepts in this cell:\n",
    "- **Logits**: raw, unnormalised scores for each possible next token.\n",
    "- **Cross-entropy loss**: measures how surprised the model is by the true next token.\n",
    "  Lower = better. With `vocab_size` random guesses we expect loss ~ `-ln(1/vocab_size)`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class BigramLanguageModel(nn.Module):\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads the next-token logits from this lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        logits = self.token_embedding_table(idx)   # (B, T, vocab_size)\n",
    "        if targets is None:\n",
    "            return logits, None\n",
    "        # cross_entropy wants (N, C), so flatten the batch & time dims together\n",
    "        B, T, C = logits.shape\n",
    "        loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B, T) of current context; extend it one token at a time\n",
    "        for _ in range(max_new_tokens):\n",
    "            logits, _ = self(idx)\n",
    "            logits = logits[:, -1, :]                 # only the last step -> (B, C)\n",
    "            probs = F.softmax(logits, dim=-1)         # scores -> probabilities\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)  # sample one token\n",
    "            idx = torch.cat((idx, idx_next), dim=1)\n",
    "        return idx\n",
    "\n",
    "model = BigramLanguageModel(vocab_size).to(device)\n",
    "logits, loss = model(xb, yb)\n",
    "expected = -torch.log(torch.tensor(1.0 / vocab_size)).item()\n",
    "print(\"logits shape:\", logits.shape)\n",
    "print(f\"initial loss: {loss.item():.4f}   (random guessing ~ {expected:.4f})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Generate *before* training\n",
    "\n",
    "Right now the weights are random, so this will be pure gibberish. We start generation\n",
    "from a single newline character (token 0). This is our \"before\" picture."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Train it\n",
    "\n",
    "The training loop is the same four steps forever:\n",
    "1. **forward** \u2014 get predictions + loss\n",
    "2. **zero_grad** \u2014 clear old gradients\n",
    "3. **backward** \u2014 backpropagation computes how each weight affected the loss\n",
    "4. **step** \u2014 the optimizer (AdamW) nudges weights to reduce the loss\n",
    "\n",
    "> Analogy: loss is \"how wrong we are.\" Backprop tells each knob which way to turn.\n",
    "> The optimizer turns all the knobs a little. Repeat thousands of times."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
    "batch_size = 32\n",
    "\n",
    "for step in range(10000):\n",
    "    xb, yb = get_batch(\"train\")\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if step % 1000 == 0:\n",
    "        print(f\"step {step:5d}: loss {loss.item():.4f}\")\n",
    "\n",
    "print(\"final loss:\", round(loss.item(), 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Generate *after* training\n",
    "\n",
    "Still not Shakespeare \u2014 a bigram model only ever looks at one character back, so it\n",
    "can't form real words. But it should now produce plausible letter/space patterns\n",
    "instead of random noise. That jump is the model *learning*.\n",
    "\n",
    "The fix for \"only one character back\" is **self-attention** \u2014 letting each token look\n",
    "at all previous tokens. That's the heart of the Transformer, and it's our next step."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(model.generate(context, max_new_tokens=400)[0].tolist()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}