Playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zsh:1: command not found: wget\n"
     ]
    }
   ],
   "source": [
    "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
    "! wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('input.txt') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of text: 1115394\n"
     ]
    }
   ],
   "source": [
    "print(\"Length of text: {}\".format(len(text)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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"
     ]
    }
   ],
   "source": [
    "print(text[:250])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 65 unique characters in the dataset\n",
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n"
     ]
    }
   ],
   "source": [
    "# The unique characters in the file\n",
    "vocab = sorted(list(set(text)))\n",
    "vocab_size = len(vocab)\n",
    "print(\"There are {} unique characters in the dataset\".format(vocab_size))\n",
    "print(\"\".join(vocab))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 53, 50, 39, 1, 51, 39, 39, 51, 47]\n",
      "hola maami\n"
     ]
    }
   ],
   "source": [
    "# Mapping of chrachters to indices and vice versa\n",
    "chr_idx_map = {c: i for i, c in enumerate(vocab)}\n",
    "idx_chr_map = {i: c for i, c in enumerate(vocab)}\n",
    "# Later we will use our custom tokenizer\n",
    "encode =  lambda x: [chr_idx_map[c] for c in x]\n",
    "decode = lambda x: ''.join([idx_chr_map[i] for i in x])\n",
    "\n",
    "print(encode('hola maami'))\n",
    "print(decode(encode('hola maami')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment with tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115394])\n",
      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47])\n"
     ]
    }
   ],
   "source": [
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape)\n",
    "print(data[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Test split of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = data.size(0)\n",
    "train_data = data[:int(n*0.9)]\n",
    "val_data = data[int(n*0.9):]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([69, 15, 26, 66, 18, 15, 75, 26])\n",
      "tensor([69]) -> 15\n",
      "tensor([69, 15]) -> 26\n",
      "tensor([69, 15, 26]) -> 66\n",
      "tensor([69, 15, 26, 66]) -> 18\n",
      "tensor([69, 15, 26, 66, 18]) -> 15\n",
      "tensor([69, 15, 26, 66, 18, 15]) -> 75\n",
      "tensor([69, 15, 26, 66, 18, 15, 75]) -> 26\n"
     ]
    }
   ],
   "source": [
    "block_size = 8\n",
    "\n",
    "x = torch.randint(1,90, (block_size,))\n",
    "print(x)\n",
    "for i in range(len(x)):\n",
    "    if i!=len(x)-1:\n",
    "        print(f\"{x[:i+1]} -> {x[i+1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([18]) -> 47\n",
      "tensor([18, 47]) -> 56\n",
      "tensor([18, 47, 56]) -> 57\n",
      "tensor([18, 47, 56, 57]) -> 58\n",
      "tensor([18, 47, 56, 57, 58]) -> 1\n",
      "tensor([18, 47, 56, 57, 58,  1]) -> 15\n",
      "tensor([18, 47, 56, 57, 58,  1, 15]) -> 47\n",
      "tensor([18, 47, 56, 57, 58,  1, 15, 47]) -> 58\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "\n",
    "for i in range(len(x)):\n",
    "    print(f\"{x[:i+1]} -> {y[i]}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data batch loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[15,  7, 42,  0, 45,  3, 15, 10],\n",
      "        [34, 40, 12, 20, 47, 36, 40, 38],\n",
      "        [ 1, 14,  9, 15, 13, 46, 22,  0],\n",
      "        [32, 11, 46,  5, 39,  4,  5, 49]])\n",
      "tensor([15]) -> 7\n",
      "tensor([15,  7]) -> 42\n",
      "tensor([15,  7, 42]) -> 0\n",
      "tensor([15,  7, 42,  0]) -> 45\n",
      "tensor([15,  7, 42,  0, 45]) -> 3\n",
      "tensor([15,  7, 42,  0, 45,  3]) -> 15\n",
      "tensor([15,  7, 42,  0, 45,  3, 15]) -> 10\n",
      "----\n",
      "tensor([34]) -> 40\n",
      "tensor([34, 40]) -> 12\n",
      "tensor([34, 40, 12]) -> 20\n",
      "tensor([34, 40, 12, 20]) -> 47\n",
      "tensor([34, 40, 12, 20, 47]) -> 36\n",
      "tensor([34, 40, 12, 20, 47, 36]) -> 40\n",
      "tensor([34, 40, 12, 20, 47, 36, 40]) -> 38\n",
      "----\n",
      "tensor([1]) -> 14\n",
      "tensor([ 1, 14]) -> 9\n",
      "tensor([ 1, 14,  9]) -> 15\n",
      "tensor([ 1, 14,  9, 15]) -> 13\n",
      "tensor([ 1, 14,  9, 15, 13]) -> 46\n",
      "tensor([ 1, 14,  9, 15, 13, 46]) -> 22\n",
      "tensor([ 1, 14,  9, 15, 13, 46, 22]) -> 0\n",
      "----\n",
      "tensor([32]) -> 11\n",
      "tensor([32, 11]) -> 46\n",
      "tensor([32, 11, 46]) -> 5\n",
      "tensor([32, 11, 46,  5]) -> 39\n",
      "tensor([32, 11, 46,  5, 39]) -> 4\n",
      "tensor([32, 11, 46,  5, 39,  4]) -> 5\n",
      "tensor([32, 11, 46,  5, 39,  4,  5]) -> 49\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4\n",
    "block_size = 8\n",
    "\n",
    "x = torch.randint(0,50,(batch_size, block_size))\n",
    "print(x)\n",
    "for b in range(batch_size):\n",
    "    for t in range(block_size):\n",
    "        if t!=0:\n",
    "            print(f\"{x[b,:t]} -> {x[b,t]}\")\n",
    "    print(\"----\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      "torch.Size([4, 8])\n",
      "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
      "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
      "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
      "        [25, 17, 27, 10,  0, 21,  1, 54]])\n",
      "targets:\n",
      "torch.Size([4, 8])\n",
      "tensor([[43, 58,  5, 57,  1, 46, 43, 39],\n",
      "        [53, 56,  1, 58, 46, 39, 58,  1],\n",
      "        [58,  1, 58, 46, 39, 58,  1, 46],\n",
      "        [17, 27, 10,  0, 21,  1, 54, 39]])\n",
      "----\n",
      "when input is [24] the target: 43\n",
      "when input is [24, 43] the target: 58\n",
      "when input is [24, 43, 58] the target: 5\n",
      "when input is [24, 43, 58, 5] the target: 57\n",
      "when input is [24, 43, 58, 5, 57] the target: 1\n",
      "when input is [24, 43, 58, 5, 57, 1] the target: 46\n",
      "when input is [24, 43, 58, 5, 57, 1, 46] the target: 43\n",
      "when input is [24, 43, 58, 5, 57, 1, 46, 43] the target: 39\n",
      "when input is [44] the target: 53\n",
      "when input is [44, 53] the target: 56\n",
      "when input is [44, 53, 56] the target: 1\n",
      "when input is [44, 53, 56, 1] the target: 58\n",
      "when input is [44, 53, 56, 1, 58] the target: 46\n",
      "when input is [44, 53, 56, 1, 58, 46] the target: 39\n",
      "when input is [44, 53, 56, 1, 58, 46, 39] the target: 58\n",
      "when input is [44, 53, 56, 1, 58, 46, 39, 58] the target: 1\n",
      "when input is [52] the target: 58\n",
      "when input is [52, 58] the target: 1\n",
      "when input is [52, 58, 1] the target: 58\n",
      "when input is [52, 58, 1, 58] the target: 46\n",
      "when input is [52, 58, 1, 58, 46] the target: 39\n",
      "when input is [52, 58, 1, 58, 46, 39] the target: 58\n",
      "when input is [52, 58, 1, 58, 46, 39, 58] the target: 1\n",
      "when input is [52, 58, 1, 58, 46, 39, 58, 1] the target: 46\n",
      "when input is [25] the target: 17\n",
      "when input is [25, 17] the target: 27\n",
      "when input is [25, 17, 27] the target: 10\n",
      "when input is [25, 17, 27, 10] the target: 0\n",
      "when input is [25, 17, 27, 10, 0] the target: 21\n",
      "when input is [25, 17, 27, 10, 0, 21] the target: 1\n",
      "when input is [25, 17, 27, 10, 0, 21, 1] the target: 54\n",
      "when input is [25, 17, 27, 10, 0, 21, 1, 54] the target: 39\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4\n",
    "block_size = 8\n",
    "\n",
    "def get_data_batch(split):\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_data_batch('train')\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): # batch dimension\n",
    "    for t in range(block_size): # time dimension\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b,t]\n",
    "        print(f\"when input is {context.tolist()} the target: {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Experimentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn \n",
    "import torch.nn.functional as F \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 65])\n",
      "4.534911632537842\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"\\nKH,FlD!Hc&jWbYUhie n PPCdpVzu\\nnH3$hAuGFKmOW!'Ns XANf;MwPpGC.!o,e ttX.!\\nzTr&higRKASjOtE\\nKu&l;zt\\n;;MH3\""
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class TimePassGPT(nn.Module):\n",
    "    def __init__(self, vocab_size) -> None:\n",
    "        super().__init__()\n",
    "        self.token_embedding = nn.Embedding(vocab_size, vocab_size)\n",
    "        \n",
    "    def forward(self, idx, targets=None):\n",
    "        # idx: (B, T), targets: (B, T)\n",
    "        logits = self.token_embedding(idx)\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\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",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx: (B, T)\n",
    "        with torch.no_grad():\n",
    "            for _ in range(max_new_tokens):\n",
    "                logits, loss = self(idx)\n",
    "                # logits: (B, T, C)\n",
    "                logits = logits[:, -1, :] \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",
    "                # print(idx)\n",
    "        return idx\n",
    "model = TimePassGPT(vocab_size)\n",
    "logits, loss = model(xb, yb)\n",
    "\n",
    "print(logits.shape)\n",
    "print(loss.item())\n",
    "decode(model.generate(idx=torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist())\n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step: 0, Loss: 4.513947010040283\n",
      "Step: 10, Loss: 4.530942440032959\n",
      "Step: 20, Loss: 4.499247074127197\n",
      "Step: 30, Loss: 4.458986282348633\n",
      "Step: 40, Loss: 4.393402576446533\n",
      "Step: 50, Loss: 4.463761806488037\n",
      "Step: 60, Loss: 4.517825126647949\n",
      "Step: 70, Loss: 4.476200103759766\n",
      "Step: 80, Loss: 4.521479606628418\n",
      "Step: 90, Loss: 4.386441707611084\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "for steps in range(100):\n",
    "    xb, yb = get_data_batch('train')\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if steps % 10 == 0:\n",
    "        print(f\"Step: {steps}, Loss: {loss.item()}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "kdpABz!FUZ?ZQUk&iGeg!VFGFj?!vznQ,TLbnRoHkGmOv-O;mg3g3db;'NzTe;bc&ELiAaXa,llW3dVNEavN'!J!oAN'BrbHuSIGxRMI$OJIgaZwz!DxbMSBhJbNQQ;NzViXaw&33!sUcOFfDH's-Ik:x:dbl-Oz\n",
      "PP&;,'ovon$VNJmmC-aOd&CRo&TUyS'?Px3;-g?s-phhj AsY.3ExFs ezTLunFvSn:TU!VsYXHLuP.\n",
      " AZkymmCXjdpIpVPNZM3mO?fHXGeugeyFlpIPpyNVWaZn'rKiglWTLFsXZygR!Xw ThieDO-PVHg3dxnx?sEEwjkvtkLQIJWu\n",
      ":gPpiZZI.Znbqp\n",
      "FlsVThyjBwQHg3QxM'VH3IinaOkOJXmOFleXg.\n",
      "XGF;l'YJ-dgTLYQjjs tII.uu\n",
      "iZE:&!,.!nHg3;?A$pg3ekftm-T$aarRZamO&gKPvbmREcMH3;-RWI!;zT&Ka&gWauGFJKlQ;;bwBrsng\n"
     ]
    }
   ],
   "source": [
    "print(decode(model.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Brute Force of averaging previous vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0, 3],\n",
      "        [8, 4],\n",
      "        [9, 7],\n",
      "        [5, 0]])\n",
      "Token 0: tensor([0, 3]) -> tensor([0., 3.])\n",
      "Token 1: tensor([8, 4]) -> tensor([4.0000, 3.5000])\n",
      "Token 2: tensor([9, 7]) -> tensor([5.6667, 4.6667])\n",
      "Token 3: tensor([5, 0]) -> tensor([5.5000, 3.5000])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randint(0,10, (4,2)) # block_size = 4, channels = 2\n",
    "# Each token is represented by a vector of size 2 and there are 4 tokens\n",
    "# Attention - average of the previous tokens across the dimensions\n",
    "print(x)\n",
    "for i in range(x.size(0)):\n",
    "    print(f\"Token {i}: {x[i]} -> {x[:i+1].mean(dim=0, dtype=torch.float)}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Matrix multiplication Tricks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2., 4.],\n",
      "        [1., 1.],\n",
      "        [2., 3.],\n",
      "        [4., 3.]])\n",
      "-----\n",
      "tensor([[1., 0., 0., 0.],\n",
      "        [1., 1., 0., 0.],\n",
      "        [1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1.]])\n",
      "-----\n",
      "tensor([[ 2.,  4.],\n",
      "        [ 3.,  5.],\n",
      "        [ 5.,  8.],\n",
      "        [ 9., 11.]])\n",
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333, 0.0000],\n",
      "        [0.2500, 0.2500, 0.2500, 0.2500]])\n",
      "tensor([[2.0000, 4.0000],\n",
      "        [1.5000, 2.5000],\n",
      "        [1.6667, 2.6667],\n",
      "        [2.2500, 2.7500]])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randint(0,5, (4,2), dtype=torch.float)\n",
    "wei = torch.tril(torch.ones(4,4))\n",
    "print(x)\n",
    "print(\"-----\")\n",
    "print(wei)\n",
    "print(\"-----\")\n",
    "# Multiplying using a lower triangular matrix will make the tokens to attend to the previous tokens\n",
    "print(wei@x)\n",
    "# We can avergae it out by summing the weights in each layer\n",
    "print(wei/torch.sum(wei, dim=1,keepdim=True))\n",
    "wei1 = wei/torch.sum(wei, dim=1,keepdim=True)\n",
    "print(wei1@x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[4., 4.],\n",
      "        [3., 3.],\n",
      "        [1., 4.],\n",
      "        [4., 4.]])\n",
      "-----\n",
      "tensor([[1., 0., 0., 0.],\n",
      "        [1., 1., 0., 0.],\n",
      "        [1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1.]])\n",
      "-----\n",
      "tensor([[ 4.,  4.],\n",
      "        [ 7.,  7.],\n",
      "        [ 8., 11.],\n",
      "        [12., 15.]])\n",
      "tensor([[1.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333, 0.0000],\n",
      "        [0.2500, 0.2500, 0.2500, 0.2500]])\n",
      "tensor([[4.0000, 4.0000],\n",
      "        [3.5000, 3.5000],\n",
      "        [2.6667, 3.6667],\n",
      "        [3.0000, 3.7500]])\n"
     ]
    }
   ],
   "source": [
    "# We can use softmax as well for finding the attention weights\n",
    "x = torch.randint(0,5, (4,2), dtype=torch.float)\n",
    "wei = torch.tril(torch.ones(4,4))\n",
    "print(x)\n",
    "print(\"-----\")\n",
    "print(wei)\n",
    "print(\"-----\")\n",
    "# Multiplying using a lower triangular matrix will make the tokens to attend to the previous tokens\n",
    "print(wei@x)\n",
    "# using softmax\n",
    "wei = wei.masked_fill(wei==0, float('-inf'))\n",
    "print(F.softmax(wei, dim = 1))\n",
    "wei2 = F.softmax(wei, dim = 1)\n",
    "print(wei2@x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Self Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 8, 16])\n",
      "torch.Size([4, 8, 16])\n",
      "torch.Size([4, 8, 8])\n",
      "torch.Size([4, 8, 16])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32\n",
    "x = torch.randn((B,T,C))\n",
    "\n",
    "### Attention = softmax(query * key / sqrt(d_k)) * value\n",
    "head_size = 16\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",
    "\n",
    "k = key(x)\n",
    "q = query(x)\n",
    "\n",
    "print(k.shape)\n",
    "print(q.shape)\n",
    "\n",
    "wei = q@k.transpose(2,1)\n",
    "print(wei.shape)\n",
    "wei = wei*head_size**-0.5\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",
    "\n",
    "out = wei@v \n",
    "print(out.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.6402,  0.8511,  0.0207, -1.6763],\n",
      "        [ 1.2977, -1.1287,  0.7366, -1.7447]])\n",
      "tensor([-0.0411, -0.2098])\n",
      "tensor([1.3126, 2.1228])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn((2,4))\n",
    "print(x)\n",
    "# Let's check the mean\n",
    "print(x.mean(dim=1))\n",
    "print(x.var(dim=1))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0161,  0.1886, -0.4992, -0.1440],\n",
      "        [-0.1002,  0.7558, -0.6298, -1.8367]])\n",
      "tensor([[ 0.3516,  1.0602, -1.3206, -0.0912],\n",
      "        [ 0.3249,  1.1139, -0.1632, -1.2756]])\n",
      "tensor([3.7253e-08, 0.0000e+00])\n",
      "tensor([1.0000, 1.0000])\n"
     ]
    }
   ],
   "source": [
    "## Let' try out layer normalization\n",
    "# Layer normalization is a normalization technique that normalizes the activations of a layer for each given example in a mini-batch.\n",
    "# out = (x - mean) / sqrt(var + eps) * gamma + beta\n",
    "# where gamma and beta are learnable parameters\n",
    "\n",
    "x = torch.randn((2,4))\n",
    "gamma = torch.ones(4)\n",
    "beta = torch.zeros(4)\n",
    "eps = 1e-6\n",
    "mean = x.mean(dim=1, keepdim=True)\n",
    "var = x.var(dim=1, keepdim=True)\n",
    "\n",
    "out = (x - mean) / torch.sqrt(var + eps) * gamma + beta\n",
    "\n",
    "print(x)\n",
    "print(out)\n",
    "\n",
    "print(out.mean(dim=1))\n",
    "print(out.var(dim=1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm1d(nn.Module):\n",
    "    def __init__(self, dim, eps = 1e-6) -> None:\n",
    "        super().__init__()\n",
    "        self.gamma = torch.ones(dim)\n",
    "        self.beta = torch.zeros(dim)\n",
    "        self.eps = eps\n",
    "        \n",
    "    def forward(self, x):\n",
    "        mean = x.mean(dim=1, keepdim=True)\n",
    "        var = x.var(dim=1, keepdim=True)\n",
    "\n",
    "        x_hat = (x - mean) / torch.sqrt(var + eps) \n",
    "        out = x_hat* self.gamma + self.beta\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.1335, -0.1059, -0.3824,  ..., -1.3422, -0.1971,  0.8795],\n",
      "        [-0.0353, -0.7440, -0.3371,  ..., -0.6276, -0.4846,  0.4557],\n",
      "        [ 0.3069, -1.5011,  1.4898,  ..., -0.6819,  0.9993,  0.8382],\n",
      "        ...,\n",
      "        [-1.6081, -1.6324, -0.7634,  ..., -0.9847,  0.0039, -0.8610],\n",
      "        [-0.2273,  0.0066, -0.2763,  ..., -0.8705, -1.2443, -0.7531],\n",
      "        [ 0.3054, -0.1505, -0.3809,  ..., -1.4962, -0.7711, -1.0681]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 100])"
      ]
     },
     "execution_count": 223,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "ln = LayerNorm1d(100)\n",
    "x = torch.randn(32, 100) # batch size 32 of 100-dimensional vectors\n",
    "x = ln(x)\n",
    "print(x)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1469), tensor(0.8803))"
      ]
     },
     "execution_count": 224,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-3.5763e-09), tensor(1.0000))"
      ]
     },
     "execution_count": 225,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multihead Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Multihead attention\n",
    "## It consists of multiple attention heads, each of which is a scaled dot-product attention mechanism\n",
    "## Each head is a linear transformation of the input followed by a scaled dot-product attention\n",
    "## The outputs of the attention heads are concatenated and linearly transformed to produce the final output\n",
    "## There is a residual connection around each of the sub-layers, followed by layer normalization\n",
    "## The output of the final multi-head attention layer is passed through a feed-forward neural network, followed by another layer normalization\n",
    "\n",
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, embed_size, head_size, dropout_val = 0.1) -> None:\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(embed_size, head_size, bias=False)\n",
    "        self.query = nn.Linear(embed_size, head_size, bias=False)\n",
    "        self.value = nn.Linear(embed_size, head_size, bias=False)\n",
    "        self.dropout = nn.Dropout(dropout_val)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape \n",
    "        k = self.key(x) # (B, T, H)\n",
    "        q = self.query(x) # (B, T, H)\n",
    "        \n",
    "        wei = q@k.transpose(2,1) # (B, T, T)\n",
    "        wei = wei / head_size**-0.5\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",
    "        wei = self.dropout(wei)\n",
    "        \n",
    "        v = self.value(x) # (B, T, H)\n",
    "        out = wei@v # (B, T, T) @ (B, T, H) -> (B, T, H)\n",
    "        return out\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 4, 16])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn((2,4,32))\n",
    "head_size = 16\n",
    "attn = SelfAttention(32, head_size)\n",
    "out = attn(x)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, num_heads, embed_dim, head_size, attn_dropout_rate = 0.1) -> None:\n",
    "        super().__init__()\n",
    "        self.attn_heads = nn.ModuleList([SelfAttention(embed_size=embed_dim, head_size=head_size) for _ in range(num_heads)])\n",
    "        self.projection = nn.Linear(num_heads*head_size, embed_dim)\n",
    "        self.dropout = nn.Dropout(attn_dropout_rate)\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.shape\n",
    "        out = torch.cat([attn_head(x) for attn_head in self.attn_heads], dim = -1)\n",
    "        out = self.dropout(self.projection(out))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 4, 32])\n"
     ]
    }
   ],
   "source": [
    "x = torch.randn((2,4,32))\n",
    "mh = MultiHeadAttention(8, 32, 16)\n",
    "out = mh(x)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads, attn_dropout_rate = 0.1) -> None:\n",
    "        super().__init__()\n",
    "        head_size = embed_dim // num_heads\n",
    "        self.mha = MultiHeadAttention(num_heads, embed_dim, head_size, attn_dropout_rate)\n",
    "        self.ln1 = nn.LayerNorm(embed_dim)\n",
    "        self.ffn = nn.Sequential(\n",
    "            nn.Linear(embed_dim, 4*embed_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4*embed_dim, embed_dim),\n",
    "            nn.Dropout(attn_dropout_rate)\n",
    "        )\n",
    "        self.ln2 = nn.LayerNorm(embed_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        B, T, C =  x.shape \n",
    "        out = self.mha(x)\n",
    "        # print(out)\n",
    "        out = self.ln1(out + x)\n",
    "        # print(out)\n",
    "        out = self.ffn(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1, 128])\n"
     ]
    }
   ],
   "source": [
    "x = torch.zeros((1, 1, 128))\n",
    "db = nn.Sequential(*[DecoderBlock(128, 8) for _ in range(6)])\n",
    "out = db(x)\n",
    "print(out.shape)\n",
    "# print(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPTScratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 266,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPTScratch(nn.Module):\n",
    "    def __init__(self, vocab_size, embed_dim, num_heads, num_layers, block_size, attn_dropout_rate = 0.1) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        self.token_embedding = nn.Embedding(vocab_size, embed_dim)\n",
    "        self.positional_embedding = nn.Embedding(block_size, embed_dim)\n",
    "        self.decoder_blocks = nn.Sequential(*[DecoderBlock(embed_dim, num_heads, attn_dropout_rate) for _ in range(num_layers)])\n",
    "        self.ln = nn.LayerNorm(embed_dim)\n",
    "        self.lm_head = nn.Linear(embed_dim, vocab_size)\n",
    "        \n",
    "    def forward(self, idx, targets = None):\n",
    "        B, T = idx.shape \n",
    "        tok_emb = self.token_embedding(idx) # (B,T) -> (B,T,C)\n",
    "        pos_emb = self.positional_embedding(torch.arange(T, device=idx.device)) # (T,C)\n",
    "        x = tok_emb + pos_emb\n",
    "        # print(x.shape)\n",
    "        x = self.decoder_blocks(x)\n",
    "        x = self.ln(x)\n",
    "        logits = self.lm_head(x)\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C) # (B,T,C) -> (B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 288,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "block_size = 32\n",
    "embed_dim = 128\n",
    "num_heads = 4\n",
    "num_layers = 8\n",
    "dropout = 0.1\n",
    "learning_rate = 0.01\n",
    "device = 'cpu'\n",
    "max_iter = 1000\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 289,
   "metadata": {},
   "outputs": [],
   "source": [
    "gptModel = GPTScratch(vocab_size, embed_dim, num_heads, num_layers, block_size, dropout).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 290,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.604161 M parameters\n"
     ]
    }
   ],
   "source": [
    "print(sum(p.numel() for p in gptModel.parameters())/1e6, 'M parameters')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 291,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(gptModel.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 292,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter: 0, Loss: 4.45413064956665\n",
      "Iter: 50, Loss: 3.4570345878601074\n",
      "Iter: 100, Loss: 3.2940514087677\n",
      "Iter: 150, Loss: 3.359337329864502\n",
      "Iter: 200, Loss: 3.4028513431549072\n",
      "Iter: 250, Loss: 3.2585270404815674\n",
      "Iter: 300, Loss: 3.381510019302368\n",
      "Iter: 350, Loss: 3.33134388923645\n",
      "Iter: 400, Loss: 3.368924140930176\n",
      "Iter: 450, Loss: 3.2642624378204346\n",
      "Iter: 500, Loss: 3.2841086387634277\n",
      "Iter: 550, Loss: 3.329723358154297\n",
      "Iter: 600, Loss: 3.315089464187622\n",
      "Iter: 650, Loss: 3.2320597171783447\n",
      "Iter: 700, Loss: 3.4392151832580566\n",
      "Iter: 750, Loss: 3.280527114868164\n",
      "Iter: 800, Loss: 3.225339412689209\n",
      "Iter: 850, Loss: 3.333913564682007\n",
      "Iter: 900, Loss: 3.422816038131714\n",
      "Iter: 950, Loss: 3.3597142696380615\n"
     ]
    }
   ],
   "source": [
    "for iter in range(max_iter):\n",
    "    xb, yb = get_data_batch('train')\n",
    "    \n",
    "    logits, loss = gptModel(xb.to(device), yb.to(device))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if iter % 50 == 0:\n",
    "        print(f\"Iter: {iter}, Loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 293,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(model, idx, max_new_tokens):\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_new_tokens):\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            # print(idx_cond)\n",
    "            logits, loss  = model(idx_cond)\n",
    "            logits = logits[:, -1, :] # (B,T,V) -> (B,V)\n",
    "            # print(logits.shape)\n",
    "            # print(logits)\n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            # print(probs)\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)\n",
    "            idx = torch.cat((idx, idx_next), dim=1)\n",
    "    return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 287,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      " sn\n",
      "wKt  tm tLi sfssufnshlaunt,vToO ntrdoo\n",
      "etryvbde lue Ue t i.eaiaaeeC ro mawtreUteskwpef Orroe dten ninoeeEa\n",
      "Ah,dOe\n",
      "aiansm \n",
      "e nocOfuE uUoemastvdegtryM sto'tdyIr nraetrh imeaereaosceewtIm ehwFaf   emrdyyr wooeht aoh uakd' ahococdta oerew nayheafodd i\n",
      "t\n",
      "\n",
      "ne\n",
      "kt\n",
      "jme eurrsiaJthneh ud!s re;  obd  tuIIre!oeL:eaeeaanRs:pne.hsml holohryiertuoay u uerftht ajyl,uW \n",
      "rNp unhulfgstluplssgIes!nTsEtt rBau'\n",
      " hhoToN\n",
      "l\n",
      "w c a e  tscsnt ifog,tcsu ht aiueshegSVao .e;t hwyfnlodeaeAts\n",
      " tt  r e ensne'a   ma  yp at.hm\n",
      "u nth LD ntr ltmo n;eT uo \n",
      "a  Y na\n",
      "d:haenuWt Bf nu hoOe uirt a\n",
      " t: tsMo u TArd A:eol r eBa.psl  iard\n",
      "nhaitn ieaa att,eda r  engn:ttr s  IhBrh\n",
      "slcseaibeDfdtsnis\n",
      ",ht dlwaIhd, pnCatppe rrednne rudlaayml\n",
      "!ehhoLo trrt trDe\n",
      "l, paeeyReIIl ,,t eo. rcde se utaachl m otacdhl daAudg nih s ahyhlf    deai e e:da'nou uhTGb\n",
      "ee Uai e,\n",
      "eudgt y ahielr w tn Ryo'un lp\n",
      "OtlIRahos pg  leeo ie aut ilcsaIttda,tssWrE,rTtdrsue \n",
      "wl?rrmykaUaebcm:l;tgt.oc,ara\n",
      "l\n",
      "pauEhlyaotratig,apugdCdarauh peoou,iookyhpiosO seghhpd\n",
      "\n",
      "tsNatNw l prneesees,lt gW eD\n",
      "RaIhh\n",
      "snrorN  uee  heO Is  ipiudrpmh  yuNtFK:nmnsyohetwff ehp ti,MelbILr.ry yaidnige toeits\n",
      "h nfo Nuenuk\n",
      "hrt td uo duy rslfhtr\n",
      " csCes-?duKoeindreE t tO\n",
      "lomh I  A igjh  ma aauaede shw aedilon adrr oehrc elh\n",
      "uoVAt \n",
      "nt\n",
      "gods,iALeueyhramt .ehs\n",
      "tAnst:Nrwomt eyd efn miFaui n tne n tbiTsaueitNurhwliSeveu\n",
      "vh\n",
      "m ,hoIloeasig veYtiCcin b Osoe eune tS nnn hia orfsist pss it eh\n",
      "n yh\n",
      "ehasnNos \n",
      "lstyuo ppuefs :luc   ,MphTa  c\n",
      "qoeyLtaE hotot s'ea\n",
      "hL Ereehcr  owetlnnasalth t SlWR\n",
      " mariv r .oE nrop lnaPIeoI sramDrhiwdsne lnksts\n",
      "dmoero satoust My\n",
      "qtweC  Itnt\n",
      "lcrfa'hdIyaseLIi a WnyttblM!s,sP en:chisiee d slo otnhramnlscng rsuhaiuuo'  dashetroo:Itf Igsudsesss Rkn seeaO    tgg\n",
      "wedrpeaup UIrf ,ptsiyifta\n",
      "oriMECmhiew?\n",
      "usy U stW mS;ebrels\n",
      "d  h ohkagcecUio Ts\n",
      "rrd AgRt,neaft    rult enegBtapahos  da.t hor i pRtokhitw rghleh:t?u.OwIa\n",
      "ksw d btoe :,i b    e fbedirT 'rrawi\n",
      "gtKhunSioG,\n",
      "msh: fo fDlhd,ehdoee mEsie faesnaaotr kvari\n",
      " oht iaanhihaNaty h \n",
      "geeLapet oUkrn guyf f \n",
      "aelouno\n",
      "l cnshwu ntk naoon d\n"
     ]
    }
   ],
   "source": [
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(context[0].tolist()))\n",
    "\n",
    "print(decode(generate(gptModel, context, max_new_tokens=2000)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}