tutorialGPT/tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ddf475dd",
   "metadata": {},
   "source": [
    "## ShakespeareGPT\n",
    "\n",
    "> based on [Let's build GPT: from scratch, in code, spelled out.](https://www.youtube.com/watch?v=kCc8FmEb1nY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "6c863d65",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "48c6be17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fbec4214050>"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1357)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1ee13a9",
   "metadata": {},
   "source": [
    "# Preparing Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "cafad384",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len(data)=1114985\n",
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "with open('./dataset/shakespeare.txt','r',encoding='utf-8') as f:\n",
    "    data = f.read()\n",
    "    \n",
    "print(f\"{len(data)=}\\n{data[:100]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32a2b9b",
   "metadata": {},
   "source": [
    "## Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "ca8d4ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CharacterLevelTokenizer:\n",
    "    def __init__(self,data):\n",
    "        self.data = data\n",
    "        self.vocab = sorted(list(set(self.data)))\n",
    "        self.VOCAB_SIZE = len(self.vocab)\n",
    "        \n",
    "        self.i_s = {i:s for i,s in enumerate(self.vocab)}\n",
    "        self.s_i = {s:i for i,s in self.i_s.items()}\n",
    "        \n",
    "    def encode(self,s):\n",
    "        return torch.tensor([self.s_i[c] for c in s],dtype=torch.long)\n",
    "\n",
    "    def decode(self,s):\n",
    "        return ''.join([self.i_s[i.item()] for i in s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "c04a4636",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['\\n', ' ', '!', \"'\", ',', '-', '.', ':', ';', '?', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']\n",
      "62\n"
     ]
    }
   ],
   "source": [
    "tokenizer = CharacterLevelTokenizer(data)\n",
    "print(tokenizer.vocab)\n",
    "print(tokenizer.VOCAB_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "755f7dfa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([40, 55,  1, 55, 56,  1, 37, 53, 56, 55, 40,  9])"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.encode('et tu brute?')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "73131193",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'et tu brute?'"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(tokenizer.encode('et tu brute?'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57b92854",
   "metadata": {},
   "source": [
    "## Config!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "id": "486ec805",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Config:\n",
    "    block_size = 8 # context-length\n",
    "    batch_size = 4 # mini-batch size\n",
    "    vocab_size = tokenizer.VOCAB_SIZE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "017291be",
   "metadata": {},
   "source": [
    "## Dataset & Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "919282cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShakespeareDataset:\n",
    "    def __init__(self,block_size:int, is_test=False) -> None:\n",
    "        self.tokenizer = CharacterLevelTokenizer(data)\n",
    "        self.is_test = is_test\n",
    "        self.full_data = self.tokenizer.encode(self.tokenizer.data)\n",
    "        if self.is_test:\n",
    "            self.data = self.full_data[int(0.9*len(self.full_data)):]\n",
    "        else:\n",
    "            self.data = self.full_data[:int(0.9*len(self.full_data))]\n",
    "        self.block_size = block_size\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.data)\n",
    "\n",
    "    def get_block_size(self) -> int:\n",
    "        return self.block_size\n",
    "\n",
    "    def get_vocab_size(self) -> int:\n",
    "        return self.tokenizer.VOCAB_SIZE\n",
    "\n",
    "    def __getitem__(self,idx):\n",
    "        item = self.data[idx:idx+self.block_size+1]\n",
    "        x = item[:-1]\n",
    "        y = item[1:]\n",
    "        return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "71f769c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_ds.get_block_size()=8\n",
      "train_ds.get_vocab_size()=62\n",
      "len(train_ds)=1003486\n",
      "len(val_ds)=111499\n"
     ]
    }
   ],
   "source": [
    "train_ds = ShakespeareDataset(Config.block_size)\n",
    "print(f'{train_ds.get_block_size()=}\\n{train_ds.get_vocab_size()=}\\n{len(train_ds)=}')\n",
    "\n",
    "val_ds = ShakespeareDataset(Config.block_size,is_test=True)\n",
    "print(f'{len(val_ds)=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "495c7cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = torch.utils.data.DataLoader(train_ds,shuffle=False,batch_size=Config.batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "6db66826",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 8]) torch.Size([4, 8])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[15, 44, 53, 54, 55,  1, 12, 44],\n",
       "         [44, 53, 54, 55,  1, 12, 44, 55],\n",
       "         [53, 54, 55,  1, 12, 44, 55, 44],\n",
       "         [54, 55,  1, 12, 44, 55, 44, 61]]),\n",
       " tensor([[44, 53, 54, 55,  1, 12, 44, 55],\n",
       "         [53, 54, 55,  1, 12, 44, 55, 44],\n",
       "         [54, 55,  1, 12, 44, 55, 44, 61],\n",
       "         [55,  1, 12, 44, 55, 44, 61, 40]]))"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs,targets=next(iter(train_dl))\n",
    "print(inputs.shape,targets.shape)\n",
    "inputs,targets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8059a12",
   "metadata": {},
   "source": [
    "# Bi-gram Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "ded27d45",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLM(nn.Module):\n",
    "    def __init__(self,vocab_size):\n",
    "        super(BigramLM,self).__init__()\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size,vocab_size)\n",
    "        \n",
    "    def forward(self,idx,targets=None):\n",
    "        logits = self.token_embedding_table(idx) # (B,T,C:vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # torch cross entropy expects B,C,T instead of B,T,C\n",
    "            # and for targets, we need B*T instead of B,T\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",
    "        \n",
    "    def generate(self,idx,total):\n",
    "        # idx (B,T) in current context\n",
    "        for _ in range(total):\n",
    "            logits,loss = self(idx)\n",
    "            # since the last element is the next character, we pluck out -1 from T\n",
    "            logits = logits[:,-1,:] # (B*T,C) -> (B,C)\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) # (B, T+=1)\n",
    "            \n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "cb329f59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([256, 62]) tensor(4.4188, grad_fn=<NllLossBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"\\nxqI E:Mp.HtfZDyhnSa!uQaSjIncCgX'xwUv-P;DzahqW.RY;ldx CmYAQgCT.noI\\nXqtX JCeZzbMPkGcbMxQAt;l--ZFlHmUSc\""
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bglm = BigramLM(tokenizer.VOCAB_SIZE)\n",
    "logits,loss = bglm(inputs,targets)\n",
    "print(logits.shape,loss)\n",
    "\n",
    "generated = bglm.generate(\n",
    "    torch.zeros((1,1),dtype=torch.long), # initial context 0\n",
    "    total=100\n",
    ")\n",
    "generated = tokenizer.decode(generated[0])\n",
    "generated"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ced2293c",
   "metadata": {},
   "source": [
    "## training the bigram LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "93b4de06",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 0 loss: 4.459361553192139\n",
      "step: 2500 loss: 3.365318536758423\n",
      "step: 5000 loss: 3.2022905349731445\n",
      "step: 7500 loss: 2.7361788749694824\n",
      "step: 10000 loss: 2.567321538925171\n",
      "step: 12500 loss: 2.655674457550049\n",
      "step: 15000 loss: 2.230267286300659\n",
      "step: 17500 loss: 2.302700996398926\n",
      "step: 20000 loss: 2.6413631439208984\n",
      "step: 22500 loss: 2.369800329208374\n"
     ]
    }
   ],
   "source": [
    "bglm = BigramLM(tokenizer.VOCAB_SIZE)\n",
    "\n",
    "optim = torch.optim.AdamW(bglm.parameters(),lr=1e-3)\n",
    "bglm_dl = torch.utils.data.DataLoader(train_ds,shuffle=False,batch_size=32)\n",
    "\n",
    "it = iter(bglm_dl)\n",
    "for steps in range(25_000):\n",
    "    inputs,targets = next(it)\n",
    "    logits,loss=bglm(inputs,targets)\n",
    "    optim.zero_grad()\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    if steps%2500==0:\n",
    "        print(f'step: {steps} loss: {loss.item()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "230d1bff",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Thel sofelie ly rouey warbl.\n",
      "Tough!\n",
      "Whe.\n",
      "Yow. m;\n",
      "DWe f of poote I tigicowe\n",
      "\n",
      "Theld,\n",
      "Prt bre'sinil my:\n",
      "MAn'lyoombur medintacot, he angsss\n",
      "Toy be?\n",
      "CEvese k, h ne thenesee se thtere ngsoupyoree akimy's t geallin tupreiespul o h weakllf ld peais LI winee;\n",
      "VXF ske, dse wavee nsth wersscor g bomalosee at: te I Rothowis t mend n cho, m an cat f o hisemisakelfl gen winer f.\n",
      "Youpand ty bait:\n",
      "Bin I\n",
      "AREdo t, whug.\n",
      "ARYou ghiry, w as s l p.\n",
      "Ter it!\n",
      "Ed hy Heenorivearshair'sthe The beself OLEENUCanthin it ayo m\n"
     ]
    }
   ],
   "source": [
    "generated = bglm.generate(\n",
    "    torch.zeros((1,1),dtype=torch.long), # initial context 0\n",
    "    total=500\n",
    ")\n",
    "generated = tokenizer.decode(generated[0])\n",
    "\n",
    "print(generated)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd989fa0",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# basic communication between tokens!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8aa82b2",
   "metadata": {},
   "source": [
    "### Toy Example\n",
    "\n",
    "we want the tokens along T to \"talk\" to each other\n",
    "AND we also want the tokens to NOT talk to tokens after them, i.e. the future tokens\n",
    "\n",
    "for now let's talk using cumulative average\n",
    "for every Tth token, calculate cumulative average upto that token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "14a87706",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 2])"
      ]
     },
     "execution_count": 154,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B,T,C = 4,8,2 # batch, time, channel\n",
    "x=torch.randn(B,T,C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "7ce3e9fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "xbow = torch.zeros((B,T,C))\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1] # (t,C)\n",
    "        xbow[b,t] = torch.mean(xprev,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "606a99e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.6892,  0.8805],\n",
       "         [-0.3160, -0.8412],\n",
       "         [-0.9974, -0.3895],\n",
       "         [-0.5201,  0.0344],\n",
       "         [-0.1666, -0.6107],\n",
       "         [ 1.4334, -0.0633],\n",
       "         [-0.2296, -0.3650],\n",
       "         [-1.4887, -0.2825]]),\n",
       " tensor([[ 0.6892,  0.8805],\n",
       "         [ 0.1866,  0.0197],\n",
       "         [-0.2080, -0.1167],\n",
       "         [-0.2860, -0.0790],\n",
       "         [-0.2622, -0.1853],\n",
       "         [ 0.0204, -0.1650],\n",
       "         [-0.0153, -0.1935],\n",
       "         [-0.1995, -0.2047]]))"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# notice how each row is an average of all the previous rows\n",
    "x[0],xbow[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "39f72130",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0.],\n",
       "        [1., 1., 0.],\n",
       "        [1., 1., 1.]])"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# tril\n",
    "torch.tril(torch.ones(3,3)) # lower triangular matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "241cdbd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a\n",
      "tensor([[1.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333]])\n",
      "b\n",
      "tensor([[0., 1.],\n",
      "        [3., 0.],\n",
      "        [8., 5.]])\n",
      "c\n",
      "tensor([[0.0000, 1.0000],\n",
      "        [1.5000, 0.5000],\n",
      "        [3.6667, 2.0000]])\n"
     ]
    }
   ],
   "source": [
    "# all the same using matrix multiplication!\n",
    "a = torch.tril(torch.ones(3,3))\n",
    "a = a / torch.sum(a,dim=1,keepdim=True)\n",
    "b = torch.randint(0,10,(3,2)).float()\n",
    "c = a@b\n",
    "print(f'a\\n{a}\\nb\\n{b}\\nc\\n{c}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "6580a93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = torch.tril(torch.ones(T,T))\n",
    "weights = weights/weights.sum(1,keepdim=True)\n",
    "xbow2 = weights @ x # (T, T) @ (B,T,C) =[add batch dim]=> (B, T, T) @ (B, T, C) = (B,T,C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "113e6538",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 157,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(xbow,xbow2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "8cd66592",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.6425, -2.0431],\n",
       "         [ 0.3159, -0.9583],\n",
       "         [ 0.4632, -0.3893],\n",
       "         [ 0.2778, -0.0274],\n",
       "         [ 0.0654, -0.4543],\n",
       "         [ 0.1753, -0.3378],\n",
       "         [ 0.2875, -0.2855],\n",
       "         [ 0.1904, -0.2079]]),\n",
       " tensor([[-0.6425, -2.0431],\n",
       "         [ 0.3159, -0.9583],\n",
       "         [ 0.4632, -0.3893],\n",
       "         [ 0.2778, -0.0274],\n",
       "         [ 0.0654, -0.4543],\n",
       "         [ 0.1753, -0.3378],\n",
       "         [ 0.2875, -0.2855],\n",
       "         [ 0.1904, -0.2079]]))"
      ]
     },
     "execution_count": 158,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xbow[0], xbow2[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54e5e16",
   "metadata": {},
   "source": [
    "### 3rd Version using SOFTMAX!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "bc65718b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 0., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 0., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 1., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tril = torch.tril(torch.ones(T,T))\n",
    "tril"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "279268d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., -inf, -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., -inf, -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., -inf, -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., 0., -inf, -inf],\n",
      "        [0., 0., 0., 0., 0., 0., 0., -inf],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0.]])\n"
     ]
    },
    {
     "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": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights = torch.zeros((T,T))\n",
    "weights = weights.masked_fill(tril==0,float('-inf')) # fill with -infinity where tril is 0\n",
    "print(weights)\n",
    "\"\"\"\n",
    "softmax converts -inf to 0 and rest is avg.\n",
    "since e^-inf = 0\n",
    "and rest 1s are averaged\n",
    "\"\"\"\n",
    "weights = F.softmax(weights,dim=-1)\n",
    "xbow3 = weights @ x\n",
    "weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "fb411e8f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 167,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(xbow,xbow3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d857cd77",
   "metadata": {},
   "source": [
    "the averages are gonna change as they're weights\n",
    "\n",
    "weighted aggregation of past elements give you affinities between tokens depending on how interesting they're to each other ALL using lower triangular matrix multiplication where every value is aggregation of all its previous values.\n",
    "\n",
    "Hiding the future tokens is known as masking or attention mask used in the decoder block of the transformer architecture -- the masking is done using tril as we saw above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24f5f54",
   "metadata": {},
   "source": [
    "# Bigram LM with token and position embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "id": "8274781e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Config:\n",
    "    block_size = 8 # context-length\n",
    "    batch_size = 32 # mini-batch size\n",
    "    vocab_size = tokenizer.VOCAB_SIZE\n",
    "    n_embed = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "f8f4ffe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLM(nn.Module):\n",
    "    def __init__(self,Config):\n",
    "        super(BigramLM,self).__init__()\n",
    "        \n",
    "        self.n_embed = Config.n_embed # number of embedding dims\n",
    "        self.block_size = Config.block_size\n",
    "        \n",
    "        self.token_embedding_table = nn.Embedding(Config.vocab_size,self.n_embed)\n",
    "        \n",
    "        self.pos_embedding_table = nn.Embedding(self.block_size, self.n_embed)\n",
    "        \n",
    "        self.lm_head = nn.Linear(self.n_embed,Config.vocab_size)\n",
    "        \n",
    "    def forward(self,idx,targets=None):\n",
    "        \n",
    "        B,T = idx.shape\n",
    "        \n",
    "        token_embs = self.token_embedding_table(idx) # (B,T,n_embed)\n",
    "        pos_embs = self.pos_embedding_table(torch.arange(T)) # (T,n_embed)\n",
    "        \n",
    "        \"\"\"\n",
    "        token_embs: B,T,n_embed\n",
    "        pos_embs:  ,T,n_embed\n",
    "               +: B,T,n_embed (broadcasted)\n",
    "               \n",
    "        so at this point, x knows the token affinities and importance of position!\n",
    "        [note: here since its a bigram model, position embeddings makes 0 sense]\n",
    "        \"\"\"\n",
    "        x = token_embs + pos_embs # (B,T,n_embed)\n",
    "        \n",
    "        logits = self.lm_head(x) # (B,T,vocab_size)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # torch cross entropy expects B,C,T instead of B,T,C\n",
    "            # and for targets, we need B*T instead of B,T\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",
    "        \n",
    "    def generate(self,idx,total):\n",
    "        # idx (B,T) in current context\n",
    "        for _ in range(total):\n",
    "            logits,loss = self(idx)\n",
    "            # since the last element is the next character, we pluck out -1 from T\n",
    "            logits = logits[:,-1,:] # (B*T,C) -> (B,C)\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) # (B, T+=1)\n",
    "            \n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "id": "4839d4be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 0 loss: 4.577418804168701\n",
      "step: 1000 loss: 2.4386796951293945\n",
      "step: 2000 loss: 2.6219401359558105\n",
      "step: 3000 loss: 2.6456706523895264\n",
      "step: 4000 loss: 2.2980926036834717\n",
      "step: 5000 loss: 3.0280601978302\n",
      "step: 6000 loss: 2.297921895980835\n",
      "step: 7000 loss: 2.7295444011688232\n",
      "step: 8000 loss: 2.1326377391815186\n",
      "step: 9000 loss: 2.271216869354248\n"
     ]
    }
   ],
   "source": [
    "bglm = BigramLM(Config)\n",
    "\n",
    "optim = torch.optim.AdamW(bglm.parameters(),lr=1e-3)\n",
    "bglm_dl = torch.utils.data.DataLoader(train_ds,shuffle=False,batch_size=32)\n",
    "\n",
    "it = iter(bglm_dl)\n",
    "for steps in range(10000):\n",
    "    inputs,targets = next(it)\n",
    "    logits,loss=bglm(inputs,targets)\n",
    "    optim.zero_grad()\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    if steps%1000==0:\n",
    "        print(f'step: {steps} loss: {loss.item()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa6ddb23",
   "metadata": {},
   "source": [
    "# Self Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "id": "5ba23b2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 32])"
      ]
     },
     "execution_count": 197,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B,T,C = 4,8,32\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "tril = torch.tril(torch.ones(T,T))\n",
    "wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril==0,float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "out = wei@x\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f1c260b",
   "metadata": {},
   "source": [
    "gather info from past in data dependent way, the way it does it:\n",
    "\n",
    "every single token emits 2 vectors: query, key\n",
    "\n",
    "query: what am i looking for?\n",
    "\n",
    "key: what do i contain?\n",
    "\n",
    "dot product: query @ key which will be our weights\n",
    "\n",
    "now the final value will be `value`\n",
    "\n",
    "value: private info for each token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "id": "492240ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 208,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# single self-attention head\n",
    "\n",
    "head_size = 16\n",
    "\n",
    "key = nn.Linear(C,head_size,bias=False)\n",
    "query = nn.Linear(C,head_size,bias=False)\n",
    "\n",
    "value = nn.Linear(C,head_size,bias=False)\n",
    "\n",
    "k = key(x) #(B,T,16)\n",
    "q = query(x) #(B,T,16)\n",
    "\n",
    "\"\"\"\n",
    "k: (B,T,16)\n",
    "q: (B,T,16)\n",
    "\n",
    "last two dims of k have to be swapped\n",
    "k: (B,16,T)\n",
    "\n",
    "therefore\n",
    "\n",
    "q: (B,T,16)\n",
    "    @\n",
    "k: (B,16,T)\n",
    "    =\n",
    "   (B,T,T)\n",
    "   \n",
    "which will be our new weights\n",
    "\n",
    "notice how each batch has its own set of weights since each batch is different\n",
    "\n",
    "hence we get data dependency\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "wei = q @ k.transpose(-2,-1)\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",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "id": "d103bb6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4974, 0.5026, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5662, 0.3141, 0.1198, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.4006, 0.0128, 0.4399, 0.1467, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1022, 0.2856, 0.4534, 0.0862, 0.0726, 0.0000, 0.0000, 0.0000],\n",
       "        [0.3116, 0.2433, 0.2295, 0.1141, 0.0351, 0.0665, 0.0000, 0.0000],\n",
       "        [0.0322, 0.0218, 0.0068, 0.1083, 0.3796, 0.4106, 0.0406, 0.0000],\n",
       "        [0.0444, 0.0112, 0.0268, 0.1021, 0.4375, 0.3219, 0.0339, 0.0222]],\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 204,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91cae851",
   "metadata": {},
   "source": [
    "### notes on attention\n",
    "\n",
    "- communication mechanism\n",
    "- like a directed graph, with weights for each edge\n",
    "- graph: 8 nodes\n",
    "    - first node is directed to itself\n",
    "    - second node is directed to itself and the first node\n",
    "    \n",
    "    ...\n",
    "    - the last node is directed to itself and all the previous nodes\n",
    "- attention acts over a set vectors of graph, no notion of node position\n",
    "- hence position embedding is important so the nodes are aware where they are in time\n",
    "- elements along batch dimension dont talk to each other\n",
    "- batch multiplication is only for parallel processing, each sample can be considered its own graph independent of all other graphs\n",
    "\n",
    "transformer:\n",
    "\n",
    "- `encoder` block: allows all tokens to communicate, no masking\n",
    "- `decoder` block: no future token communication, via masking\n",
    "\n",
    "difference between self/cross attention:\n",
    "- self: k,q,v are from same x\n",
    "- cross: q from one x and k,v from other x (like translation task for example)\n",
    "\n",
    "\n",
    "- scaled dot-product attention\n",
    "    - to keep variance ~= 1\n",
    "    - if weights have large variance, then softmax on it makes it more like one-hot vectors\n",
    "    - we need good smooth affinities after softmax\n",
    "\n",
    "```\n",
    "attention = softmax((q@k.T)/sqrt(head_size)) @ v\n",
    "\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}