playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read to inspect the dataset\n",
    "with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters:  1115394\n"
     ]
    }
   ],
   "source": [
    "print(\"length of dataset in characters: \", len(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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",
      "All:\n",
      "We know't, we know't.\n",
      "\n",
      "First Citizen:\n",
      "Let us kill him, and we'll have corn at our own price.\n",
      "Is't a verdict?\n",
      "\n",
      "All:\n",
      "No more talking on't; let it be done: away, away!\n",
      "\n",
      "Second Citizen:\n",
      "One word, good citizens.\n",
      "\n",
      "First Citizen:\n",
      "We are accounted poor citizens, the patricians good.\n",
      "What authority surfeits on would relieve us: if they\n",
      "would yield us but the superfluity, while it were\n",
      "wholesome, we might guess they relieved us humanely;\n",
      "but they think we are too dear: the leanness that\n",
      "afflicts us, the object of our misery, is as an\n",
      "inventory to particularise their abundance; our\n",
      "sufferance is a gain to them Let us revenge this with\n",
      "our pikes, ere we become rakes: for the gods know I\n",
      "speak this in hunger for bread, not in thirst for revenge.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "65\n"
     ]
    }
   ],
   "source": [
    "# unique characters that occur in the dataset\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(\"\".join(chars))\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
      "hii there\n"
     ]
    }
   ],
   "source": [
    "# creating mapping from characters to integers\n",
    "# we are building character level model, we will work on individual characters\n",
    "\n",
    "stoi = {ch: i for i, ch in enumerate(chars)}\n",
    "itos = {i: ch for i, ch in enumerate(chars)}\n",
    "\n",
    "# tokenizer\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: takes a string, output a list of integers\n",
    "decode = lambda l: \"\".join([itos[i] for i in l]) # decoder: takes a list of integers, outputs a string\n",
    "\n",
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115394]) torch.int64\n",
      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
      "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
      "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
      "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
      "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
      "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50, 50,\n",
      "         1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1, 58,\n",
      "        53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51, 47,\n",
      "        57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43, 42,\n",
      "         8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57, 58,\n",
      "         1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1, 63,\n",
      "        53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56, 41,\n",
      "        47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51, 63,\n",
      "         1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0, 13,\n",
      "        50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,  1,\n",
      "        49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57,  1, 49, 47, 50, 50,  1,\n",
      "        46, 47, 51,  6,  1, 39, 52, 42,  1, 61, 43,  5, 50, 50,  1, 46, 39, 60,\n",
      "        43,  1, 41, 53, 56, 52,  1, 39, 58,  1, 53, 59, 56,  1, 53, 61, 52,  1,\n",
      "        54, 56, 47, 41, 43,  8,  0, 21, 57,  5, 58,  1, 39,  1, 60, 43, 56, 42,\n",
      "        47, 41, 58, 12,  0,  0, 13, 50, 50, 10,  0, 26, 53,  1, 51, 53, 56, 43,\n",
      "         1, 58, 39, 50, 49, 47, 52, 45,  1, 53, 52,  5, 58, 11,  1, 50, 43, 58,\n",
      "         1, 47, 58,  1, 40, 43,  1, 42, 53, 52, 43, 10,  1, 39, 61, 39, 63,  6,\n",
      "         1, 39, 61, 39, 63,  2,  0,  0, 31, 43, 41, 53, 52, 42,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 27, 52, 43,  1, 61, 53, 56, 42,  6,  1, 45, 53,\n",
      "        53, 42,  1, 41, 47, 58, 47, 64, 43, 52, 57,  8,  0,  0, 18, 47, 56, 57,\n",
      "        58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 35, 43,  1, 39, 56, 43,  1,\n",
      "        39, 41, 41, 53, 59, 52, 58, 43, 42,  1, 54, 53, 53, 56,  1, 41, 47, 58,\n",
      "        47, 64, 43, 52, 57,  6,  1, 58, 46, 43,  1, 54, 39, 58, 56, 47, 41, 47,\n",
      "        39, 52, 57,  1, 45, 53, 53, 42,  8,  0, 35, 46, 39, 58,  1, 39, 59, 58,\n",
      "        46, 53, 56, 47, 58, 63,  1, 57, 59, 56, 44, 43, 47, 58, 57,  1, 53, 52,\n",
      "         1, 61, 53, 59, 50, 42,  1, 56, 43, 50, 47, 43, 60, 43,  1, 59, 57, 10,\n",
      "         1, 47, 44,  1, 58, 46, 43, 63,  0, 61, 53, 59, 50, 42,  1, 63, 47, 43,\n",
      "        50, 42,  1, 59, 57,  1, 40, 59, 58,  1, 58, 46, 43,  1, 57, 59, 54, 43,\n",
      "        56, 44, 50, 59, 47, 58, 63,  6,  1, 61, 46, 47, 50, 43,  1, 47, 58,  1,\n",
      "        61, 43, 56, 43,  0, 61, 46, 53, 50, 43, 57, 53, 51, 43,  6,  1, 61, 43,\n",
      "         1, 51, 47, 45, 46, 58,  1, 45, 59, 43, 57, 57,  1, 58, 46, 43, 63,  1,\n",
      "        56, 43, 50, 47, 43, 60, 43, 42,  1, 59, 57,  1, 46, 59, 51, 39, 52, 43,\n",
      "        50, 63, 11,  0, 40, 59, 58,  1, 58, 46, 43, 63,  1, 58, 46, 47, 52, 49,\n",
      "         1, 61, 43,  1, 39, 56, 43,  1, 58, 53, 53,  1, 42, 43, 39, 56, 10,  1,\n",
      "        58, 46, 43,  1, 50, 43, 39, 52, 52, 43, 57, 57,  1, 58, 46, 39, 58,  0,\n",
      "        39, 44, 44, 50, 47, 41, 58, 57,  1, 59, 57,  6,  1, 58, 46, 43,  1, 53,\n",
      "        40, 48, 43, 41, 58,  1, 53, 44,  1, 53, 59, 56,  1, 51, 47, 57, 43, 56,\n",
      "        63,  6,  1, 47, 57,  1, 39, 57,  1, 39, 52,  0, 47, 52, 60, 43, 52, 58,\n",
      "        53, 56, 63,  1, 58, 53,  1, 54, 39, 56, 58, 47, 41, 59, 50, 39, 56, 47,\n",
      "        57, 43,  1, 58, 46, 43, 47, 56,  1, 39, 40, 59, 52, 42, 39, 52, 41, 43,\n",
      "        11,  1, 53, 59, 56,  0, 57, 59, 44, 44, 43, 56, 39, 52, 41, 43,  1, 47,\n",
      "        57,  1, 39,  1, 45, 39, 47, 52,  1, 58, 53,  1, 58, 46, 43, 51,  1, 24,\n",
      "        43, 58,  1, 59, 57,  1, 56, 43, 60, 43, 52, 45, 43,  1, 58, 46, 47, 57,\n",
      "         1, 61, 47, 58, 46,  0, 53, 59, 56,  1, 54, 47, 49, 43, 57,  6,  1, 43,\n",
      "        56, 43,  1, 61, 43,  1, 40, 43, 41, 53, 51, 43,  1, 56, 39, 49, 43, 57,\n",
      "        10,  1, 44, 53, 56,  1, 58, 46, 43,  1, 45, 53, 42, 57,  1, 49, 52, 53,\n",
      "        61,  1, 21,  0, 57, 54, 43, 39, 49,  1, 58, 46, 47, 57,  1, 47, 52,  1,\n",
      "        46, 59, 52, 45, 43, 56,  1, 44, 53, 56,  1, 40, 56, 43, 39, 42,  6,  1,\n",
      "        52, 53, 58,  1, 47, 52,  1, 58, 46, 47, 56, 57, 58,  1, 44, 53, 56,  1,\n",
      "        56, 43, 60, 43, 52, 45, 43,  8,  0,  0])\n"
     ]
    }
   ],
   "source": [
    "# let's now tokenize the entire text of dataset and store it into tesor\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000]) # the data that we look earlier will look like this into GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's now split up the data into train and val sets\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_size = 8 # aka context length\n",
    "data[:block_size + 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when input is tensor([18]) the target is : 47\n",
      "when input is tensor([18, 47]) the target is : 56\n",
      "when input is tensor([18, 47, 56]) the target is : 57\n",
      "when input is tensor([18, 47, 56, 57]) the target is : 58\n",
      "when input is tensor([18, 47, 56, 57, 58]) the target is : 1\n",
      "when input is tensor([18, 47, 56, 57, 58,  1]) the target is : 15\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target is : 47\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target is : 58\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"when input is {context} the target is : {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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 tensor([24]) the target is : 43\n",
      "when input is tensor([24, 43]) the target is : 58\n",
      "when input is tensor([24, 43, 58]) the target is : 5\n",
      "when input is tensor([24, 43, 58,  5]) the target is : 57\n",
      "when input is tensor([24, 43, 58,  5, 57]) the target is : 1\n",
      "when input is tensor([24, 43, 58,  5, 57,  1]) the target is : 46\n",
      "when input is tensor([24, 43, 58,  5, 57,  1, 46]) the target is : 43\n",
      "when input is tensor([24, 43, 58,  5, 57,  1, 46, 43]) the target is : 39\n",
      "when input is tensor([44]) the target is : 53\n",
      "when input is tensor([44, 53]) the target is : 56\n",
      "when input is tensor([44, 53, 56]) the target is : 1\n",
      "when input is tensor([44, 53, 56,  1]) the target is : 58\n",
      "when input is tensor([44, 53, 56,  1, 58]) the target is : 46\n",
      "when input is tensor([44, 53, 56,  1, 58, 46]) the target is : 39\n",
      "when input is tensor([44, 53, 56,  1, 58, 46, 39]) the target is : 58\n",
      "when input is tensor([44, 53, 56,  1, 58, 46, 39, 58]) the target is : 1\n",
      "when input is tensor([52]) the target is : 58\n",
      "when input is tensor([52, 58]) the target is : 1\n",
      "when input is tensor([52, 58,  1]) the target is : 58\n",
      "when input is tensor([52, 58,  1, 58]) the target is : 46\n",
      "when input is tensor([52, 58,  1, 58, 46]) the target is : 39\n",
      "when input is tensor([52, 58,  1, 58, 46, 39]) the target is : 58\n",
      "when input is tensor([52, 58,  1, 58, 46, 39, 58]) the target is : 1\n",
      "when input is tensor([52, 58,  1, 58, 46, 39, 58,  1]) the target is : 46\n",
      "when input is tensor([25]) the target is : 17\n",
      "when input is tensor([25, 17]) the target is : 27\n",
      "when input is tensor([25, 17, 27]) the target is : 10\n",
      "when input is tensor([25, 17, 27, 10]) the target is : 0\n",
      "when input is tensor([25, 17, 27, 10,  0]) the target is : 21\n",
      "when input is tensor([25, 17, 27, 10,  0, 21]) the target is : 1\n",
      "when input is tensor([25, 17, 27, 10,  0, 21,  1]) the target is : 54\n",
      "when input is tensor([25, 17, 27, 10,  0, 21,  1, 54]) the target is : 39\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4 # how many independet sequences we will process in parallel\n",
    "block_size = 8 # what is the maximum context length of the predictions?\n",
    "\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == \"train\" else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size, ))\n",
    "    # print(\"ix\")\n",
    "    # print(ix)\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1: i+block_size+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch(\"train\")\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 dimention\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b,t]\n",
    "        print(f\"when input is {context} the target is : {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "this contains 32 examples packed into a single batch. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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"
     ]
    }
   ],
   "source": [
    "# our input to transformer\n",
    "print(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 8, 65])\n",
      "tensor(4.8786, grad_fn=<NllLossBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "        \n",
    "    def forward(self, idx, targets=None):\n",
    "        # idx = inputs\n",
    "        # idx and targets are both (B, T) tensors of integers\n",
    "        logits = self.token_embedding_table(idx) # B,T,C\n",
    "        # we will interpret the value of lookup table as logits which will corresponds to the score for the next character in the sequence\n",
    "        if targets is None:\n",
    "            loss=None\n",
    "        else:\n",
    "            B,T,C = logits.shape\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits.view(B*T,C), targets) # need to reshape logits, because pytorch expects `C` to be 2nd dimension \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B,T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx)\n",
    "            # focus only on last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=1) # B,C\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # B,1\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # B, T+1\n",
    "        return idx \n",
    "    \n",
    "m = BigramLanguageModel(vocab_size)\n",
    "logits, loss= m(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How this works is: \"24\" will go the embedding table and will pluck out the 24th row. This row will be of dim `vocab_size`. Since we have 8 examples in one row, the output will be of dim `8 x vocab_size`. And we have 4 such examples in one batch, so the final output will be of dim `4 x 8 x vocab_size`.\n",
    "Then 43 will go to the embedding table and will pluck out the 43rd row. This is how the embedding table works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Sr?qP-QWktXoL&jLDJgOLVz'RIoDqHdhsV&vLLxatjscMpwLERSPyao.qfzs$Ys$zF-w,;eEkzxjgCKFChs!iWW.ObzDnxA Ms$3\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m.generate(idx, max_new_tokens=100)[0].tolist())) # since generation will be in batches, we need to pluck the first row and then decode to get back text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training simple model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.4735846519470215\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "for steps in range(10000):\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch(\"train\")\n",
    "    # evaluate the loss\n",
    "    logits, loss = m(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "print(loss.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "above cell is ran multiple times to get the final output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "H: ond fat d'sacairu:\n",
      "Theaive blshig; m Cam ancrthtinous hece meburgour ye.\n",
      "Bu, d alles : ppe saces se,\n",
      "FLe as, g s souranod'TE CLen w seseer urer? anthecach DUndooure sethe\n",
      "FRDUS:\n",
      "TAUpet m torenongs hiloroous heis RABe ve d owourithe hed, havid a t nd w f dingitlonet thed conond,\n",
      "\n",
      "F oo he thas hesce f bthe m ly y!\n",
      "\n",
      "KI d y my tine\n",
      "hy I wiperesouthe leant ghituthace h Yoopar, s t at s; f, wius ju k y thiean my.\n",
      "Maly it: e shiulitheaganory I d'tthilo phal tro hedr opue, ilvind whe tmeerorlvarind, \n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1,1), dtype=torch.long)\n",
    "print(decode(m.generate(idx, max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently we took only 1 previous token to predict the next token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's go to transformer model 🚀"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mathematical trick in self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 2])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# consider following toy example:\n",
    "torch.manual_seed(1337)\n",
    "B, T, C = 4, 8, 2 # batch, time, channel\n",
    "x = torch.randn(B,T,C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# best way to communicate with the previous token is to take avg of previous tokens till the current token\n",
    "\n",
    "# we want x[b,t] = mean_{i <= t} x x[b,i]\n",
    "xbow = torch.zeros((B,T,C)) # bow = bag of words\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": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1808, -0.0700],\n",
       "        [-0.3596, -0.9152],\n",
       "        [ 0.6258,  0.0255],\n",
       "        [ 0.9545,  0.0643],\n",
       "        [ 0.3612,  1.1679],\n",
       "        [-1.3499, -0.5102],\n",
       "        [ 0.2360, -0.2398],\n",
       "        [-0.9211,  1.5433]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0] # first batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1808, -0.0700],\n",
       "        [-0.0894, -0.4926],\n",
       "        [ 0.1490, -0.3199],\n",
       "        [ 0.3504, -0.2238],\n",
       "        [ 0.3525,  0.0545],\n",
       "        [ 0.0688, -0.0396],\n",
       "        [ 0.0927, -0.0682],\n",
       "        [-0.0341,  0.1332]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xbow[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=\n",
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.],\n",
      "        [1., 1., 1.]])\n",
      "--\n",
      "b=\n",
      "tensor([[2., 7.],\n",
      "        [6., 4.],\n",
      "        [6., 5.]])\n",
      "--\n",
      "c=\n",
      "tensor([[14., 16.],\n",
      "        [14., 16.],\n",
      "        [14., 16.]])\n",
      "--\n"
     ]
    }
   ],
   "source": [
    "# we wrote 2 for loops which are very inefficient, \n",
    "# much more efficient way is matrix multiplication, that's the trick\n",
    "\n",
    "# consider toy example\n",
    "torch.manual_seed(42)\n",
    "a = torch.ones((3,3))\n",
    "b = torch.randint(0, 10, (3,2)).float()\n",
    "c = a @ b\n",
    "print(\"a=\")\n",
    "print(a)\n",
    "print(\"--\")\n",
    "print(\"b=\")\n",
    "print(b)\n",
    "print(\"--\")\n",
    "print(\"c=\")\n",
    "print(c)\n",
    "print(\"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=\n",
      "tensor([[1., 0., 0.],\n",
      "        [1., 1., 0.],\n",
      "        [1., 1., 1.]])\n",
      "--\n",
      "b=\n",
      "tensor([[2., 7.],\n",
      "        [6., 4.],\n",
      "        [6., 5.]])\n",
      "--\n",
      "c=\n",
      "tensor([[ 2.,  7.],\n",
      "        [ 8., 11.],\n",
      "        [14., 16.]])\n",
      "--\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones((3,3))) # tril return lower triangular part\n",
    "b = torch.randint(0, 10, (3,2)).float()\n",
    "c = a @ b\n",
    "print(\"a=\")\n",
    "print(a)\n",
    "print(\"--\")\n",
    "print(\"b=\")\n",
    "print(b)\n",
    "print(\"--\")\n",
    "print(\"c=\")\n",
    "print(c)\n",
    "print(\"--\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "see how `c` has been calculated in the cell above. `2`, then `2+6` then `2+6+6`. So like it keeps on adding the previous values. This is the trick in self-attention. \n",
    "Now similarly we can calculate average also, just by normalizing the values i.e dividing each row by the sum of the rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "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",
      "--\n",
      "b=\n",
      "tensor([[2., 7.],\n",
      "        [6., 4.],\n",
      "        [6., 5.]])\n",
      "--\n",
      "c=\n",
      "tensor([[2.0000, 7.0000],\n",
      "        [4.0000, 5.5000],\n",
      "        [4.6667, 5.3333]])\n",
      "--\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones((3,3))) # tril return lower triangular part\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(\"a=\")\n",
    "print(a)\n",
    "print(\"--\")\n",
    "print(\"b=\")\n",
    "print(b)\n",
    "print(\"--\")\n",
    "print(\"c=\")\n",
    "print(c)\n",
    "print(\"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizing above for loop\n",
    "wei = torch.tril(torch.ones(T,T))\n",
    "wei = wei / wei.sum(1, keepdim=True) # wei = a\n",
    "xbow2 = wei @ x # (T,T) @ (B, T, C) ----> (B, T, C)\n",
    "# xbow2 will be identical to xbow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1808, -0.0700],\n",
       "        [-0.0894, -0.4926],\n",
       "        [ 0.1490, -0.3199],\n",
       "        [ 0.3504, -0.2238],\n",
       "        [ 0.3525,  0.0545],\n",
       "        [ 0.0688, -0.0396],\n",
       "        [ 0.0927, -0.0682],\n",
       "        [-0.0341,  0.1332]])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xbow2[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "# another version of writing the same thing: using softmax\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float(\"-inf\")) # in wei matrix, where tril == 0, make it -ve inf ; # future can't communicate with the past\n",
    "wei = F.softmax(wei, dim=-1) # it is also normalization operation, so we get similarly wei matrix\n",
    "xbow3 = wei @ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(xbow2, xbow3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "self attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 32])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# self attention\n",
    "torch.manual_seed(1337)\n",
    "B, T, C = 4, 8, 32 # batch, time, channel\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",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently we are initialzing weight matrices with zeros, but we don't want it to be all uniform. We want it to be data dependent. We want to gather information from the past but want to do it in a data dependent way. This is the problem that self attention solves.\n",
    "\n",
    "This is how self attention solves:\n",
    "- Every single node/token at each position will emit two vectors: a query and a key vector.\n",
    "  - Query vector: what am I looking for?\n",
    "  - Key vector: what do I contain?\n",
    "  - Then we take the dot product of the query and key vector. That dot product will now become our weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "B, T, C = 4, 8, 32 # batch, time, channel\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# implementing self attention for single head\n",
    "head_size = 16 # hyperparameter\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",
    "k = key(x)   # B, T, 16 (head_size)\n",
    "q = query(x) # B, T, 16\n",
    "\n",
    "wei = q @ k.transpose(-2, -1) # B, T, 16 @ B, 16, T ----> B, T, T\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",
    "out = wei @ v\n",
    "# out = wei @ x\n",
    "\n",
    "out.shape\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTES:\n",
    "\n",
    "- Attention is a communication mechanism.\n",
    "- Attention simply acts over a set of vectors. There is no notion of space. This is why we need to positionally encode the tokens. \n",
    "- Each example across batch dimension is independent and never \"talk\" to each other.\n",
    "- self-attention means that the keys and queries are produced from the same source. In cross-attention, the keys and queries are produced from different sources.\n",
    "- \"Scaled Attention\" additionally divides wei by 1/sqrt(head_size). This makes it so when input Q,K are unit variance, wei will be unit variance too and Softmax will stay diffuse and not saturate too much. Illustration below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = torch.randn(B,T, head_size)\n",
    "q = torch.randn(B,T, head_size)\n",
    "wei = q @ k.transpose(-2, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9006)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.0037)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(15.9317)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei.var()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Variance will be of head_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = torch.randn(B,T, head_size)\n",
    "q = torch.randn(B,T, head_size)\n",
    "wei = q @ k.transpose(-2, -1) * (head_size)**-0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(1.0966), tensor(0.9416), tensor(1.0065))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k.var(), q.var(), wei.var()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we divide by sqrt of head_size then var of wei will be 1. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this is important because\n",
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1) # if values of softmax will be close to 0 then output of softmax will be diffused"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # if inputs are extereme, then output values will be extreme. That is it will be too peaky"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm1D:\n",
    "    def __init__(self, dim, eps=1e-5):\n",
    "        self.eps = eps\n",
    "        self.gamma = torch.ones(dim)\n",
    "        self.beta = torch.ones(dim)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        # calculate the forward pass   \n",
    "            # we normalize row-wise in layer norm\n",
    "        xmean = x.mean(1, keepdim=True) # batch mean\n",
    "        xvar = x.var(1, keepdim=True) # batch variance\n",
    "        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)\n",
    "        self.out = self.gamma * xhat + self.beta\n",
    "        \n",
    "        return self.out\n",
    "    \n",
    "    def parameters(self):\n",
    "        return [self.gamma, self.beta]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (latest)",
   "language": "python",
   "name": "latest"
  },
  "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.1.-1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}