play.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of characters:  2579948\n",
      "Three Rings for the Elven-kings under the sky,\n",
      "               Seven for the Dwarf-lords in their hal\n"
     ]
    }
   ],
   "source": [
    "with open('data/input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "print(\"number of characters: \", len(text))\n",
    "print(text[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "82"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab = sorted(list(set(text)))\n",
    "len(vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2579849, 2579885, 2579886, 2579888, 2579859, 2579885, 2579886, 2579877, 2579824, 2579889, 2579870, 2579890, 2579889, 2579827, 2579715, 2579886, 2579888, 2579827, 2579886, 2579888, 2579715, 2579889]\n",
      "we are building an agi\n"
     ]
    }
   ],
   "source": [
    "class Tokenizer:\n",
    "    def __init__(self, vocab):\n",
    "        self.stoi = { ch:i for i,ch in enumerate(vocab) }\n",
    "        self.itos = { i:ch for i,ch in enumerate(vocab) }\n",
    "\n",
    "    def encode(self, text):\n",
    "        ids = [self.stoi[c] for c in text]\n",
    "        return ids\n",
    "\n",
    "    def decode(self, ids):\n",
    "        text = ''.join([self.itos[i] for i in ids])\n",
    "        return text\n",
    "\n",
    "tokenizer = Tokenizer(text)\n",
    "text = \"\"\"we are building an agi\"\"\"\n",
    "print(tokenizer.encode(text))\n",
    "print(tokenizer.decode(tokenizer.encode(text)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "just use tiktoken"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "enc = tiktoken.get_encoding(\"o200k_base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[854, 553, 6282, 448, 1017, 72]\n",
      "we are building an agi\n"
     ]
    }
   ],
   "source": [
    "text = \"we are building an agi\"\n",
    "ids = enc.encode(text)\n",
    "print(ids)\n",
    "print(enc.decode(ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "build a dataset class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "\n",
    "    def __init__(self, text, tokenizer, max_length, stride):\n",
    "\n",
    "        self.input_ids = []\n",
    "        self.target_ids = []\n",
    "\n",
    "        token_ids = tokenizer.encode(text)\n",
    "\n",
    "        for i in range(0, len(token_ids) - max_length, stride):\n",
    "\n",
    "            input_chunk = token_ids[i: i+max_length]\n",
    "            target_chunk = token_ids[i+1: i+max_length+1]\n",
    "\n",
    "            self.input_ids.append(torch.tensor(input_chunk))\n",
    "            self.target_ids.append(torch.tensor(target_chunk))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.target_ids[idx]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = CustomDataset(text, enc, max_length=4, stride=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "623287"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "dataloader = DataLoader(dataset=data, \n",
    "                        batch_size=1,\n",
    "                        num_workers=0,\n",
    "                        drop_last=True, # drops last batch if its shorter than specified batch_size\n",
    "                        shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 553, 6282,  448, 1017]]), tensor([[6282,  448, 1017,   72]]))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_iter = iter(dataloader)\n",
    "x, y = next(data_iter)\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 4, 256])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 50257\n",
    "output_dim = 256\n",
    "embedding_layer = torch.nn.Embedding(vocab_size, output_dim)\n",
    "token_embd = embedding_layer(x)\n",
    "\n",
    "max_length=4\n",
    "pos_embedding_layer = torch.nn.Embedding(max_length, output_dim)\n",
    "pos_embd = pos_embedding_layer(torch.arange(max_length)) # 0 1 ... max_length-1\n",
    "\n",
    "input_embd = token_embd + pos_embd\n",
    "input_embd.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "attention mechanism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],\n",
      "        [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],\n",
      "        [0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],\n",
      "        [0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],\n",
      "        [0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],\n",
      "        [0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])\n",
      "tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],\n",
      "        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],\n",
      "        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],\n",
      "        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],\n",
      "        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],\n",
      "        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])\n",
      "tensor([[0.4421, 0.5931, 0.5790],\n",
      "        [0.4419, 0.6515, 0.5683],\n",
      "        [0.4431, 0.6496, 0.5671],\n",
      "        [0.4304, 0.6298, 0.5510],\n",
      "        [0.4671, 0.5910, 0.5266],\n",
      "        [0.4177, 0.6503, 0.5645]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89],\n",
    "   [0.55, 0.87, 0.66],\n",
    "   [0.57, 0.85, 0.64],\n",
    "   [0.22, 0.58, 0.33],\n",
    "   [0.77, 0.25, 0.10],\n",
    "   [0.05, 0.80, 0.55]]\n",
    ")\n",
    "\n",
    "# step 1\n",
    "attn_scr = inputs @ inputs.T # efficient matmul\n",
    "print(attn_scr)\n",
    "\n",
    "# step 2\n",
    "attn_w = torch.softmax(attn_scr, dim=-1)\n",
    "print(attn_w)\n",
    "\n",
    "# step 3\n",
    "context_vec = attn_w @ inputs\n",
    "print(context_vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "attention mechanism with trainable matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0033, -0.0760],\n",
       "        [ 0.0025, -0.0773],\n",
       "        [ 0.0024, -0.0775],\n",
       "        [ 0.0010, -0.0804],\n",
       "        [ 0.0001, -0.0821],\n",
       "        [ 0.0019, -0.0786]], grad_fn=<MmBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89],\n",
    "   [0.55, 0.87, 0.66],\n",
    "   [0.57, 0.85, 0.64],\n",
    "   [0.22, 0.58, 0.33],\n",
    "   [0.77, 0.25, 0.10],\n",
    "   [0.05, 0.80, 0.55]]\n",
    ")\n",
    "\n",
    "class SelfAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out):\n",
    "\n",
    "        super().__init__()\n",
    "        # instead of nn.Parameter(torch.rand(d_in, d_out)) as nn.Linear has \n",
    "        # optimized weight initialization scheme and it is more efficient\n",
    "        # in matmul ops when bias=False\n",
    "        self.w_q = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_k = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_v = nn.Linear(d_in, d_out, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        keys = self.w_k(x)\n",
    "        queries = self.w_q(x)\n",
    "        values = self.w_v(x)\n",
    "\n",
    "        attn_scores = queries @ keys.T\n",
    "\n",
    "        d_k = keys.shape[-1]\n",
    "        attn_weights = torch.softmax(attn_scores / d_k ** 0.5, dim=-1) # original paper d_k=64\n",
    "        \n",
    "        context_vectors = attn_weights @ values\n",
    "\n",
    "        return context_vectors \n",
    "\n",
    "d_in = 3\n",
    "d_out = 2\n",
    "sa = SelfAttention(d_in, d_out)\n",
    "sa(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "causal attention: prevent model from accessing future tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0000,  0.0000],\n",
       "         [-0.2049, -0.1062],\n",
       "         [-0.1351, -0.0700],\n",
       "         [ 0.0655, -0.0176],\n",
       "         [-0.2211, -0.0101],\n",
       "         [ 0.0513, -0.0541]],\n",
       "\n",
       "        [[ 0.5709,  0.1466],\n",
       "         [ 0.2761,  0.0709],\n",
       "         [ 0.0471, -0.0232],\n",
       "         [-0.0386, -0.0638],\n",
       "         [-0.1129, -0.0938],\n",
       "         [-0.0184, -0.0903]]], grad_fn=<UnsafeViewBackward0>)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89],\n",
    "   [0.55, 0.87, 0.66],\n",
    "   [0.57, 0.85, 0.64],\n",
    "   [0.22, 0.58, 0.33],\n",
    "   [0.77, 0.25, 0.10],\n",
    "   [0.05, 0.80, 0.55]]\n",
    ")\n",
    "\n",
    "class CasualAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, bias=False):\n",
    "\n",
    "        super().__init__()\n",
    "        # instead of nn.Parameter(torch.rand(d_in, d_out)) as nn.Linear has \n",
    "        # optimized weight initialization scheme and it is more efficient\n",
    "        # in matmul ops when bias=False\n",
    "        self.w_q = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_k = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_v = nn.Linear(d_in, d_out, bias=False)\n",
    "\n",
    "        self.d_out = d_out\n",
    "        self.dropout = nn.Dropout(0.5) # prevent overfitting\n",
    "\n",
    "        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # mask to prevent model from seeing future tokens\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        b, tokens, d_in = x.shape\n",
    "        \n",
    "        keys = self.w_k(x)\n",
    "        queries = self.w_q(x)\n",
    "        values = self.w_v(x)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(1, 2)\n",
    "\n",
    "        masked = attn_scores.masked_fill(self.mask.bool()[:tokens, :tokens], -torch.inf) # masking with -inf\n",
    "\n",
    "        d_k = keys.shape[-1]\n",
    "        attn_weights = torch.softmax(masked / d_k ** 0.5, dim=-1) # original paper d_k=64\n",
    "\n",
    "        # drop by 50% and remaining are rescale by a factor of 1/0.5 -> 2\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vectors = attn_weights @ values\n",
    "\n",
    "        return context_vectors\n",
    "\n",
    "d_out = 2\n",
    "batch = torch.stack((inputs, inputs), dim=0)\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "ca = CasualAttention(d_in, d_out, context_length)\n",
    "ca(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "multi-head attention: split attention into multiple heads where each head learns different aspect of the data and then combine the outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [-0.9600, -0.6369, -0.2028,  0.6374],\n",
       "         [-1.7226, -0.9330,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000, -0.2691,  0.9411],\n",
       "         [-0.9658, -0.6532, -0.2174,  0.7612],\n",
       "         [-0.3829, -0.1564, -0.1982,  0.5920]],\n",
       "\n",
       "        [[ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [-0.9600, -0.6369, -0.3205,  1.2088],\n",
       "         [-1.7226, -0.9330, -0.2739,  0.8554],\n",
       "         [-0.7407, -0.5015, -0.1639,  0.6102],\n",
       "         [-0.2062, -0.2165, -0.0476,  0.2311],\n",
       "         [-0.6981, -0.5547, -0.2338,  0.7444]],\n",
       "\n",
       "        [[ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000, -0.3205,  1.2088],\n",
       "         [-0.6381, -0.4233, -0.1374,  0.4262],\n",
       "         [-0.4829, -0.3204, -0.1597,  0.4945],\n",
       "         [-1.1720, -0.8698, -0.2396,  0.7670],\n",
       "         [-0.3193, -0.2118, -0.1086,  0.4078]]], grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89],\n",
    "   [0.55, 0.87, 0.66],\n",
    "   [0.57, 0.85, 0.64],\n",
    "   [0.22, 0.58, 0.33],\n",
    "   [0.77, 0.25, 0.10],\n",
    "   [0.05, 0.80, 0.55]]\n",
    ")\n",
    "\n",
    "class MultiHeadAttentionStack(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, n_head, bias=False):\n",
    "\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([CasualAttention(d_in, d_out, context_length) for _ in range(n_head)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.cat([head(x) for head in self.heads], dim=-1)\n",
    "\n",
    "d_in = 3\n",
    "d_out = 2\n",
    "batch = torch.stack((inputs, inputs, inputs), dim=0)\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "mhas = MultiHeadAttentionStack(d_in, d_out, context_length, n_head=2)\n",
    "mhas(batch)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "lets combine causal attention and multi-head attention code to compute the attention in parallel. currently we're just stacking multiple causal attention blocks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.6093, -0.2380],\n",
       "         [ 0.6680, -0.1577],\n",
       "         [ 0.6637, -0.3462],\n",
       "         [ 0.4764, -0.4829],\n",
       "         [ 0.4008, -0.6265],\n",
       "         [ 0.4819, -0.5453]],\n",
       "\n",
       "        [[ 0.6528, -0.4079],\n",
       "         [ 0.5989, -0.3880],\n",
       "         [ 0.6899, -0.4989],\n",
       "         [ 0.6011, -0.6095],\n",
       "         [ 0.5705, -0.3075],\n",
       "         [ 0.5086, -0.5835]],\n",
       "\n",
       "        [[ 0.5433, -0.2040],\n",
       "         [ 0.5737, -0.3353],\n",
       "         [ 0.5032, -0.5368],\n",
       "         [ 0.6032, -0.2582],\n",
       "         [ 0.5619, -0.4527],\n",
       "         [ 0.6547, -0.3854]]], grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "inputs = torch.tensor(\n",
    "  [[0.43, 0.15, 0.89],\n",
    "   [0.55, 0.87, 0.66],\n",
    "   [0.57, 0.85, 0.64],\n",
    "   [0.22, 0.58, 0.33],\n",
    "   [0.77, 0.25, 0.10],\n",
    "   [0.05, 0.80, 0.55]]\n",
    ")\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_in, d_out, context_length, n_heads, bias=False):\n",
    "\n",
    "        super().__init__()\n",
    "        \n",
    "        # d_out must be divisible by n_heads to distribute it perfectly across heads\n",
    "        self.d_out = d_out\n",
    "        self.n_heads = n_heads\n",
    "        self.head_dim = d_out // n_heads\n",
    "\n",
    "        self.w_q = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_k = nn.Linear(d_in, d_out, bias=False)\n",
    "        self.w_v = nn.Linear(d_in, d_out, bias=False)\n",
    "\n",
    "        self.out_proj = nn.Linear(d_out, d_out)\n",
    "        \n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        b, n_tokens, d_in = x.shape\n",
    "\n",
    "        keys = self.w_k(x)\n",
    "        queries = self.w_q(x)\n",
    "        values = self.w_v(x)\n",
    "\n",
    "        # split d_out into n_heads and head_dim\n",
    "        keys = keys.view(b, n_tokens, self.n_heads, self.head_dim)\n",
    "        queries = queries.view(b, n_tokens, self.n_heads, self.head_dim)\n",
    "        values = values.view(b, n_tokens, self.n_heads, self.head_dim)\n",
    "\n",
    "        # align all the three matrices\n",
    "        keys = keys.transpose(1, 2)\n",
    "        queries = queries.transpose(1, 2)\n",
    "        values = values.transpose(1, 2)\n",
    "\n",
    "        attn_scores = queries @ keys.transpose(2, 3)\n",
    "\n",
    "        # create mask of shape (6, 6) with true and false\n",
    "        mask = self.mask.bool()[:n_tokens, :n_tokens]\n",
    "        \n",
    "        # use mask to fill the attention scores matrix \n",
    "        attn_scores.masked_fill(mask, -torch.inf)\n",
    "        \n",
    "        d_k = keys.shape[-1]\n",
    "        attn_weights = torch.softmax(attn_scores / d_k ** 0.5, dim=-1)\n",
    "\n",
    "        attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        context_vec = (attn_weights @ values).transpose(1, 2) # transpose back to the shape (b, n_tokens, n_heads, head_dim)\n",
    "\n",
    "        context_vec = context_vec.contiguous().view(b, n_tokens, self.d_out)\n",
    "\n",
    "        context_vec = self.out_proj(context_vec) # linear projection\n",
    "\n",
    "        return context_vec\n",
    "\n",
    "d_out = 2\n",
    "batch = torch.stack((inputs, inputs, inputs), dim=0)\n",
    "batch_size, context_length, d_in = batch.shape\n",
    "mha = MultiHeadAttention(d_in, d_out, context_length, n_heads=2)\n",
    "mha(batch)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}