minGPT_student.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tiny GPT Implementation\n",
    "\n",
    "Welcome to this walkthrough on how to implement GPT from scratch! \n",
    "\n",
    "Much of this notebook is taken from Andrej Karpathy's video [\"Let's build GPT\"](https://www.youtube.com/watch?v=kCc8FmEb1nY&ab_channel=AndrejKarpathy) and its corresponding resources. However, this notebook is reworked to provide a more notebook-first experience, to aid hands-on learning.\n",
    "\n",
    "This notebook will cover basic concepts such as attention and next-token prediction that are crucial to understanding how GPT works. This walkthrough will not cover much of the finer details of reproducing GPT performance. We'll be using a smaller dataset and single-GPU training.\n",
    "\n",
    "## Dataset\n",
    "Let's download our dataset that we will be training on. GPT-2 and later iterations of GPT were trained on closed-source, large, web-scale datasets. We'll instead be using a much smaller dataset for now for instructional and practical purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EDA\n",
    "Let's take a look at our dataset and what it looks like. First, we need to open it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's take a look at the length, some example text, and the alphabet we're dealing with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"======= Dataset Length =======\")\n",
    "print(\"Length of the dataset in characters:\", len(text))\n",
    "\n",
    "print(\"======= Sample Text =======\")\n",
    "print(text[:500])\n",
    "\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(\"======= Alphabet =======\")\n",
    "print(\"Alphabet:\", \"\".join(chars))\n",
    "print(\"Alphabet Size:\", vocab_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that our alphabet is consisted of uppercase and lowercase alphabetical characters, in addition to some punctuation and special characters. However, we need to convert these characters into some sort of number representation to feed them into our language model. This process is called **tokenization**. A simple way to transform our characters into tokens is just to take a simple mapping of them, assigning each unique character a number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate through all the numbers and create a bidirectional mapping\n",
    "# Start with the string -> integer mapping\n",
    "stoi = ...\n",
    "# Then create the integer -> string mapping\n",
    "itos = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's create an `encode(input: str) -> list[int]` function and a `decode(input: list[int]) -> str` function. `encode` will use our mapping to convert a string to a list of tokens representing the characters, and `decode` will convert a list of tokens to a string representing the tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encode = ... # encoder: take a string, output a list of integers\n",
    "decode = ... # decoder: take a list of integers, output a string"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try encoding and decoding a piece of text!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = \"Hello world!\"\n",
    "print(f\"Input: {input}\")\n",
    "print(\"==> encoding\")\n",
    "encoded = encode(input)\n",
    "print(encoded)\n",
    "print(\"<== decoding\")\n",
    "decoded = decode(encoded)\n",
    "print(decoded)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have just created our first encoder and decoder! Normally, tokenization is done with groups of characters rather than single characters like we are doing here. There's a trade off between accuracy and complexity when tokenizing. We'll discuss this more in depth when we talk about tokenization.\n",
    "\n",
    "Now that we have our encoder, we want to use it to preprocess our training data. Let's try encoding the entire dataset using our encoder!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's wrap it into a torch tensor to allow us to perform efficient matrix operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "data = torch.tensor(data, dtype=torch.long)\n",
    "print(data.shape, data.type())\n",
    "print(data[:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our data tokenized, let's split it into a training and validation split. Let's use a 90% train/val split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's now split up the data into train and validation sets\n",
    "percent_train = 0.9\n",
    "n = ... # first 90% will be train, rest val\n",
    "train_data = ...\n",
    "val_data = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we train a model, we want to have some sort of loss that we try to minimize, and something we want to predict. Many large language models today use next-token prediction. Just like the name suggests, at every step in the generation process we want to predict the token that comes next. We can then combine all these tokens to get our final output. \n",
    "\n",
    "Let's take a look at what this means. Let's say we have a `block_size` of 8. This is essentially how large our context window is. For example, when predicting the next token, with a `block_size` of 8, we will not look at any tokens more than 8 tokens behind the one we are trying to predict. In practice, block sizes are much bigger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This single block of 8 tokens is actually providing us a lot of input/output pairs we can use to train our model. For example, we know that `47` is the \"right answer\", or our target when the input is `[18]`. Extending this, `56` is our target when our input is `[18, 47]`, and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = train_data[:block_size+1]\n",
    "y = train_data[1:block_size+1]\n",
    "print(\"Block:\", x)\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"====\\nInput: {context}\\nTarget: {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is this the best we can do? We can actually increase our efficiency by utilizing the parallelism provided by GPUs to process multiple blocks at the same time. The number of blocks we process at the same time is called our batch size. Let's see what this looks like."
   ]
  },
  {
   "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 = ...\n",
    "    ix = ...\n",
    "    x = ...\n",
    "    y = ...\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "print(\"Training:\")\n",
    "print(xb)\n",
    "print(\"Target:\")\n",
    "print(yb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is the relation between the Training tensor and the Target tensor? Do you notice any similarities?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Model\n",
    "Now, let's build out the actual model that will allow us to generate text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "\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",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx) # (B,T,C)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    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 = ...\n",
    "            # focus only on the last time step\n",
    "            logits = ... # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = ... # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = ... # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = ... # (B, T+1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, instantiate your model and generate 100 text tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ...\n",
    "logits, loss = ...\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "print(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, this produces gibberish. We haven't actually trained our model, so let's do that. We will follow the paper \"Attention is All You Need\" to implement multi-headed self-attention. But first, what is attention?\n",
    "\n",
    "# Attention\n",
    "According to the paper \"Attention is All You Need\", attention is defined as the expression below (more specifically, scaled dot-product attention).\n",
    "$$\\text{Attention}(Q, K ,V) = \\text{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V$$\n",
    "But what does this mean exactly? When we want to predict the next token, we need to find a way to use the tokens that came before it as information. \n",
    "\n",
    "There are many ways to do this with varying complexity, but the simplest way is to take a weighted average of the tokens that came before the current token. For example, we want the 5th token to obtain information (\"pay attention to\") the 4th, 3rd, 2nd, and 1st tokens, but not the 6th or 7th token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can implement this in a for loop for each token, but doing so is very inefficient. Instead, we can use matrices to utilize the parallelism provided by NumPy matrix operations. First, we can devise a matrix that allows each token to only communicate with tokens that appear prior to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "a = ...\n",
    "a = ...\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what this does when multiplied with another matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = ...\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = ...\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 1st row of matrix $c$ is equal to the 1st row of matrix $b$, the 2nd row is equal to the average of the 1st and 2nd rows, and the 3rd row is the average of the 1st, 2nd, and 3rd rows. You can verify this for yourself.\n",
    "\n",
    "We can also do this as a batched operation, by extending this into a 3 dimensional array. PyTorch infers extra dimensions automatically, so we don't need to change our weight matrix. \n",
    "\n",
    "For example, if we have a weight matrix with shape $(T, T)$ (representing relative weights in the time dimension), and multiply this with a matrix $A$ with shape $(B, T, C)$ (batch, time, channels), PyTorch will infer the outermost $B$ dimension for our weight matrix. The result will be a $(B, T, T) \\times (B, T, C) \\rightarrow (B, T, C)$ matrix.\n",
    "\n",
    "What a $(B,T,C)$ dimensional matrix actually means in this contex is that we have a $B \\times T$ arrangement of tokens, while the information contained in each token is $C$ dimensional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "B,T,C = ... # batch, time, channels\n",
    "wei = ...\n",
    "wei = ...\n",
    "x = ...\n",
    "wei = ...\n",
    "print(x.shape) # we expect this to be (B, T, C) = (4, 8, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also construct this weighted matrix in a variety of different ways. For example, we can use a softmax, which does the same thing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 4\n",
    "tril = ...\n",
    "wei = ...\n",
    "wei = ...\n",
    "wei = ...\n",
    "wei"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Keys, Queries, Values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's revisit how we aggregate information. Previously, we took a simple average of all tokens prior to our current token. However, this causes each token to be weighted uniformly equally. In practice, certain tokens are probably more important for certain words than others. For example, in the sentence \"the red dog\", the word \"red\" is probably more important to the word \"dog\" than the word \"the\". As such, we want to be able to model this.\n",
    "\n",
    "A Transformer does this with 3 sets of learnable values. These are called keys, queries, and values. You might have seen these represented as $(K, Q, V)$.\n",
    "\n",
    "When we want to predict a token, we send out a Query, which can be interpreted as a sort of request. We take dot-products of the Query with different Keys, which act as a sort of label for what they contain. The higher this dot-product, the more we pay attention to that word when computing the current token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll start with single-headed attention first, then move to multi-headed attention. First, we'll create a random matrix with dimensions $(B,T,C)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "B,T,C = ... # batch, time, channels\n",
    "x = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Keys, queries, values are learned values. We can represent them as a simple linear layer. We want this to be a simple linear transform, so we don't want to set a bias term. This will have dimensions $(C, \\text{head size})$. TO DO: explain $\\text{head size}$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "head_size = 16\n",
    "key = ...\n",
    "query = ...\n",
    "value = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's pass our input data through $K$ and $Q$ individually. Notice that in this stage, they don't communicate yet. This output will have dimensions $(B, T, C) \\times (B, C, \\text{head size}) \\rightarrow (B, C, \\text{head size})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = ...   # (B, T, 16)\n",
    "q = ... # (B, T, 16)\n",
    "print(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we want to take the dot product. The efficient way to do this is through a matrix multiply. However, both are $(B, C, \\text{head size})$. We need to transpose one of them before multiplying to get a matrix of shape $(B, C, C)$. These will be our new weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wei = ...\n",
    "wei.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, there is a problem. Try printing the first row. Current tokens attend to future tokens, which should not happen. As a result, we need to apply the same transform as before to mask future tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tril = ...\n",
    "wei = ...\n",
    "print(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to before, we want to take a softmax over this to make each row resemble a probability distribution that sums to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wei = ...\n",
    "print(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that while the general structure of the matrix is the same, prior tokens on each row are no longer equally weighted, but instead weighted according to the dot-product of their keys with the query.\n",
    "\n",
    "Next, we need to pass our inputs through $V$, or our value matrix. Then, similar to the previous portions, we want to weight our values. In other words, we want to multiply our weights with our values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = ...\n",
    "out = ...\n",
    "print(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a few things to notice here.\n",
    "1. This is positionally invariant, at least for now. There is no notion of position. We can rearrange the words in the sentence and it will still work, although our predictions would similarly be misordered.\n",
    "2. Batch dimensions are always independent. There is no communication across different batches.\n",
    "\n",
    "This is also a good time to explain some different concepts:\n",
    "1. The above example is a \"decoder\" block. This is a decoder block because we mask future tokens. \"Encoder\" blocks operate in the same way, but without masking. This is because when encoding we can see future tokens, but when generating (decoding) we can't.\n",
    "2. \"Self-Attention\" means that the keys and values are produced from the same source as queries. Other types of attention such as cross-attention the queries are produced from x but the keys and values come from some other source, such as an encoder."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nanoGPT",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}