lets_learn_gpt.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building our own GPT\n",
    "\n",
    "Today we looked through the Transformers model. To show it via code, we will be building our GPT model. \n",
    "\n",
    "Chat-GPT was trained on all the publically available internet until 2021. That is an enormous amount of data. We don't have that much capacity to train our model. So we will be using a smaller dataset. For today, we will be using the TinyShakespeare dataset, which contains all the plays written by Shakespeare in a dialogue sort of format.\n",
    "\n",
    "At the end of the day, we want to be able to generate text like Shakespeare from thin air.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](image/shakespeare.png \"Title\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Exploration\n",
    "\n",
    "We already have the data, lets explore it to get a basic understanding of the data. Things to look for:\n",
    "1. How it looks like?\n",
    "2. What sort of structure is it in?\n",
    "3. Other things that you can think of?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read it in to inspect it\n",
    "with open('dataset/input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"length of dataset in characters: \", len(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's look at the first 1000 characters\n",
    "print(text[:1000])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question\n",
    "What other things would you like to explore in the data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_data_exploration_2.py\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that all the alphabets that exist in the English Literature are present in our dataset. Aside from the alphabets, we also have punctuations in our dataset. In total we have 65 characters in our vocabulary.\n",
    "\n",
    "From this point onwards, I am going to refer them as tokens."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, as the goal for today our GPT will output(generate) output as a single token by single token.\n",
    "Since we are dealing with NLP and Sequence Modelling, our output should depend on the previous tokens.\n",
    "But before we get into it, we have to process our data. Do you remember what was the first step in the Transformers architecture?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "\n",
    "Machine Learning models can't work with raw data. We have to convert it into a format that the model can understand. In this case, we have to convert our text into numbers. For today's case we will use the most simple encoding possible. We will assign a number to each token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a mapping from characters to integers\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] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DEMO\n",
    "Go back to slides and take a look at the other tokenization methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tiktoken\n",
    "enc = tiktoken.get_encoding(\"gpt2\")\n",
    "enc.n_vocab"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tiktoken is developed by OpenAI and most probably used by Chat-GPT as well. The thing that Chat-GPT generates will not be charcter by character but by token by token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](image/9000.jpeg \"Title\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are only 26 characters in English alphabets, then how is this number soo high?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc.encode(\"hii there\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc.decode([71, 4178, 612])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Things to note**: You can make a tradeoff between very long sequence of integers with a very small vocabulary or a very small sequence of integers with a very large vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
    "import torch # we use PyTorch: https://pytorch.org\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000]) # the 1000 characters we looked at earier will to the GPT look like this"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "0 is the newline and 1 is the space."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loading\n",
    "Before we get into training we should split our data into training and test. Because we dont want the model to just memorize the data. We want it to output, Shakespeare like text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's now split up the data into train and validation 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:]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After splitting the data, we have to convert it into a format that the model can understand. And as you know we cannot just simply feed the entire dataset to the model as it would be computationally expensive. So we have to split our data into chunks. And for today I will be referring to them as the block_size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To reiterate the concept of sequence modelling, we want our model to output a token (character) based on the previous sequence of tokens. In this case we want 47 to be referred from 18. If 18,47,56 appear in a sequence model should learn to output 57. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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: {target}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, from a block size of 8 we generate 8 examples that we want the model to train upon. This can be as small as 1 to the entire block size. This way we teach the transformer to not only look at the previous token but also the previous sequence of tokens."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batch time\n",
    "We have looked into one dimension of the sequence modelling. Now we have to look into the other dimension. We have to split our data into batches. This is because we cannot feed the entire dataset to the model at once. We have to split it into batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4 # how many independent sequences will we process in parallel?\n",
    "block_size = 8 # what is the maximum context length for 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",
    "    # batch_size random starting points\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()} the target: {target}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What we are trying to do is to get random chunks of data in the form of block_size from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337) # Bonus question: Why 1337?\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "    \n",
    "    def __init__(self) -> None:\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):\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx) # (B,T,C) -> (4,8,65)\n",
    "        loss = F.cross_entropy(logits, targets)\n",
    "        return logits\n",
    "\n",
    "m = BigramLanguageModel()\n",
    "loss = m(xb, yb)\n",
    "print(loss)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question\n",
    "Damn an error. What do we do now?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_vanilla_bigram_model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = BigramLanguageModel()\n",
    "logits, loss = m(xb, yb)\n",
    "print(loss)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Are we on the right track?\n",
    "A good yardstick to check if our first step is correct is to check the loss.\n",
    "Since the vocab_size is 65 and we have a model which hasn't been trained yet, the loss should be around -ln(1/65). If it is not, then we have to check our code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate ln(1/65)\n",
    "import math\n",
    "-1* math.log(1/65)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Will add this method in the bigram model to aid us in generating text\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 the 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_token = torch.zeros((1, 1), dtype=torch.long)\n",
    "m = BigramLanguageModel()\n",
    "print(decode(m.generate(my_token, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question- 2x\n",
    "Damn. Another error! What do you think is the problem? If you need hint, you just have to say the magic word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_complete_bigram_model.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, this is almost shit. But atleast we are now able to generate the output."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time to train\n",
    "\n",
    "![alt text](image/train.gif \"Title\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "\n",
    "# TODO(Ambuj): Run this multiple times to see the loss decrease\n",
    "\n",
    "for steps in range(1000): # increase number of steps for good results... \n",
    "    \n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_token = torch.zeros((1, 1), dtype=torch.long)\n",
    "print(decode(m.generate(my_token, max_new_tokens=100)[0].tolist()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### INTERMEZZO\n",
    "We trained our model, and this is a super simple Bigram model which is just a lookup table. We can do a lot better than this. We can use a Transformer model.\n",
    "We are not taking advantage of the history (**yet!**). But before we do that, we have to understand how the Transformer model works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# consider the following toy example:\n",
    "\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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I dont want my 5th token to communicate with the 6,7,8th token. I want it to communicate with the 1st,2nd,3rd,4th token. We only want information to flow from one direction(from the past) and not look into the future."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a simplistic toy example, I just want to calculate the mean of all the previous characters. This is not ideal because we are just taking the mean and the mean of the previous characters is not an ideal metric to calculate the next character. But for the sake of simplicity, lets just do it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
    "xbow = torch.zeros((B,T,C)) #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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE(self): Explain xbow\n",
    "xbow[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\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(\"b =\")\n",
    "print(b)\n",
    "print(\"c =\")\n",
    "print(c)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This is the mathematical trick in Self-Attention\n",
    "The Matrix a was the boring matrix in this case of all 1's. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
    "torch.manual_seed(1337)\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets do it efficiently."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version 2: Making it the pytorch way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# version 2: using matrix multiply for a weighted aggregation\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xbow[0], xbow2[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version 3: Adding softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# version 3: use Softmax\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) # Softmax=Normalization operation\n",
    "wei"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# version 3: use Softmax\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T,T)) # Affinity\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)\n",
    "# wei"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version 4: Self-Attention way!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each token will emit 3 vectors: Query, Key, Value.\n",
    "\n",
    "\n",
    "Query: What am I looking for?\n",
    "\n",
    "\n",
    "Key: What do I contain?\n",
    "\n",
    "\n",
    "Value: If you find me interesting, what should I propagate to the next layer?\n",
    "\n",
    "For attention, we will be using the dot product of the query and key. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](image/qkv.gif \"Title\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The formula for the self-attention is as follows:\n",
    "\n",
    "$$A(K,Q,V)=\\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# version 4: self-attention!\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# let's see a single Head perform 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 = 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": null,
   "metadata": {},
   "outputs": [],
   "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",
    "- 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",
    "- In an \"encoder\" attention block just delete the single line that does masking with `tril`, allowing all tokens to communicate. This block here 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",
    "- \"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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE(Ambuj): Comment about encoder block and decoder block. \n",
    "\n",
    "What changes should be done about the Sentiment Analysis model?\n",
    "\n",
    "Why 1/sqrt(head_size)?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wei.var()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conversation about diffused input and softmax saturation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # gets too peaky, converges to one-hot"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's take the attention for a spin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "batch_size = 16 # how many independent sequences will we process in parallel?\n",
    "block_size = 32 # what is the maximum context length for predictions?\n",
    "max_iters = 5000\n",
    "eval_interval = 100\n",
    "learning_rate = 1e-3\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "eval_iters = 200\n",
    "n_embd = 64\n",
    "n_head = 4\n",
    "n_layer = 4\n",
    "dropout = 0.0\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            logits, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time to make a choice\n",
    "\n",
    "![alt text](image/matrix.jpeg \"Choice\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Red Pill| Blue pill |\n",
    "| --- | ----------- |\n",
    "| Overall model provided | 2 Individual components provided |\n",
    "| Construct individual components | Construct Overall model |"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Head\n",
    "Contruct a Pytorch Module which has 3 tabs: Query, Key, Value. The output of the module should be the attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_head.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MultiHead\n",
    "Construct a pytorch model which takes n_heads and head_size as input. It should concatenate the output of the Head module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_multihead.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets look back at the architecture of the Transformer model.\n",
    "\n",
    "![alt text](image/architecture.png \"Title\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What are we still missing?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Feedforward\n",
    "After constructing Head and Multihead blocks only other individual components remaining is the Feedforward block. The Feedforward block is a simple 2 layer neural network with ReLU activation in between. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/feedforward.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Blocks\n",
    "Now that we have all the individual components, we can construct the blocks. The blocks are the building blocks of the Transformer model. The blocks are made up of the following components: Multihead, Feedforward, LayerNorm, Dropout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_blocks.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GPT\n",
    "Now that we have all the individual components, we can construct the GPT model. The GPT model is made up of the following components: Embedding, Positional Encoding, Blocks, Linear, loss function (log likelihood). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load answers/load_gpt.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Transformer()\n",
    "m = model.to(device)\n",
    "# print the number of parameters in the model\n",
    "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "\n",
    "    # every once in a while evaluate the loss on train and val sets\n",
    "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question\n",
    "The loss is not going down. What is the problem? Can you solve it?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results time\n",
    "Does our model really write like Shakespeare? Time to put it to test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate from the model\n",
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "GPT 1, 2, 3 are all decoder transformer models. They are autoregressive models, i.e. models that generates the next token based on the previous tokens. \n",
    "Now after building our first decoder model I want to put things into perspective for you guys. Our model has 4 layers, 4 heads, 64 dimensions embedding."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](image/all_models.png \"Title\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](image/shakespeare.png \"Title\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ab35b32f68116c0216e57a17ae4d2a651f25e42e7896b86e875bcfc913c2896e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}