toy-gpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the tiny shakespeare dataset\n",
    "#!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in the dataset\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "characters in dataset: 1115394\n"
     ]
    }
   ],
   "source": [
    "print('characters in dataset:', 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",
      "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": [
    "# inspect at the first 1000 characters\n",
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "Vocab Size: 65\n"
     ]
    }
   ],
   "source": [
    "# get all unique characters in text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print('Vocab Size:', vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20, 43, 50, 50, 53, 1, 35, 53, 56, 50, 42, 2]\n",
      "Hello World!\n"
     ]
    }
   ],
   "source": [
    "# map characters to integers (tokenizer)\n",
    "stoi = {ch:i for i,ch in enumerate(chars)}\n",
    "itos = {i:ch for i,ch in enumerate(chars)}\n",
    "encode = lambda s: [stoi[c] for c in s] # input: string, output: list of ints\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # input: list of ints,\n",
    "                                                 # output: string\n",
    "\n",
    "print(encode('Hello World!'))\n",
    "print(decode(encode('Hello World!')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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": [
    "# encode the entire dataset and store in a PyTorch tensor\n",
    "import torch\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train-validation split\n",
    "n = int(0.9*(len(data))) # first 90% of data will be training data\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# need to break training data into chunks to reduce computational load\n",
    "block_size = 8\n",
    "train_data[:block_size+1] # with n+1 tokens, we can predict n next-word tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when input is tensor([18]) the target: 47\n",
      "when input is tensor([18, 47]) the target: 56\n",
      "when input is tensor([18, 47, 56]) the target: 57\n",
      "when input is tensor([18, 47, 56, 57]) the target: 58\n",
      "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
      "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n"
     ]
    }
   ],
   "source": [
    "# inspect inputs and targets for a single chunk\n",
    "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: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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] target is: 43\n",
      "when input is [24, 43] target is: 58\n",
      "when input is [24, 43, 58] target is: 5\n",
      "when input is [24, 43, 58, 5] target is: 57\n",
      "when input is [24, 43, 58, 5, 57] target is: 1\n",
      "when input is [24, 43, 58, 5, 57, 1] target is: 46\n",
      "when input is [24, 43, 58, 5, 57, 1, 46] target is: 43\n",
      "when input is [24, 43, 58, 5, 57, 1, 46, 43] target is: 39\n",
      "when input is [44] target is: 53\n",
      "when input is [44, 53] target is: 56\n",
      "when input is [44, 53, 56] target is: 1\n",
      "when input is [44, 53, 56, 1] target is: 58\n",
      "when input is [44, 53, 56, 1, 58] target is: 46\n",
      "when input is [44, 53, 56, 1, 58, 46] target is: 39\n",
      "when input is [44, 53, 56, 1, 58, 46, 39] target is: 58\n",
      "when input is [44, 53, 56, 1, 58, 46, 39, 58] target is: 1\n",
      "when input is [52] target is: 58\n",
      "when input is [52, 58] target is: 1\n",
      "when input is [52, 58, 1] target is: 58\n",
      "when input is [52, 58, 1, 58] target is: 46\n",
      "when input is [52, 58, 1, 58, 46] target is: 39\n",
      "when input is [52, 58, 1, 58, 46, 39] target is: 58\n",
      "when input is [52, 58, 1, 58, 46, 39, 58] target is: 1\n",
      "when input is [52, 58, 1, 58, 46, 39, 58, 1] target is: 46\n",
      "when input is [25] target is: 17\n",
      "when input is [25, 17] target is: 27\n",
      "when input is [25, 17, 27] target is: 10\n",
      "when input is [25, 17, 27, 10] target is: 0\n",
      "when input is [25, 17, 27, 10, 0] target is: 21\n",
      "when input is [25, 17, 27, 10, 0, 21] target is: 1\n",
      "when input is [25, 17, 27, 10, 0, 21, 1] target is: 54\n",
      "when input is [25, 17, 27, 10, 0, 21, 1, 54] target is: 39\n"
     ]
    }
   ],
   "source": [
    "# get mini-batches containing multiple chunks of data from the entire dataset.\n",
    "# we use mini-batches instead of single chunks to increase parallelizability\n",
    "torch.manual_seed(1337)\n",
    "batch_size = 4 # number of independent sequences (chunks) we process in parallel\n",
    "block_size = 8 # maximum context length (chunk token limit) for predictions\n",
    "\n",
    "def get_batch(split):\n",
    "    # generate a 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",
    "    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 dimension\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b, t]\n",
    "        print(f\"when input is {context.tolist()} target is: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": [
    "print(xb) # input to transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 65])\n",
      "tensor(4.8786, grad_fn=<NllLossBackward0>)\n",
      "\n",
      "SKIcLT;AcELMoTbvZv C?nq-QE33:CJqkOKH-q;:la!oiywkHjgChzbQ?u!3bLIgwevmyFJGUGp\n",
      "wnYWmnxKWWev-tDqXErVKLgJ\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "    # predictions for next token are made based only on current token\n",
    "    \n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads logits for next token from lookup table.\n",
    "        # each token has a unique row of size vocab_size to read logits from.\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "        \n",
    "    def forward(self, idx, targets=None):\n",
    "        \n",
    "        # idx and targets are both int tensors of shape (Batch, Time).\n",
    "        # for each token in each context window \n",
    "        # (Time -> num. of tokens in context window) \n",
    "        # in each chunk (Batch -> number of chunks in batch),\n",
    "        # pull out the corresponding logits (Channel -> vocab_size/feature size)\n",
    "        # from the lookup table\n",
    "        logits = self.token_embedding_table(idx) # (Batch, Time, Channel)\n",
    "        \n",
    "        # measure quality of logits based on target (how well are we predicting\n",
    "        # next character)\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            # reshape logits for pytorch: (Batch, Channel, Time)\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",
    "        # logits represent probability that a character is the correct next-word\n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B, T) array of indices in current context window.\n",
    "        # we are generating B next-word predictions and appending each\n",
    "        # prediction to the appropriate context window in the batch B\n",
    "        for _ in range(max_new_tokens):\n",
    "            # get 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 distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "model = BigramLanguageModel(vocab_size)\n",
    "logits, loss = model(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "idx = torch.zeros((1, 1), dtype=torch.long) # idx holds input sequence for \n",
    "                                            # generate, 0 represents start of\n",
    "                                            # sequence\n",
    "print(decode(model.generate(idx, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.382369041442871\n"
     ]
    }
   ],
   "source": [
    "# train the Bigram model\n",
    "batch_size = 32\n",
    "for steps in range(10000):\n",
    "    \n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "    \n",
    "    # evaluate loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "lso br. ave aviasurf my, yxMPZI ivee iuedrd whar ksth y h bora s be hese, woweee; the! KI 'de, ulseecherd d o blllando;LUCEO, oraingofof win!\n",
      "RIfans picspeserer hee tha,\n",
      "TOFonk? me ain ckntoty ded. bo'llll st ta d:\n",
      "ELIS me hurf lal y, ma dus pe athouo\n",
      "BEY:! Indy; by s afreanoo adicererupa anse tecorro llaus a!\n",
      "OLeneerithesinthengove fal amas trr\n",
      "TI ar I t, mes, n IUSt my w, fredeeyove\n",
      "THek' merer,\n"
     ]
    }
   ],
   "source": [
    "idx = torch.zeros((1, 1), dtype=torch.long)\n",
    "print(decode(model.generate(idx, max_new_tokens=400)[0].tolist()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mathematical trick in self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 2])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# consider following example\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,2 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# version 1: use for loops to average output.\n",
    "# want x[b,t] = mean_{i<=t} 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) # average over t to get avg C based on\n",
    "                                         # prev tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0.],\n",
       "        [1., 1., 0.],\n",
       "        [1., 1., 1.]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tril(torch.ones(3, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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"
     ]
    }
   ],
   "source": [
    "# torch.tril allows for fast running sum -> can be adapted for running average\n",
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones(3, 3))\n",
    "a = a / torch.sum(a, 1, keepdim=True)\n",
    "b = torch.randint(0,10,(3,2)).float()\n",
    "c = a @ b\n",
    "print('a=')\n",
    "print(a)\n",
    "print('--')\n",
    "print('b=')\n",
    "print(b)\n",
    "print('--')\n",
    "print('c=')\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# parallelize version 1 (averaging previous channels)\n",
    "wei = torch.tril(torch.ones(T, T))\n",
    "wei = wei / wei.sum(1, keepdim=True)\n",
    "xbow2 = wei @ x # ((B,) T, T) @ (B, T, C) -> (B, T, C)\n",
    "torch.allclose(xbow, xbow2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.1808, -0.0700],\n",
       "         [-0.0894, -0.4926],\n",
       "         [ 0.1490, -0.3199],\n",
       "         [ 0.3504, -0.2238],\n",
       "         [ 0.3525,  0.0545],\n",
       "         [ 0.0688, -0.0396],\n",
       "         [ 0.0927, -0.0682],\n",
       "         [-0.0341,  0.1332]]),\n",
       " 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": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xbow[0], xbow2[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.3333, 0.3333, 0.3333, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000],\n",
       "        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.0000],\n",
       "        [0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250, 0.1250]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T,T))  # weights represent \"interaction strength\" (or\n",
    "                          # affinities) with prev. tokens in self-attention.\n",
    "                          # tells us how much to weight prev. tokens info\n",
    "                          # when aggregating thru weighted sum\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))  # future tokens can't \n",
    "                                                 # communicate with past\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "wei"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 3: use Softmax (identical to Ver 1 and 2)\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",
    "xbow3 = wei @ x\n",
    "torch.allclose(xbow, xbow3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 32])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 4: self-attention - allows for data-dependent information aggregation\n",
    "# from past tokens. \n",
    "# every token emits 2 vectors: query and key. query represents \"what am I \n",
    "# interested in?\", key represents \"what do I contain?\"\n",
    "# key and query are emitted based on token's content and position (embeddings).\n",
    "# affinities are calculated through dot product of current token's query\n",
    "# and past tokens' keys -> becomes wei.\n",
    "# keys and queries that are \"aligned\" produce a high affinity, allowing \n",
    "# the current token to learn more about that past token as opposed to tokens\n",
    "# with lower affinity (high affinity = more info aggregated from that past token.)\n",
    "# tokens also emit value vector, represents \"if you find me interesting, here's\n",
    "# the info I will communicate to you.\" value vector is what gets aggregated \n",
    "# between different nodes (tokens).\n",
    "# by making output dependent on value vector instead of x, we can vary what info\n",
    "# a token will communicate to other tokens based on which attention head we are \n",
    "# using.\n",
    "# token information is kept in x (\"private\" to the token)\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32  # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# implement a single Head of self-attention\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",
    "k = key(x)   # (B, T, 16)\n",
    "q = query(x) # (B, T, 16)\n",
    "wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) -> (B, T, T)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "#wei = torch.zeros((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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5792, 0.1187, 0.1889, 0.1131, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0294, 0.1052, 0.0469, 0.0276, 0.7909, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0176, 0.2689, 0.0215, 0.0089, 0.6812, 0.0019, 0.0000, 0.0000],\n",
       "        [0.1691, 0.4066, 0.0438, 0.0416, 0.1048, 0.2012, 0.0329, 0.0000],\n",
       "        [0.0210, 0.0843, 0.0555, 0.2297, 0.0573, 0.0709, 0.2423, 0.2391]],\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notes:\n",
    "- Attention is a **communication mechanism**. Can be seen as nodes in a directed graph looking at each other and aggregating information with a weighted sum from all nodes that point to them, with data-dependent weights.\n",
    "  - First node points to itself, second node is pointed to by first node and itself, etc.\n",
    "  - Attention can be applied to any directed graph as a communication mechanism between the nodes\n",
    "- There is no notion of space. Attention simply acts over a set of vectors. This is why we need to positionally encode tokens.\n",
    "- Each example across batch dimension is of course processed completely independently and never \"talk\" to each other\n",
    "  - 4 seperate pools of 8 nodes\n",
    "- In an \"encoder\" attention block just delete the single line that does masking with `tril`, allowing all tokens to communicate. Block we are using is called a \"decoder\" attention block because it has triangular masking, and is usually used in autoregressive settings, like language modeling.\n",
    "- \"self-attention\" just means that the keys and values are produced from the same source as queries. In \"cross-attention\", the queries still get produced from x, but the keys and values come from some other, external source (e.g. an encoder module)\n",
    "  - cross-attention pulls information from a separate source of nodes\n",
    "- \"Scaled\" attention additional 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": 42,
   "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 # divide by sqrt of head_size\n",
    "                                                # to reduce wei's variance to 1.\n",
    "                                                # if wei's variance is left high,\n",
    "                                                # softmax will converge to a\n",
    "                                                # one-hot vector (especially \n",
    "                                                # important to avoid at \n",
    "                                                # initialization, otherwise each\n",
    "                                                # token will only learn from 1 \n",
    "                                                # other token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.0700)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.9006)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.1277)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 100])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class LayerNorm1d: \n",
    "    # normalizes rows to gaussian distribution w/ mean 0, variance 1\n",
    "    # (BatchNorm normalizes cols instead)\n",
    "  \n",
    "    def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
    "        self.eps = eps\n",
    "        self.gamma = torch.ones(dim)\n",
    "        self.beta = torch.zeros(dim)\n",
    "  \n",
    "    def __call__(self, x):\n",
    "        # calculate the forward pass\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) # normalize to unit variance\n",
    "        self.out = self.gamma * xhat + self.beta\n",
    "        return self.out\n",
    "  \n",
    "    def parameters(self):\n",
    "        return [self.gamma, self.beta]\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "module = LayerNorm1d(100)\n",
    "x = torch.randn(32, 100) # batch size 32 of 100-dimensional vectors\n",
    "x = module(x)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1469), tensor(0.8803))"
      ]
     },
     "execution_count": 49,
     "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": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-3.5763e-09), tensor(1.0000))"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch incl. all features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}