derive-gpt-from-scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install \"torch>=2.1\" numpy --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Shakespeare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('shakespeare.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters: 1115395\n"
     ]
    }
   ],
   "source": [
    "print(f\"length of dataset in characters: {len(text)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You are all resolved rather to die than to famish?\n",
      "\n",
      "All:\n",
      "Resolved. resolved.\n",
      "\n",
      "First Citizen:\n",
      "First, you know Caius Marcius is chief enemy to the people.\n",
      "\n",
      "All:\n",
      "We know't, we know't.\n",
      "\n",
      "First Citizen:\n",
      "Let us kill him, and we'll have corn at our own price.\n",
      "Is't a verdict?\n",
      "\n",
      "All:\n",
      "No more talking on't; let it be done: away, away!\n",
      "\n",
      "Second Citizen:\n",
      "One word, good citizens.\n",
      "\n",
      "First Citizen:\n",
      "We are accounted poor citizens, the patricians good.\n",
      "What authority surfeits on would relieve us: if they\n",
      "would yield us but the superfluity, while it were\n",
      "wholesome, we might guess they relieved us humanely;\n",
      "but they think we are too dear: the leanness that\n",
      "afflicts us, the object of our misery, is as an\n",
      "inventory to particularise their abundance; our\n",
      "sufferance is a gain to them Let us revenge this with\n",
      "our pikes, ere we become rakes: for the gods know I\n",
      "speak this in hunger for bread, not in thirst for revenge.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "65\n"
     ]
    }
   ],
   "source": [
    "# all the unique chars that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[58, 56, 39, 52, 57, 44, 53, 56, 51, 43, 56, 57]\n",
      "transformers\n"
     ]
    }
   ],
   "source": [
    "# Implement a simple chaacter level tokenization schema. More sophisticated tokenizers include SentencePiece / tiktoken\n",
    "stoi = { ch:i for i,ch in enumerate(chars) } # str to int mapping\n",
    "itos = { i:ch for i,ch in enumerate(chars) } # int to str mapping\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])\n",
    "\n",
    "print(encode(\"transformers\"))\n",
    "print(decode(encode(\"transformers\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115395]) torch.int64\n",
      "tensor([ 0, 18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43,\n",
      "        44, 53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52,\n",
      "        63,  1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,\n",
      "         1, 57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39,\n",
      "        49,  6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15,\n",
      "        47, 58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50,\n",
      "        50,  1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1,\n",
      "        58, 53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51,\n",
      "        47, 57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43,\n",
      "        42,  8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57,\n",
      "        58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1,\n",
      "        63, 53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56,\n",
      "        41, 47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51,\n",
      "        63,  1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0,\n",
      "        13, 50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,\n",
      "         1, 49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
      "        58, 47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57,  1, 49, 47, 50, 50,\n",
      "         1, 46, 47, 51,  6,  1, 39, 52, 42,  1, 61, 43,  5, 50, 50,  1, 46, 39,\n",
      "        60, 43,  1, 41, 53, 56, 52,  1, 39, 58,  1, 53, 59, 56,  1, 53, 61, 52,\n",
      "         1, 54, 56, 47, 41, 43,  8,  0, 21, 57,  5, 58,  1, 39,  1, 60, 43, 56,\n",
      "        42, 47, 41, 58, 12,  0,  0, 13, 50, 50, 10,  0, 26, 53,  1, 51, 53, 56,\n",
      "        43,  1, 58, 39, 50, 49, 47, 52, 45,  1, 53, 52,  5, 58, 11,  1, 50, 43,\n",
      "        58,  1, 47, 58,  1, 40, 43,  1, 42, 53, 52, 43, 10,  1, 39, 61, 39, 63,\n",
      "         6,  1, 39, 61, 39, 63,  2,  0,  0, 31, 43, 41, 53, 52, 42,  1, 15, 47,\n",
      "        58, 47, 64, 43, 52, 10,  0, 27, 52, 43,  1, 61, 53, 56, 42,  6,  1, 45,\n",
      "        53, 53, 42,  1, 41, 47, 58, 47, 64, 43, 52, 57,  8,  0,  0, 18, 47, 56,\n",
      "        57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 35, 43,  1, 39, 56, 43,\n",
      "         1, 39, 41, 41, 53, 59, 52, 58, 43, 42,  1, 54, 53, 53, 56,  1, 41, 47,\n",
      "        58, 47, 64, 43, 52, 57,  6,  1, 58, 46, 43,  1, 54, 39, 58, 56, 47, 41,\n",
      "        47, 39, 52, 57,  1, 45, 53, 53, 42,  8,  0, 35, 46, 39, 58,  1, 39, 59,\n",
      "        58, 46, 53, 56, 47, 58, 63,  1, 57, 59, 56, 44, 43, 47, 58, 57,  1, 53,\n",
      "        52,  1, 61, 53, 59, 50, 42,  1, 56, 43, 50, 47, 43, 60, 43,  1, 59, 57,\n",
      "        10,  1, 47, 44,  1, 58, 46, 43, 63,  0, 61, 53, 59, 50, 42,  1, 63, 47,\n",
      "        43, 50, 42,  1, 59, 57,  1, 40, 59, 58,  1, 58, 46, 43,  1, 57, 59, 54,\n",
      "        43, 56, 44, 50, 59, 47, 58, 63,  6,  1, 61, 46, 47, 50, 43,  1, 47, 58,\n",
      "         1, 61, 43, 56, 43,  0, 61, 46, 53, 50, 43, 57, 53, 51, 43,  6,  1, 61,\n",
      "        43,  1, 51, 47, 45, 46, 58,  1, 45, 59, 43, 57, 57,  1, 58, 46, 43, 63,\n",
      "         1, 56, 43, 50, 47, 43, 60, 43, 42,  1, 59, 57,  1, 46, 59, 51, 39, 52,\n",
      "        43, 50, 63, 11,  0, 40, 59, 58,  1, 58, 46, 43, 63,  1, 58, 46, 47, 52,\n",
      "        49,  1, 61, 43,  1, 39, 56, 43,  1, 58, 53, 53,  1, 42, 43, 39, 56, 10,\n",
      "         1, 58, 46, 43,  1, 50, 43, 39, 52, 52, 43, 57, 57,  1, 58, 46, 39, 58,\n",
      "         0, 39, 44, 44, 50, 47, 41, 58, 57,  1, 59, 57,  6,  1, 58, 46, 43,  1,\n",
      "        53, 40, 48, 43, 41, 58,  1, 53, 44,  1, 53, 59, 56,  1, 51, 47, 57, 43,\n",
      "        56, 63,  6,  1, 47, 57,  1, 39, 57,  1, 39, 52,  0, 47, 52, 60, 43, 52,\n",
      "        58, 53, 56, 63,  1, 58, 53,  1, 54, 39, 56, 58, 47, 41, 59, 50, 39, 56,\n",
      "        47, 57, 43,  1, 58, 46, 43, 47, 56,  1, 39, 40, 59, 52, 42, 39, 52, 41,\n",
      "        43, 11,  1, 53, 59, 56,  0, 57, 59, 44, 44, 43, 56, 39, 52, 41, 43,  1,\n",
      "        47, 57,  1, 39,  1, 45, 39, 47, 52,  1, 58, 53,  1, 58, 46, 43, 51,  1,\n",
      "        24, 43, 58,  1, 59, 57,  1, 56, 43, 60, 43, 52, 45, 43,  1, 58, 46, 47,\n",
      "        57,  1, 61, 47, 58, 46,  0, 53, 59, 56,  1, 54, 47, 49, 43, 57,  6,  1,\n",
      "        43, 56, 43,  1, 61, 43,  1, 40, 43, 41, 53, 51, 43,  1, 56, 39, 49, 43,\n",
      "        57, 10,  1, 44, 53, 56,  1, 58, 46, 43,  1, 45, 53, 42, 57,  1, 49, 52,\n",
      "        53, 61,  1, 21,  0, 57, 54, 43, 39, 49,  1, 58, 46, 47, 57,  1, 47, 52,\n",
      "         1, 46, 59, 52, 45, 43, 56,  1, 44, 53, 56,  1, 40, 56, 43, 39, 42,  6,\n",
      "         1, 52, 53, 58,  1, 47, 52,  1, 58, 46, 47, 56, 57, 58,  1, 44, 53, 56,\n",
      "         1, 56, 43, 60, 43, 52, 45, 43,  8,  0])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train-Test Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1003855 111540\n"
     ]
    }
   ],
   "source": [
    "n = int(0.9*len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "print(len(train_data), len(val_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0, 18, 47, 56, 57, 58,  1, 15, 47])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "When input content is tensor([0]), the targer is 18\n",
      "When input content is tensor([ 0, 18]), the targer is 47\n",
      "When input content is tensor([ 0, 18, 47]), the targer is 56\n",
      "When input content is tensor([ 0, 18, 47, 56]), the targer is 57\n",
      "When input content is tensor([ 0, 18, 47, 56, 57]), the targer is 58\n",
      "When input content is tensor([ 0, 18, 47, 56, 57, 58]), the targer is 1\n",
      "When input content is tensor([ 0, 18, 47, 56, 57, 58,  1]), the targer is 15\n",
      "When input content is tensor([ 0, 18, 47, 56, 57, 58,  1, 15]), the targer is 47\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"When input content is {context}, the targer is {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      "torch.Size([4, 8])\n",
      "tensor([[59, 56,  6,  0, 24, 43, 58,  1],\n",
      "        [50, 39, 47, 51,  1, 58, 46, 39],\n",
      "        [47, 52, 45,  1, 58, 53,  1, 57],\n",
      "        [14, 43, 47, 52, 45,  1, 46, 53]])\n",
      "targets:\n",
      "torch.Size([4, 8])\n",
      "tensor([[56,  6,  0, 24, 43, 58,  1, 61],\n",
      "        [39, 47, 51,  1, 58, 46, 39, 58],\n",
      "        [52, 45,  1, 58, 53,  1, 57, 39],\n",
      "        [43, 47, 52, 45,  1, 46, 53, 50]])\n",
      "----\n",
      "when input is tensor([59]), target is 56\n",
      "when input is tensor([59, 56]), target is 6\n",
      "when input is tensor([59, 56,  6]), target is 0\n",
      "when input is tensor([59, 56,  6,  0]), target is 24\n",
      "when input is tensor([50]), target is 39\n",
      "when input is tensor([50, 39]), target is 47\n",
      "when input is tensor([50, 39, 47]), target is 51\n",
      "when input is tensor([50, 39, 47, 51]), target is 1\n",
      "when input is tensor([47]), target is 52\n",
      "when input is tensor([47, 52]), target is 45\n",
      "when input is tensor([47, 52, 45]), target is 1\n",
      "when input is tensor([47, 52, 45,  1]), target is 58\n",
      "when input is tensor([14]), target is 43\n",
      "when input is tensor([14, 43]), target is 47\n",
      "when input is tensor([14, 43, 47]), target is 52\n",
      "when input is tensor([14, 43, 47, 52]), target is 45\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4\n",
    "block_size = 8\n",
    "\n",
    "def get_batch(split, batch_size, block_size):\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train', batch_size, block_size)\n",
    "print('inputs:')\n",
    "print(xb.shape)\n",
    "print(xb)\n",
    "print('targets:')\n",
    "print(yb.shape)\n",
    "print(yb)\n",
    "\n",
    "print('----')\n",
    "\n",
    "for b in range(batch_size):\n",
    "    for t in range(batch_size):\n",
    "        context = xb[b, :t+1]\n",
    "        target =yb[b,t]\n",
    "        print(f\"when input is {context}, target is {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bigram LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 65])\n",
      "tensor(4.6299, grad_fn=<NllLossBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # the whole model is simply a square lookup table.\n",
    "        # For each char (token) a in the whole char set, we maintain the probability of char b appearing after a.\n",
    "        # So the size is vocab_size x vocab_size\n",
    "\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        # idx and targets are both (B, T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx)  # (B, T, C) -> Batch=4(batch size), Time=8(block_size), Channel=65(vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits & targets into 2D to cater for F.corss_entropy\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens): # idx is (B, T) array of indices in the current context\n",
    "        \n",
    "        # generate max_new_tokens tokens iteratively by looking at only the last token each time\n",
    "        for _ in range(max_new_tokens):\n",
    "            \n",
    "            # get the predictions\n",
    "            logits, loss = self(idx)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "    \n",
    "m = BigramLanguageModel(vocab_size)\n",
    "logits, loss = m(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "SKIcLT;AcELMoTbvZv C?nq-QE33:CJqkOKH-q;:la!oiywkHjgChzbQ?u!3bLIgwevmyFJGUGp\n",
      "wnYWmnxKWWev-tDqXErVKLgJ\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m.generate(idx, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Bigram LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, steps, batch_size, block_size, lr=1e-3):\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr)\n",
    "    for steps in range(steps):\n",
    "        xb, yb = get_batch('train', batch_size, block_size)\n",
    "\n",
    "        logits, loss = model(xb, yb)\n",
    "        optimizer.zero_grad(set_to_none=True) # clear grads from the previous step\n",
    "        loss.backward() # calculate grads for all params\n",
    "        optimizer.step() # update params\n",
    "\n",
    "    print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.572626829147339\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 32\n",
    "train(m, 5000, batch_size, block_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "gnd ous fo,\n",
      "\n",
      "Thathoree w flat ze, dr q? vmy'NCIN stocowinsth,\n",
      "\n",
      "TI! R:Cak.\n",
      "EOUe.\n",
      "T!\n",
      "PUCENI herd toug\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m.generate(idx, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deriving self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 2])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# consider this toy example batch:\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B, T, C = 4, 8, 2 # batch, time, channels\n",
    "x = torch.randn(B, T, C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, the bigram model is not communicating / paying attention at the (n-1, n-2, ... , 1)th tokens when predicting the (n+1)th token from the nth. Majority of the context info is lost.\n",
    "\n",
    "We need to derive a mechanism for the model to attend to previous tokens when predicting the future token.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive aggregation: averaging past tokens (weakest from of \"communication\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 1: by naive for loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.1808, -0.0700],\n",
       "         [-0.0894, -0.4926],\n",
       "         [ 0.1490, -0.3199],\n",
       "         [ 0.3504, -0.2238],\n",
       "         [ 0.3525,  0.0545],\n",
       "         [ 0.0688, -0.0396],\n",
       "         [ 0.0927, -0.0682],\n",
       "         [-0.0341,  0.1332]],\n",
       "\n",
       "        [[ 1.3488, -0.1396],\n",
       "         [ 0.8173,  0.4127],\n",
       "         [-0.1342,  0.4395],\n",
       "         [ 0.2711,  0.4774],\n",
       "         [ 0.2421,  0.0694],\n",
       "         [ 0.0084,  0.0020],\n",
       "         [ 0.0712, -0.1128],\n",
       "         [ 0.2527,  0.2149]],\n",
       "\n",
       "        [[-0.6631, -0.2513],\n",
       "         [ 0.1735, -0.0649],\n",
       "         [ 0.1685,  0.3348],\n",
       "         [-0.1621,  0.1765],\n",
       "         [-0.2312, -0.0436],\n",
       "         [-0.1015, -0.2855],\n",
       "         [-0.2593, -0.1630],\n",
       "         [-0.3015, -0.2293]],\n",
       "\n",
       "        [[ 1.6455, -0.8030],\n",
       "         [ 1.4985, -0.5395],\n",
       "         [ 0.4954,  0.3420],\n",
       "         [ 1.0623, -0.1802],\n",
       "         [ 1.1401, -0.4462],\n",
       "         [ 1.0870, -0.4071],\n",
       "         [ 1.0430, -0.1299],\n",
       "         [ 1.1138, -0.1641]]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
    "xbow = torch.zeros((B, T, C)) # x bag of words\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1]\n",
    "        xbow[b, t] = torch.mean(xprev, 0)\n",
    "xbow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=tensor([[1.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333]])\n",
      "b=tensor([[2., 7.],\n",
      "        [6., 4.],\n",
      "        [6., 5.]])\n",
      "c=tensor([[2.0000, 7.0000],\n",
      "        [4.0000, 5.5000],\n",
      "        [4.6667, 5.3333]])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones(3, 3))\n",
    "a = a/ torch.sum(a, 1, keepdim=True)\n",
    "b = torch.randint(0,10,(3,2)).float()\n",
    "c = a @ b\n",
    "print(f'a={a}')\n",
    "print(f'b={b}')\n",
    "print(f'c={c}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 2: by matrix multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The operation of x[b,t] = mean_{i<=t} x[b,i] can be simplified & optimized by a row-normalized (each row sums to 1) lower triangular matric @ x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version 2.1: get weights by dividing by row sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wei: tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],\n",
      "        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],\n",
      "        [0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])\n",
      "xbow2: tensor([[[ 0.1808, -0.0700],\n",
      "         [-0.0894, -0.4926],\n",
      "         [ 0.1490, -0.3199],\n",
      "         [ 0.3504, -0.2238],\n",
      "         [ 0.3525,  0.0545],\n",
      "         [ 0.0688, -0.0396],\n",
      "         [ 0.0927, -0.0682],\n",
      "         [-0.0341,  0.1332]],\n",
      "\n",
      "        [[ 1.3488, -0.1396],\n",
      "         [ 0.8173,  0.4127],\n",
      "         [-0.1342,  0.4395],\n",
      "         [ 0.2711,  0.4774],\n",
      "         [ 0.2421,  0.0694],\n",
      "         [ 0.0084,  0.0020],\n",
      "         [ 0.0712, -0.1128],\n",
      "         [ 0.2527,  0.2149]],\n",
      "\n",
      "        [[-0.6631, -0.2513],\n",
      "         [ 0.1735, -0.0649],\n",
      "         [ 0.1685,  0.3348],\n",
      "         [-0.1621,  0.1765],\n",
      "         [-0.2312, -0.0436],\n",
      "         [-0.1015, -0.2855],\n",
      "         [-0.2593, -0.1630],\n",
      "         [-0.3015, -0.2293]],\n",
      "\n",
      "        [[ 1.6455, -0.8030],\n",
      "         [ 1.4985, -0.5395],\n",
      "         [ 0.4954,  0.3420],\n",
      "         [ 1.0623, -0.1802],\n",
      "         [ 1.1401, -0.4462],\n",
      "         [ 1.0870, -0.4071],\n",
      "         [ 1.0430, -0.1299],\n",
      "         [ 1.1138, -0.1641]]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "wei = torch.tril(torch.ones(T, T)) # weight - the row-normalized lower triangular matrix\n",
    "wei = wei / wei.sum(1, keepdim=True)\n",
    "\n",
    "xbow2 = wei @ x # (B (auto broadcasted by torch), T, T) @ (B, T, C) --> (B, T, C)\n",
    "print(f\"wei: {wei}\")\n",
    "print(f\"xbow2: {xbow2}\")\n",
    "torch.allclose(xbow, xbow2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version 2.2: get weights by softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "by turning all zeros to -inf in a lower tri mat, then applying softmax to row, we can get the same weights\n",
    "- another advantage of softmax is that it ensures **non negative weights**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
       "        [1., 1., -inf, -inf, -inf, -inf, -inf, -inf],\n",
       "        [1., 1., 1., -inf, -inf, -inf, -inf, -inf],\n",
       "        [1., 1., 1., 1., -inf, -inf, -inf, -inf],\n",
       "        [1., 1., 1., 1., 1., -inf, -inf, -inf],\n",
       "        [1., 1., 1., 1., 1., 1., -inf, -inf],\n",
       "        [1., 1., 1., 1., 1., 1., 1., -inf],\n",
       "        [1., 1., 1., 1., 1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei = torch.tril(torch.ones(T, T)) # T by T lower tri mat\n",
    "wei = wei.masked_fill(wei==0, float('-inf'))\n",
    "wei"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],\n",
       "        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],\n",
       "        [0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei = F.softmax(wei, dim=1)\n",
    "wei"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementing the averaging head module "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AveragingHead(nn.Module):\n",
    "    \"\"\" one head of naive aggregation \"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape\n",
    "        wei = torch.tril(torch.ones(T, T))\n",
    "        wei = F.softmax(wei.masked_fill(wei==0, float('-inf')), dim=1) # T by T lower tri mat\n",
    "        agg_x = wei @ x # (T, T) @ (B, T, C) --> (B, T, C)\n",
    "        return agg_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: tensor([[[0.5820, 0.1338, 0.7995],\n",
      "         [0.3071, 0.6526, 0.6105],\n",
      "         [0.1575, 0.6983, 0.7883]]])\n",
      "agg_x: tensor([[[0.5820, 0.1338, 0.7995],\n",
      "         [0.4446, 0.3932, 0.7050],\n",
      "         [0.3489, 0.4949, 0.7328]]])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(500)\n",
    "x = torch.rand(1, 3, 3)\n",
    "\n",
    "ah = AveragingHead()\n",
    "agg_x = ah(x)\n",
    "print(f'x: {x}')\n",
    "print(f'agg_x: {agg_x}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding aggregation head to Bigram LM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### A few changes on BigramLanguageModelV2 from BigramLanguageModel\n",
    "1. add position_embedding_table (along T / block_size axis) to capture positional info\n",
    "2. parameterize n_embed in embedding tables to configure # dimensions of embedding vectors\n",
    "3. add AggregationHead to establish the weakest form of communication between upper context of text\n",
    "4. Now that we have implemented positional embedding, we cannot feed idx longer than block_size, else we will get index out of range when accessing the positional embedding table. idx is cropped to the last block_size block during generate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLanguageModelWithAveragingHead(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, n_embed):\n",
    "        super().__init__()\n",
    "        \n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
    "        self.head = AveragingHead()\n",
    "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        # idx and targets are both (B, T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B, T, n_embed)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T)) # (T, n_embed)\n",
    "        x = tok_emb + pos_emb # (B, T, C)\n",
    "        agg_x = self.head(x) # (B, C, C)\n",
    "        logits = self.lm_head(agg_x)  # (B, T, vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits & targets into 2D to cater for F.corss_entropy\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens): # idx is (B, T) array of indices in the current context\n",
    "        \n",
    "        # generate max_new_tokens tokens iteratively by looking at only the last token each time\n",
    "        for _ in range(max_new_tokens):\n",
    "            \n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cropped = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cropped)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OptimizedModule(\n",
       "  (_orig_mod): BigramLanguageModelWithAveragingHead(\n",
       "    (token_embedding_table): Embedding(65, 32)\n",
       "    (position_embedding_table): Embedding(8, 32)\n",
       "    (head): AveragingHead()\n",
       "    (lm_head): Linear(in_features=32, out_features=65, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1 = torch.compile(BigramLanguageModelWithAveragingHead(vocab_size, n_embed=32))\n",
    "m1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.906489133834839\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "train(m1, 5000, batch_size, block_size, lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the training result is worse than vanilla bigram, indicating averaging is a bad communication mechanism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "UCHERMNYA:I,\n",
      "\n",
      " rCSRhree\n",
      "e flaf , Wadh chuvmoa,s t srboowe ntl,\n",
      "\n",
      "hIa Rmicangeeke.\n",
      ",,\n",
      "PUOEAIFUTRN IOWT\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m1.generate(idx, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Version 4: self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 32])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "B, T, C = 4, 8, 32\n",
    "x = torch.randn(B, T, C)\n",
    "\n",
    "# let's see a single Head performing self-attention\n",
    "head_size = 16 # output size of the Linear layers\n",
    "key = nn.Linear(C, head_size, bias=False)\n",
    "query = nn.Linear(C, head_size, bias=False)\n",
    "k = key(x) # (B, T, 16)\n",
    "q = query(x) # (B, T, 16)\n",
    "wei = q @ k.transpose(-2,-1) # (B, T, 16) @ (B, 16, T) --> (B, T, T)\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",
    "out = wei @ x\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5792, 0.1187, 0.1889, 0.1131, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0294, 0.1052, 0.0469, 0.0276, 0.7909, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0176, 0.2689, 0.0215, 0.0089, 0.6812, 0.0019, 0.0000, 0.0000],\n",
       "        [0.1691, 0.4066, 0.0438, 0.0416, 0.1048, 0.2012, 0.0329, 0.0000],\n",
       "        [0.0210, 0.0843, 0.0555, 0.2297, 0.0573, 0.0709, 0.2423, 0.2391]],\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementing self-attention head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "    def __init__(self, n_embed, head_size, block_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embed, head_size, bias=False) # (C, head_size)\n",
    "        self.query = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape\n",
    "        k = self.key(x) # (B, T, C)\n",
    "        q = self.query(x) # (B, T, C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2, -1) # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
    "        wei *= C**-0.5 # Scaled attention\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=2) # (B, T, T)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B, T, C)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding self-attention to bigram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLanguageModelWithSelfAttentionHead(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, n_embed):\n",
    "        super().__init__()\n",
    "        \n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
    "        self.self_attention_head = Head(n_embed, n_embed, block_size)\n",
    "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        # idx and targets are both (B, T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B, T, n_embed)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T)) # (T, n_embed)\n",
    "        x = tok_emb + pos_emb # (B, T, C)\n",
    "        x = self.self_attention_head(x) # single head self-attention (B, T, C)\n",
    "        logits = self.lm_head(x)  # (B, T, vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits & targets into 2D to cater for F.corss_entropy\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens): # idx is (B, T) array of indices in the current context\n",
    "        \n",
    "        # generate max_new_tokens tokens iteratively by looking at only the last token each time\n",
    "        for _ in range(max_new_tokens):\n",
    "            \n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cropped = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cropped)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OptimizedModule(\n",
       "  (_orig_mod): BigramLanguageModelWithSelfAttentionHead(\n",
       "    (token_embedding_table): Embedding(65, 32)\n",
       "    (position_embedding_table): Embedding(8, 32)\n",
       "    (self_attention_head): Head(\n",
       "      (key): Linear(in_features=32, out_features=32, bias=False)\n",
       "      (query): Linear(in_features=32, out_features=32, bias=False)\n",
       "      (value): Linear(in_features=32, out_features=32, bias=False)\n",
       "    )\n",
       "    (lm_head): Linear(in_features=32, out_features=65, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m2 = torch.compile(BigramLanguageModelWithSelfAttentionHead(vocab_size, n_embed=32))\n",
    "m2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.420104503631592\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "train(m2, 5000, batch_size, block_size, lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "training loss is significantly lower than the previous 2 versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Prn our fo,\n",
      "A gat hree w fo f hat dr chevesars t sto the nto,\n",
      "YTI shancange iem t, thand ther ater\n",
      "Se m,\n",
      "S:\n",
      "Bawe th my:\n",
      "O:\n",
      "Ad onosnienkfe wick eanicth xjom\n",
      "Ange wheso aly, sod pre Gd, sho hall.\n",
      "Withalle fo.\n",
      "\n",
      "whes nde shan ld\n",
      "Th men ro toheroans, t\n",
      "ISLAShor perre, dh troussa,\n",
      "s. Wine.\n",
      "\n",
      "Avised whto ter\n",
      "He elovereivin ghad for te I!\n",
      "Fimpgenck meas mecon. Whansat, has by blarnet fine barry iveam famirera matoke, ay bret hatout in ser!\n",
      "Mous hainisithan meare dooutanth, tharave wan,\n",
      "O NGRlo I ay degorse ardithatsilyofnom, dis,\n",
      "Thory;\n",
      "Ih wat th tuy an faveencolery mivomild, pre matt dee Gr\n",
      "Wharf take ourgther ot, fi\n",
      "\n",
      "ot\n",
      "llbe the ler ary: FiciO Lowrens owrocan dor fr.\n",
      "\n",
      "That ongatinof thesr pay y Cane you: than\n",
      "Bind retel;\n",
      "Gul noume sen:\n",
      "\n",
      "JUSTor V:\n",
      "s melove whakis ar tearem Mank sefat, dyot. Whef oubve fathan hpir l:\n",
      "D We sthithise kned wove thilaw wimard ols the,\n",
      "BO, st\n",
      "Jlf the secth I takl feat burcerielert gt oo, plathound t-y mausam ney oorthe.\n",
      "\n",
      "F: dos wilerdo osidoud eothal yea thirog tis\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m2.generate(idx, max_new_tokens=1000)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementing multi-headed self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHead(nn.Module):\n",
    "    \"\"\" multiple heads of self-attention \"\"\"\n",
    "    def __init__(self, num_heads, n_embed, head_size, block_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(n_embed, head_size, block_size) for _ in range(num_heads)])\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([h(x) for h in self.heads], dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adding multi-headed self-attention to bigram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLanguageModelWithMultiHeadedSelfAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, n_embed, num_heads):\n",
    "        super().__init__()\n",
    "        \n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
    "        self.self_attention_heads = MultiHead(num_heads, n_embed, n_embed//num_heads, block_size) # output dimension of MultiHead --> num_heads * _embed//num_heads = n_embed\n",
    "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        # idx and targets are both (B, T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B, T, n_embed)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T)) # (T, n_embed)\n",
    "        x = tok_emb + pos_emb # (B, T, C)\n",
    "        x = self.self_attention_heads(x) # single head self-attention (B, T, C)\n",
    "        logits = self.lm_head(x)  # (B, T, vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits & targets into 2D to cater for F.corss_entropy\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens): # idx is (B, T) array of indices in the current context\n",
    "        \n",
    "        # generate max_new_tokens tokens iteratively by looking at only the last token each time\n",
    "        for _ in range(max_new_tokens):\n",
    "            \n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cropped = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cropped)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OptimizedModule(\n",
       "  (_orig_mod): BigramLanguageModelWithMultiHeadedSelfAttention(\n",
       "    (token_embedding_table): Embedding(65, 32)\n",
       "    (position_embedding_table): Embedding(8, 32)\n",
       "    (self_attention_heads): MultiHead(\n",
       "      (heads): ModuleList(\n",
       "        (0-3): 4 x Head(\n",
       "          (key): Linear(in_features=32, out_features=8, bias=False)\n",
       "          (query): Linear(in_features=32, out_features=8, bias=False)\n",
       "          (value): Linear(in_features=32, out_features=8, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (lm_head): Linear(in_features=32, out_features=65, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m3 = torch.compile(BigramLanguageModelWithMultiHeadedSelfAttention(vocab_size, n_embed=32, num_heads=4))\n",
    "m3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.334955930709839\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "train(m3, 5000, batch_size, block_size, lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "PUERENY:\n",
      "I,\n",
      "A gat hree wofl fu, Wa:\n",
      "Thy ves', cresto the now,\n",
      "Youg hancangeeke.\n",
      "\n",
      "And lit ther ate wis my\n",
      "Surcake them sarabud ongingenkfe wick elficty ximm\n",
      "An powe so all be disgen da shooush hing younto you whis nome hat nor so'se rove heroansten\n",
      "I nou of perrt, du trousst,\n",
      "st gine.\n",
      "\n",
      "Dvired?\n",
      "\n",
      "So ter\n",
      "beeel asesivin ghakset. te I!\n",
      "ENF:\n",
      "Olom meds! For mast beat, has by bartnet fie? batry if?\n",
      "man miresa matwes, an bret heacut inf mathus ifh weistuthape hin dooutarth, de Cave want nothtlarenay de of e and thatss coffor, dis,\n",
      "Thow wore wat ther yiue cke mecelerth honcill, prepeatt dee Gut;\n",
      "Anfirk'stos Mand you, wither\n",
      "labe the lis 'sy: UCcie adwe.\n",
      "\n",
      "K\n",
      "BET:\n",
      "By ou fromranto ongatinof thest pay youcen you: thaven homr teod\n",
      "she mry bus me of and galds meleven hak ande teare.\n",
      "MBak sefar, delt. Wheito bve fath and in your ge sabithis okned wove thiln cot and of nowe,\n",
      "sinss lall the se they de lis waist cers lert gon of plath munst-yerdusampeey oortve.\n",
      "\n",
      "Ford shis bret on dour eot su ye, Ving go in\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m3.generate(idx, max_new_tokens=1000)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Remaining transformer components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feedforward layers of transformer block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, n_embed):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embed, n_embed*4), # following \"Attention is all you need\" -> hidden layer size = 4 * input size\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(n_embed*4, n_embed),\n",
    "        )\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The transformer block\n",
    "\n",
    "For each block, we self attend, then feed forward, interspersed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, n_embed, n_head):\n",
    "        super().__init__()\n",
    "        head_size = n_embed // n_head\n",
    "        self.self_attention = MultiHead(n_head, n_embed, head_size, block_size)\n",
    "        self.feed_forward = FeedForward(n_embed)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.self_attention(x)\n",
    "        x = self.feed_forward(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Residual Connections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, n_embed, n_head):\n",
    "        super().__init__()\n",
    "        head_size = n_embed // n_head\n",
    "        self.self_attention = MultiHead(n_head, n_embed, head_size, block_size)\n",
    "        self.feed_forward = FeedForward(n_embed)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.self_attention(x) # residual\n",
    "        x = x + self.feed_forward(x) # residual\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHead(nn.Module):\n",
    "    \n",
    "    def __init__(self, num_heads, n_embed, head_size, block_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(n_embed, head_size, block_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embed, n_embed) # linear transformation to project self-attention outputs back to the residual pathway\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        x = self.proj(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layernorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, n_embed, n_head):\n",
    "        super().__init__()\n",
    "        head_size = n_embed // n_head\n",
    "        self.self_attention = MultiHead(n_head, n_embed, head_size, block_size)\n",
    "        self.feed_forward = FeedForward(n_embed)\n",
    "        self.layer_norm_1 = nn.LayerNorm(n_embed)\n",
    "        self.layer_norm_2 = nn.LayerNorm(n_embed)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.self_attention(self.layer_norm_1(x)) # applies layer norm BEFORE attention (studies suggest this is better than the original architecture)\n",
    "        x = x + self.feed_forward(self.layer_norm_2(x)) # applies layer norm BEFORE ffwd (studies suggest this is better than the original architecture\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, n_embed, dropout):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embed, n_embed*4), # following \"Attention is all you need\" -> hidden layer size = 4 * input size\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(n_embed*4, n_embed),\n",
    "        )\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "    def __init__(self, n_embed, head_size, block_size, dropout):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "        self.dropout = nn.Dropout(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",
    "\n",
    "        wei = q @ k.transpose(-2, -1)\n",
    "        wei *= C**-0.5\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        wei = F.softmax(wei, dim=2)\n",
    "        wei = self.dropout(wei) # dropout\n",
    "        \n",
    "        v = self.value(x)\n",
    "        out = wei @ v\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHead(nn.Module):\n",
    "    \n",
    "    def __init__(self, num_heads, n_embed, head_size, block_size, dropout):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(n_embed, head_size, block_size, dropout) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embed, n_embed) # linear transformation to project self-attention outputs back to the residual pathway\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        x = self.proj(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, n_head, n_embed, block_size, dropout):\n",
    "        super().__init__()\n",
    "        head_size = n_embed // n_head\n",
    "        self.self_attention = MultiHead(n_head, n_embed, head_size, block_size, dropout)\n",
    "        self.feed_forward = FeedForward(n_embed, dropout)\n",
    "        self.layer_norm_1 = nn.LayerNorm(n_embed)\n",
    "        self.layer_norm_2 = nn.LayerNorm(n_embed)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.self_attention(self.layer_norm_1(x)) # applies layer norm BEFORE attention (studies suggest this is better than the original architecture)\n",
    "        x = self.dropout(x)\n",
    "        x = x + self.feed_forward(self.layer_norm_2(x)) # applies layer norm BEFORE ffwd (studies suggest this is better than the original architecture\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT(nn.Module):\n",
    "\n",
    "    def __init__(self, block_size, n_embed, n_head, n_block, dropout, device):\n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
    "        self.register_buffer('positional_intervals', torch.arange(block_size, device=self.device))\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
    "        self.blocks = nn.Sequential(\n",
    "            *[Block(n_head, n_embed, block_size, dropout) for _ in range(n_block)]\n",
    "        )\n",
    "        self.ln_f = nn.LayerNorm(n_embed) # final layer norm\n",
    "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        # idx and targets are both (B, T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
    "        pos_emb = self.position_embedding_table(self.positional_intervals[:T]) # (T,C)\n",
    "        x = tok_emb + pos_emb # (B,T,C)\n",
    "        x = self.blocks(x)\n",
    "        x = self.ln_f(x)\n",
    "        logits = self.lm_head(x)  # (B, T, vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits & targets into 2D to cater for F.corss_entropy\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens): # idx is (B, T) array of indices in the current context\n",
    "        \n",
    "        # generate max_new_tokens tokens iteratively by looking at only the last token each time\n",
    "        for _ in range(max_new_tokens):\n",
    "            \n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cropped = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cropped)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, steps, batch_size, block_size, lr, eval_interval, eval_iters, device):\n",
    "    @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,  batch_size, block_size)\n",
    "                X, Y = X.to(device), Y.to(device)\n",
    "                logits, loss = model(X, Y)\n",
    "                losses[k] = loss.item()\n",
    "            out[split] = losses.mean()\n",
    "        model.train()\n",
    "        return out\n",
    "\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr)\n",
    "    for iter in range(steps):\n",
    "\n",
    "        # every once in a while evaluate the loss on train and val sets\n",
    "        if iter % eval_interval == 0:\n",
    "            losses = estimate_loss()\n",
    "            print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "\n",
    "        # sample a batch of data\n",
    "        xb, yb = get_batch('train', batch_size, block_size)\n",
    "        xb, yb = xb.to(device), yb.to(device)\n",
    "        \n",
    "        # evaluate the loss\n",
    "        logits, loss = model(xb, yb)\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.606145 M parameters\n"
     ]
    }
   ],
   "source": [
    "# hyperparameters\n",
    "batch_size = 64\n",
    "block_size = 256\n",
    "max_iters = 20000\n",
    "eval_interval = 500\n",
    "learning_rate = 3e-4\n",
    "device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'\n",
    "eval_iters = 200\n",
    "n_embed = 160\n",
    "n_head = 5\n",
    "n_block = 5\n",
    "dropout = 0.2\n",
    "# ------------\n",
    "\n",
    "babyGPT = GPT(block_size, n_embed, n_head, n_block, dropout, device)\n",
    "babyGPT.to(device)\n",
    "print(sum(p.numel() for p in babyGPT.parameters())/1e6, 'M parameters')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 4.3397, val loss 4.3301\n",
      "step 500: train loss 2.4637, val loss 2.4719\n",
      "step 1000: train loss 2.3547, val loss 2.3718\n",
      "step 1500: train loss 2.2414, val loss 2.2755\n",
      "step 2000: train loss 2.1253, val loss 2.1733\n",
      "step 2500: train loss 2.0056, val loss 2.0844\n",
      "step 3000: train loss 1.9044, val loss 2.0143\n",
      "step 3500: train loss 1.8153, val loss 1.9428\n",
      "step 4000: train loss 1.7483, val loss 1.8954\n",
      "step 4500: train loss 1.6949, val loss 1.8594\n",
      "step 5000: train loss 1.6516, val loss 1.8243\n",
      "step 5500: train loss 1.6170, val loss 1.7879\n",
      "step 6000: train loss 1.5824, val loss 1.7640\n",
      "step 6500: train loss 1.5572, val loss 1.7419\n",
      "step 7000: train loss 1.5296, val loss 1.7198\n",
      "step 7500: train loss 1.5132, val loss 1.7023\n",
      "step 8000: train loss 1.4948, val loss 1.6882\n",
      "step 8500: train loss 1.4728, val loss 1.6734\n",
      "step 9000: train loss 1.4636, val loss 1.6648\n",
      "step 9500: train loss 1.4457, val loss 1.6466\n",
      "step 10000: train loss 1.4340, val loss 1.6336\n",
      "step 10500: train loss 1.4220, val loss 1.6319\n",
      "step 11000: train loss 1.4055, val loss 1.6193\n",
      "step 11500: train loss 1.4102, val loss 1.6197\n",
      "step 12000: train loss 1.3944, val loss 1.6069\n",
      "step 12500: train loss 1.3887, val loss 1.6000\n",
      "step 13000: train loss 1.3774, val loss 1.5886\n",
      "step 13500: train loss 1.3758, val loss 1.5932\n",
      "step 14000: train loss 1.3664, val loss 1.5804\n",
      "step 14500: train loss 1.3625, val loss 1.5799\n",
      "step 15000: train loss 1.3473, val loss 1.5718\n",
      "step 15500: train loss 1.3455, val loss 1.5761\n",
      "step 16000: train loss 1.3401, val loss 1.5634\n",
      "step 16500: train loss 1.3413, val loss 1.5594\n",
      "step 17000: train loss 1.3336, val loss 1.5579\n",
      "step 17500: train loss 1.3291, val loss 1.5606\n",
      "step 18000: train loss 1.3202, val loss 1.5526\n",
      "step 18500: train loss 1.3154, val loss 1.5542\n",
      "step 19000: train loss 1.3171, val loss 1.5508\n",
      "step 19500: train loss 1.3087, val loss 1.5407\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "train(babyGPT, max_iters, batch_size, block_size, learning_rate, eval_interval, eval_iters, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "broathed's construns of that love, which are--\n",
      "\n",
      "GLOUCESTER:\n",
      "And he own cuntest of his hounds, poor country,-\n",
      "With honest lose, he was it sincer.\n",
      "\n",
      "NORTLANNE:\n",
      "How not to says from his lord's guilt.\n",
      "Nothing munis is toesth heme,\n",
      "Rounce to nare have cold.\n",
      "\n",
      "GLOUCESTER:\n",
      "But how ip, nor morue way dear who spay, and I\n",
      "Thou, somence that dath knowns, and for it the Claudio;\n",
      "Thou clook, and, way to merry, hell a Volst's mock of\n",
      "a grepard and the king of sorrower the treto the fhight\n",
      "varter's habste son; a sunce of yet combriness are gone stouch,\n",
      "my poor to-mooth, sir have, oil.\n",
      "Now officery, proner, we men you come in the vient tower,\n",
      "Know have a kised in succeite used\n",
      "To Onch: no childs a parsuest, where your goodlish!\n",
      "\n",
      "HENRY PELIZABOL:\n",
      "Ay, a city molece took dayying worse more:\n",
      "Camilonius long pose it; farewell.\n",
      "\n",
      "QUEEN MARGARET:\n",
      "Come, goes my ging trough, from you my head.\n",
      "\n",
      "KING EDWE RICHARD IV:\n",
      "Not-none. Boy: shief, my love,\n",
      "Lord lord, art it is nothier are than he your had,\n",
      "With hast him\n",
      "The\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long, device=device)\n",
    "print(decode(babyGPT.generate(idx, max_new_tokens=1000)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(babyGPT, f'babyGPT_{max_iters}_steps.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closing notes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few departures from the GPT models used in OpenAI, for simplicity's sake:\n",
    "- In the feedforward layer, ReLU is used here instead of GELU\n",
    "- a simple character level tokenizer is used here, while OpenAI uses the more sophisticated tiktoken"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GPT-from-Scratch",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}