3-neural-nlp/Section 4.3 - Generative Pretrained Transformer (GPT).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6ec22c80-be7e-4dc2-80fa-5423bc8f1ccc",
   "metadata": {},
   "source": [
    "<img src='data/images/section-notebook-header.png' />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dfc3b51-89a6-429b-81be-af6b190987a0",
   "metadata": {},
   "source": [
    "**Disclaimer:** This source code of the GPT-2 model is adopted from Andrej Karpathy's [minGPT](https://github.com/karpathy/minGPT) and [nanoGPT](https://github.com/karpathy/nanoGPT) implementation. Overall, the code remains almost unchanged but there are some very minor modifications. The code has been slightly re-organized and some variable names and values have been modified to make their purpose more intuitive.\n",
    "\n",
    "The main additions to the implementation of GPT-2 are the added explanations and discussion of the code. This should make it easier to understand the individual components of the architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3823ed2a-c60a-498c-be8d-1970b56500e2",
   "metadata": {},
   "source": [
    "# Generative Pre-trained Transformer (GPT)\n",
    "\n",
    "### Overview\n",
    "\n",
    "The Generative Pretrained Transformer (GPT) architecture stands as a landmark in the field of natural language processing (NLP), representing a paradigm shift in how machines comprehend and generate human language. Developed by OpenAI, GPT is built upon the Transformer model, a neural network architecture introduced by Vaswani et al. in 2017. However, GPT extends and refines this architecture, leveraging the power of unsupervised learning to achieve remarkable proficiency in a wide range of language tasks.\n",
    "\n",
    "At its core, GPT harnesses the principles of transfer learning, where a model is first pretrained on vast amounts of text data in an unsupervised manner and then fine-tuned on specific tasks with labeled data. This approach enables GPT to capture intricate linguistic patterns, semantic relationships, and syntactic structures inherent in natural language, thereby endowing it with a deep understanding of human discourse. One of the defining features of the GPT architecture is its generative capability, allowing it to produce coherent and contextually relevant text based on a given prompt or input. This remarkable ability has found myriad applications across various domains, including language translation, text summarization, question answering, and creative writing.\n",
    "\n",
    "In this introduction, we delve into the intricacies of the GPT architecture, exploring its underlying mechanisms, training methodologies, and real-world applications. By understanding the inner workings of GPT, we gain insight into the transformative potential it holds for revolutionizing human-machine interaction and advancing the frontiers of artificial intelligence. More specifically, we look into the GPT-2 architecture. The GPT-2 model architecture introduced several notable characteristics and advancements, building upon its predecessor, GPT-1, to further enhance performance and versatility in natural language processing tasks. Some of the special characteristics of the GPT-2 model architecture include:\n",
    "\n",
    "* **Scale:** GPT-2 significantly increased the scale of its predecessor, with up to 1.5 billion parameters in its largest variant. This scale allowed GPT-2 to capture more intricate patterns and nuances in language, leading to improved performance across various tasks.\n",
    "\n",
    "* **Multi-layered Transformer Architecture:** Like GPT-1, GPT-2 relies on a multi-layered Transformer architecture, consisting of multiple encoder and decoder layers. These layers facilitate efficient processing of sequential data and enable the model to capture dependencies across different parts of the input text.\n",
    "\n",
    "* **Self-Attention Mechanism:** GPT-2 utilizes self-attention mechanisms within its Transformer architecture to weigh the importance of different words in a sentence based on their contextual relevance. This mechanism enables the model to capture long-range dependencies and maintain coherence in generated text.\n",
    "\n",
    "* **Unsupervised Pretraining:** Similar to GPT-1, GPT-2 undergoes unsupervised pretraining on large corpora of text data, such as books, articles, and websites. This pretraining phase allows the model to learn representations of language in an unsupervised manner, capturing diverse linguistic patterns and structures.\n",
    "\n",
    "* **Fine-Tuning for Specific Tasks:** While pretrained on a diverse range of text data, GPT-2 can be fine-tuned on specific tasks using supervised learning. This fine-tuning process involves providing labeled data for tasks such as language translation, text summarization, sentiment analysis, etc., enabling GPT-2 to adapt its learned representations to the nuances of the target task.\n",
    "\n",
    "* **Conditional Text Generation:** GPT-2 supports conditional text generation, where users can provide prompts or context to guide the generation of text. This feature allows for more controlled and targeted text generation, making GPT-2 suitable for various applications, including content creation, storytelling, and dialogue systems.\n",
    "\n",
    "* **High-Quality Text Generation:** GPT-2 is renowned for its ability to generate high-quality, coherent text across diverse topics and styles. The model exhibits impressive fluency, coherence, and semantic relevance in its generated outputs, often indistinguishable from human-written text, especially in smaller generations.\n",
    "\n",
    "Overall, these characteristics make the GPT-2 model architecture a powerful tool for a wide range of natural language processing tasks, contributing to advancements in AI-driven text generation and understanding. The focus in GPT-2 has two main reasons. Firstly, the model is still small enough in terms of the number of trainable parameters to run on a  consumer-grade desktop computer or a personal laptop. And secondly, OpenAI made the pretrained models of GPT-2 publicly available. This means that we do not have to train our implementation of GPT-2 from scratch but can copy over the pretrained weights from available models.\n",
    "\n",
    "**Important side note:** The ability to copy over the weights from existing pretrained models requires that our GPT-2 implementation must match the structure of the implementation from OpenAI. To some extent this will make the code a bit more difficult to understand, but enough details and explanations will be provided throughout the notebook to help with the overall understanding\n",
    "\n",
    "### Background: Transformers\n",
    "\n",
    "As the name Generative Pre-trained Transformer suggests, GPT is based on the Transformer architecture. It is therefore recommended that you first checkout the available Transformer notebooks. GPT is in some sense only \"half a transformer\" as it is considered a decoder-only architecture primarily because of its autoregressive nature and its reliance on generating text. In the original Transformer architecture introduced by Vaswani et al., both encoder and decoder components were utilized for tasks like machine translation, where the model needed to both understand the input text (encoder) and generate the translated text (decoder).\n",
    "\n",
    "However, in the case of GPT, the model is designed specifically for text generation tasks. During training, GPT learns to predict the next word in a sequence given the previous words (i.e., language modeling). This autoregressive process is inherently a decoding task, where the model generates text word-by-word based on the context provided by preceding words. Therefore, GPT lacks an explicit encoder component; instead, it consists solely of decoder layers. The figure below illustrates this but shows the decoder part of the \"full\" Transformer architecture. Of course, since there is no encoder, the decoder is missing the cross-attention component between the output of the encoder and the intermediate state of the decoder.\n",
    "\n",
    "<img src=\"data/images/gpt-transformer-architecture.png\" />\n",
    "\n",
    "In this decoder-only architecture, each token generated by the model depends on the previously generated tokens, mimicking the left-to-right generation process characteristic of language generation tasks. This design choice simplifies the architecture and training process for text generation while still achieving impressive results across a variety of language tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6019b983-ef50-4ea4-be7d-d9e6f4c45a21",
   "metadata": {},
   "source": [
    "## Setting up the Notebook\n",
    "\n",
    "### Import Required Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b3f0bb8-0a27-4ea5-83b8-5471a2e43dc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from src.gpt import CausalAttention, CausalMultiHeadAttention, MLPLayer, TransformerBlock, GPT, gpt_base_configs\n",
    "from src.bpe import BPETokenizer\n",
    "\n",
    "from src.utils import Dict2Class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e9fe93-bab7-47db-8279-307ae0d02115",
   "metadata": {},
   "source": [
    "### Checking/Setting the Computation Device\n",
    "\n",
    "PyTorch allows to train neural networks on supported GPUs to significantly speed up the training process. If you have a support GPU, feel free to utilize it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2517a4-86e1-48f0-a522-5424530f6756",
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "#use_cuda = False\n",
    "device = torch.device(\"cuda:0\" if use_cuda else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b641b9b6-fa1c-4825-b065-fdac83415fe1",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5bc36dc-0421-4cc6-b2c2-bb9bd8f5eb45",
   "metadata": {},
   "source": [
    "## Preliminaries\n",
    "\n",
    "GPT-2 is not a single model but a family of models that differ in their size with respect to the number of trainable parameters. GPT-2 comes in several different versions, each with varying sizes in terms of the number of parameters. These versions are:\n",
    "\n",
    "* **GPT-2 Small:** This version has 117M parameters and is the smallest variant of GPT-2. It is suitable for tasks where computational resources are limited or where a smaller model suffices.\n",
    "\n",
    "* **GPT-2 Medium:** With 345M parameters, GPT-2 Medium offers a balance between model size and performance. It provides improved performance compared to GPT-2 Small while remaining relatively computationally efficient.\n",
    "\n",
    "* **GPT-2 Large:** This version has 774M parameters, making it significantly larger than the Medium variant. GPT-2 Large achieves better performance on various NLP tasks due to its increased capacity to capture complex language patterns.\n",
    "\n",
    "* **GPT-2 XL:** With 1.5B parameters, GPT-2 XL is the largest publicly available version of GPT-2. It offers the highest capacity for capturing nuanced language patterns and has demonstrated superior performance on tasks requiring extensive context understanding.\n",
    "\n",
    "These different versions of GPT-2 cater to a wide range of use cases and computational requirements, allowing researchers and practitioners to choose the variant that best suits their needs. Additionally, they provide a continuum of performance and resource trade-offs, enabling experimentation with different model sizes depending on the specific task and available resources.\n",
    "\n",
    "The number of trainable parameters derives from the various configuration parameters such as the number of heads, the number of transformer blocks, the maximum allowed input text size, and more. For convenience, the file `gpt.py` (bottom), contains the configurations for all 4 GPT-versions for easy use. It also includes a toy configuration `gpt-tiny` yielding a very small model we can use for checking the individual components of the GPT architecture throughout the notebook.\n",
    "\n",
    "Let's have a quick look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca0882d2-e552-4e79-a8ab-0e8d1cd8edf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Dict2Class(gpt_base_configs['gpt-tiny'])\n",
    "\n",
    "for attribute, value in config.__dict__.items():\n",
    "    print(f\"{attribute}: {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd04420a-9cc5-4a22-8ce4-2d51dcbb7a0a",
   "metadata": {},
   "source": [
    "The main parameters (`num_heads`, `num_layers`, `block_size`, `mlp_factor`) will be explained in more detail later on. All other parameters are less relevant, but you should be able to easily check with the code where they are used to understand their meaning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "731d33dd-4d88-4622-a0bd-bcaec1c7c981",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db81b410-6ce6-4f43-adba-111cab3e3e3e",
   "metadata": {},
   "source": [
    "## GPT Architecture Components\n",
    "\n",
    "### Causal Attention\n",
    "\n",
    "Causal attention, also known as masked attention, is a key component of the GPT architecture that ensures the autoregressive nature of the model during training and inference. In the context of GPT, causal attention is used to prevent the model from attending to future tokens when generating each token in a sequence. The basis for the Causal Attention is the Scaled Dot-Product Attention from the Transformer architecture, where the attention scores are computed between each token in a sequence and all other tokens. These attention scores determine the importance of each token with respect to every other token in the sequence. The Scaled Dot-Product Attention mechanism is defined as follows:\n",
    "\n",
    "Given a query matrix $Q$, a key matrix $Q$, and a value matrix $V$, the attention scores $Attention(Q,K,V)$ are computed as:\n",
    "\n",
    "$$Attention(Q,K,V) = softmax(\\frac{QK^T}{\\sqrt{d_k}})V$$\n",
    "\n",
    "where $d_k$ is the dimensionality of the key vectors, and the softmax function is applied row-wise to obtain attention weights. For more details about the Scaled Dot-Product Attention, you can check out the notebook covering the basic Transformer architecture. In the context of GPT, causal attention extends this mechanism by applying a mask to the attention scores before the softmax operation. This mask ensures that each token can only attend to previous tokens in the sequence, and not to future tokens, preserving the autoregressive property required for text generation tasks.\n",
    "\n",
    "The causal mask is typically implemented by setting the attention scores corresponding to future positions to a large negative value (or negative infinity) before applying the softmax function. As a result, the softmax operation effectively assigns a probability of zero to future tokens, preventing the model from attending to them during training and inference.\n",
    "\n",
    "#### Worked Example for Causal Masking\n",
    "\n",
    "Let's create an instance of `CausalAttention`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e643c2-0672-4e2f-b5e3-d15499624fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Dict2Class(gpt_base_configs['gpt-tiny'])\n",
    "\n",
    "attention = CausalAttention(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2550ff-eb6f-46a8-9ee8-916c64bd7a81",
   "metadata": {},
   "source": [
    "Let's first have a look at the shape attention matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baa39664-ad97-4a6f-bd66-0a43f1094d5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Shape of Attention Mask: {attention.mask.shape}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c501ddf3-f67b-4e11-8543-5a97d3851323",
   "metadata": {},
   "source": [
    "The shape of the attention matrix is `(1, 1, block_size, block_size`) where `block_size` is the longest possible sequence that we can input to the GPT Transformer decoder. For our `gpt-tiny` configuration `block_size=64`, for all other support GPT-2 architectures, `block_size=1024`.\n",
    "\n",
    "Notice that our attention matrix has 2 additional dimensions. This is because the class `CausalAttention` attention will compute the attention matrices\n",
    "\n",
    "* For *all sequences* in a batch **and**\n",
    "* For *all heads* for a sequence\n",
    "\n",
    "We can also have a look at the attention mask itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a95fe8e-59c7-4fa2-bf6f-4792323ecb2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(attention.mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a4a051-92eb-4b30-ab2e-e57a8c6c96e6",
   "metadata": {},
   "source": [
    "Apart from the additional first 2 dimensions, the attention matrix is a matrix with the top-right half being all 0's and the bottom-left half incl. the diagonal all 1's. In a nutshell, this will ensure that:\n",
    "\n",
    "* The 1st token can only attend to itself\n",
    "* The 2nd token can only attend to the 1st token and itself (and vice versa)\n",
    "* The 3rd token can only attend to the 1st token, the 2nd token and itself (and vice versa)\n",
    "* ...\n",
    "\n",
    "In other words, each token can only attend to previous tokens but to future tokens.\n",
    "\n",
    "Let's actually show this by creating a random attention matrix, assuming a batch size of 8 and a 4 heads for the decoder. The code cell below creates a tensor containing $8\\cdot 4 = 32$ matrices -- one for each sequence and head -- representing an attention matrix for an input sequence of 5 tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c786a070-9c53-4d7e-a2cb-2bb2c110dec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "batch_size = 8\n",
    "seq_len    = 5\n",
    "\n",
    "attention_matrices = torch.rand(batch_size, config.num_heads, seq_len, seq_len)\n",
    "\n",
    "# Print the attenion matrix for the first sequence and the first head\n",
    "print(attention_matrices[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbd6ecf2-e5cd-4027-97bc-cfca3a9e6d3d",
   "metadata": {},
   "source": [
    "We can now apply the causal attention mask to 32 attention matrices in the tensor `attention_matrices`. This is done by setting each entry of `attention_matrices` to `-inf` where there is $0$ in the corresponding position in the attention mask. The single line in the code cell below accomplishes this. 2 additional comments:\n",
    "\n",
    "* Since our attention mask is too large (here: $64 \\times 64$), we need to shrink it to the size of the sequence length (here: $4$). We can simply do this conveniently using slicing: `attention.mask[:,:,:seq_len,:seq_len]`\n",
    "\n",
    "* Since the shape of the attention mask is now `(1, 1, seq_len, seq_len)` the command below applies the attention mask to all 24 attention matrices in the tensor `attention_matrices`. This is automatically done using [broadcasting](https://pytorch.org/docs/stable/notes/broadcasting.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06cd7e41-1b12-4ae5-a37c-9c062ebcb190",
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_matrices_masked = attention_matrices.masked_fill(attention.mask[:,:,:seq_len,:seq_len] == 0, float('-inf'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e279143f-d5f8-409d-8922-a4f36193fb88",
   "metadata": {},
   "source": [
    "We can now look at the same attention matrix again (see above), but now after the attention mask is applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "803ed3f3-2e84-4bde-92a5-7d66a3e9c53b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(attention_matrices_masked[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7f6eca-b2ab-4d80-bc8d-3c59c533f98c",
   "metadata": {},
   "source": [
    "Notice that this attention matrix -- and all attention matrices in `attention_matrices_masked` -- have now `-inf` values where there was a $0$ in the attention mask.\n",
    "\n",
    "The only core step missing is to compute the Softmax for all masked attention matrices to get the final attention weights. Recall that the attention weights w.r.t. each token have to sum up to $1$. An input attention weight of `-inf` will ensure that the value of the attention weight after the Softmax will be $0$, and all other weights summing up to $1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e02390-11e9-4711-9932-ecbcfcb77a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "attention_weights = F.softmax(attention_matrices_masked, dim=-1)\n",
    "\n",
    "# Print normalized weights of the same attention matrix (see above)\n",
    "print(attention_weights[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c0560c4-26b5-4108-9ded-d6dd3173a20d",
   "metadata": {},
   "source": [
    "#### Worked Example for `CausalAttention`\n",
    "\n",
    "In practice the `forward()` method of class `CausalAttention` expects a tensors of shape `(batch_size, num_heads, seq_len, embed_size)` since we compute the attention with respect to all sequences in the batch and all heads for each sequence. So let's create an example tensor of the correct shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b70ef2df-75a8-454f-b428-5489811c3419",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "seq_len    = 10\n",
    "\n",
    "causal_attention_input = torch.rand(batch_size, config.num_heads, seq_len, config.embed_size)\n",
    "\n",
    "print(causal_attention_input.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c323c6-8e6c-4482-8f0d-676c9a04ebff",
   "metadata": {},
   "source": [
    "Since we are still using the `gpt-tiny` model architecture, the number of heads is $4$ and the embedding size is $32$. This we can give as input to the `forward()` method of the `CausalAttention` class, which perform self-attention over the input sequences -- this is why we have 3x `causal_attention_input` -- as well as causal masking for each attention matrix as illustrated above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc59c66f-3b4d-4139-a9fc-5abfeb74f2e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "causal_attention_output = attention(causal_attention_input, causal_attention_input, causal_attention_input, seq_len)\n",
    "\n",
    "print(causal_attention_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8810f48b-39bc-40a7-b7f9-bd3ff580a28e",
   "metadata": {},
   "source": [
    "Of course, the output shape matches the input shape as the attention mechanism \"only\" changes embedding values for each token."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e19ad1ae-4673-4829-8a4f-3e72fdaee70a",
   "metadata": {},
   "source": [
    "### Multi-Head Casual Attention\n",
    "\n",
    "The class `CausalMultiHeadAttention` handles to computation of the self-attention for all heads by\n",
    "\n",
    "* Performing the linear transformation of the queries, keys, and values -- which in case of self-attention are all identical\n",
    "\n",
    "* Calling `CausalAttention` to perform to actual Scaled-Dot Product Attention incl. the application of the Causal Attention Matrix\n",
    "\n",
    "Note that there is only a single `nn.Linear` layer to handle queries, keys, and values; hence the definition as `nn.Linear(config.embed_size, 3*config.embed_size)`. This is conceptually the same as having 3 `nn.Linear` layers defined as `nn.Linear(config.embed_size, config.embed_size)`, but shows better performance in practice.\n",
    "\n",
    "Let's test this class by first creating some random input. Compared to the example input for `CausalAttention`, this tensor is missing the number of heads as dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fa2c14-935d-4882-ad67-33e37d995981",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_head_attention_input = torch.rand(batch_size, seq_len, config.embed_size)\n",
    "\n",
    "print(multi_head_attention_input.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48add8d9-fc02-4067-9755-d72519fe8e7f",
   "metadata": {},
   "source": [
    "Now we can define a `CausalMultiHeadAttention` layer to compute the output. Of course, the shape of the output is the same as the shape of the input. All the handling of multiple heads is done within this layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "012768f2-8a53-4535-873d-ec2cc33c8c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_head_attention = CausalMultiHeadAttention(config)\n",
    "\n",
    "multi_head_attention_output = multi_head_attention(multi_head_attention_input)\n",
    "\n",
    "print(multi_head_attention_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba09b246-01e3-4d08-b5b8-e6315c6e5088",
   "metadata": {},
   "source": [
    "### MLP Layer\n",
    "\n",
    "In the Transformer architecture, the feed-forward layer is a crucial component within the transformer blocks, which are the building blocks of the Transformer model. The feed-forward layer is applied independently to each position in the sequence, which is typically a sequence of word embeddings in natural language processing tasks. The feed-forward layer consists of two linear transformations with a non-linear activation function applied in between.\n",
    "\n",
    "The feedforward layer consists of two linear transformations with a ReLU activation function in between. The first linear transformation projects the input vector to a higher-dimensional space, and the second linear transformation maps it back to the original dimensionality. The output of the feedforward layer is computed as follows:\n",
    "\n",
    "$$FFN(x) = max(0, xW_1 + b_1)W_2 + b_2$$\n",
    "\n",
    "where $x$ is the input vector, $W_1$ and $W_2$ are learnable weight matrices, $b_1$ and $b_2$ are learnable bias vectors, and $max(0, \\cdot)$ represents the ReLU activation function.\n",
    "\n",
    "The primary purpose of the feed-forward layer in the Transformer architecture is to provide a mechanism for the model to learn complex, nonlinear relationships within the input sequence. By applying two linear transformations with a non-linear activation function in between, the feed-forward layer can model intricate patterns in the data, allowing the model to capture rich representations of the input sequences. This helps the Transformer model in tasks such as language modeling, machine translation, and other sequence-to-sequence tasks. Additionally, the feed-forward layer introduces flexibility and expressiveness to the Transformer architecture, enabling it to handle a wide range of natural language processing tasks effectively.\n",
    "\n",
    "The class `MLPLayer` implements this component in the GPT-2 architecture, but with 2 noteworthy differences\n",
    "\n",
    "* The size of the linear layers is expressed as a multiple of the embedding size instead as a fixed size. By default, the linear layers in the MLP / Feed-Forward layer have a size 4 times the embedding size.\n",
    "\n",
    "* GPT-2 uses as activation function between the 2 linear layers the [Gaussian Error Linear Unit (GeLU)](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html) activation function. GeLU is a non-linear activation function that was introduced as an alternative to the rectified linear unit (ReLU) in deep learning models. It is defined as:\n",
    "$$\\text{GeLU}(x) = x \\cdot \\Phi(x)$$\n",
    "where \\Phi(x) is the cumulative distribution function (CDF) of the standard normal distribution, defined as:\n",
    "$$\\Phi(x) = \\frac{1}{2} \\left(1 + \\text{erf}\\left(\\frac{x}{\\sqrt{2}}\\right)\\right)$$\n",
    "and erf denotes the error function.\n",
    "\n",
    "To give an example, we can run the output of the Multi-Head Attention class through the MLP layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13e9012-bad3-41c0-b498-750877aea3bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_layer = MLPLayer(config)\n",
    "\n",
    "mpl_layer_out = mlp_layer(multi_head_attention_output)\n",
    "\n",
    "print(mpl_layer_out.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f37037a2-840d-4bb5-863c-170ebfca58fd",
   "metadata": {},
   "source": [
    "Again, the shape of the output is the same as of the input."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "badf0ba0-1cea-4b08-b268-494957059299",
   "metadata": {},
   "source": [
    "### Transformer Block\n",
    "\n",
    "The `TransformerBlock` layer combines the individual components of Multi-Head Attention Layer and MLP Layer as well as\n",
    "\n",
    "* Adding layer normalization steps *and*\n",
    "\n",
    "* Residual connections\n",
    "\n",
    "in line with the basic Transformer architecture. Let's create an example input for the `TransformerBlock` layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc13090e-2ed2-4813-80ef-60b20e285ad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "block_input = torch.rand(batch_size, seq_len, config.embed_size)\n",
    "\n",
    "print(block_input.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d801d239-7771-4c0c-8ae5-d40683665f31",
   "metadata": {},
   "source": [
    "We can now define a `TransformerBlock layer` (which include `MultiHeadCausalAttention` and `MLPLayer`) and give it the example input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9a115df-225d-4d41-986c-6d03870eebdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_block = TransformerBlock(config)\n",
    "\n",
    "block_output = transformer_block(block_input)\n",
    "\n",
    "print(block_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3004f57c-455e-4228-99cd-20d73c45e903",
   "metadata": {},
   "source": [
    "### Final GPT Class\n",
    "\n",
    "The class `GPT` includes all components beyond the series of `TransformerBlock` layers like in the basic Transformer architecture, most notably:\n",
    "\n",
    "* The positional embedding\n",
    "\n",
    "* The final output layer (incl. an additional layer normalization steps)\n",
    "\n",
    "Let's have a look at the model for the configuration `gpt-tiny`. Notice the line `(0-2): 3 x TransformerBlock` indicating that this configuration uses 3 Transformer blocks. Apart from that, you should recognize all the components we have discussed individually before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b3ed1f-f2ec-4e74-b9c6-c7742dfab7ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.vocab_size = 10000 # some example value required for GPT class\n",
    "\n",
    "gpt = GPT(config)\n",
    "\n",
    "print(gpt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d3d61ab-537f-4053-82ba-0d9bd5738a28",
   "metadata": {},
   "source": [
    "We now have a complete implementation of the GPT-2 architecture implemented in PyTorch. In principle, you could use this implementation to train your own GPT-based Large Language Model. However, to achieve any meaningful performance, this would require to collect and process huge amounts of training data, as well as huge computational resources (ideally very large GPU-based computing clusters). These requirements make it virtually prohibitively expensive for individual users to train their own Large Language model such as GPT from scratch.\n",
    "\n",
    "Thus, instead of training a model from scratch, in this notebook, we make our implementation GPT \"practically usable\" by copying all trained weights from a pretrained model. Recall that this was the reason for some of the implementation details -- for example, we could have used 3 different `nn.Linear` layers to transform the queries, keys, and values, but this would have made the copying over of pretrained weights much more complicated. Of course, in practice it would be much more reasonable to directly use such a pretrained model to begin with. However, the goal of this notebook is to look under the hood of GPT and not to treat it like a block-box model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91ae7825-9240-4003-89dc-4da66b7b574e",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e298f216-8704-4391-a19e-ef75e4ffe52a",
   "metadata": {},
   "source": [
    "## Using GPT\n",
    "\n",
    "To make our own implementation usable, we have not address to issues:\n",
    "\n",
    "* We need to get a pretrained GPT-2 model to copy over all the weights into our instance of GPT-2\n",
    "\n",
    "* Preprocess and input we give our model in the same way as it was done for creating the pretrained model.\n",
    "\n",
    "### Create Instance from Pretrained Model\n",
    "\n",
    "The class `GPT` contains the auxiliary method that performs the following main steps\n",
    "\n",
    "* Create an instance of the `GPT` class based on the given configuration (e.g., `gpt2`, `gpt2-medium`, etc.).\n",
    "\n",
    "* Download the pretrained model of the same configuration and create an instance of that model; under the hood this is done by utility methods of the `transformer` package.\n",
    "\n",
    "* Copy the pretrained weights over to our instance of the `GPT` class. Notice that this step essentially just involves iterating in parallel through all layers in mode models and copies the weights over. This naturally requires both implementations to be \"equal enough\" in terms of having the same number of layers and in the same order.\n",
    "\n",
    "**Side note:** Which configuration you will be able to load depends on your available memory. For example, with 16GB of memory loading `gpt2-xl` is very likely to fail due to insufficient memory. The problem is that we need to load 2 models, the pretrained and our own \"copy\". But again, this notebook is not about state-of-the-art results -- after all, GPT-2 is already long obsolete -- but about understanding the more nitty-gritty details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e089ac59-5b82-4f93-81fd-5721634ba29a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GPT.from_pretrained('gpt2')         # 124M parameters\n",
    "#model = GPT.from_pretrained('gpt2-medium')  # 355M parameters\n",
    "#model = GPT.from_pretrained('gpt2-large')   # 774M parameters\n",
    "#model = GPT.from_pretrained('gpt2-xl')      # 1558M parameters -- fails to load; not enough memory :(\n",
    "#model = GPT.from_pretrained('distilgpt2')   # 82 parameters\n",
    "\n",
    "# Model model to device and set to eval mode\n",
    "model.to(device)\n",
    "\n",
    "# We never train this model, so let's just set it to evaluation mode\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fad1f6cb-359a-40ba-9985-c1fd52bb2d5d",
   "metadata": {},
   "source": [
    "### Input Tokenization using BPE\n",
    "\n",
    "Like most modern deep learning models for handling text as input, GPT-2 relies on subword-based tokenization. Subword-based tokenization is a technique used in natural language processing (NLP) to break down words into smaller subword units, which are then treated as tokens. This approach is particularly useful for handling out-of-vocabulary (OOV) words, rare words, and morphologically rich languages where word boundaries are not always clear. The basic idea is to create a vocabulary of subword units based on the input text corpus, which allows the model to represent both frequent and rare words more effectively.\n",
    "\n",
    "In traditional tokenization, words are treated as atomic units, meaning each word is considered as a single token. However, this approach can be problematic when dealing with morphologically complex languages or when encountering rare or unseen words during inference. Subword-based tokenization addresses these challenges by breaking down words into smaller units, such as character n-grams or other linguistically motivated subword units. This enables the model to generalize better to unseen words or morphological variations because it can still recognize and compose subword units that it has encountered during training.\n",
    "\n",
    "In contrast to traditional tokenization, where an input string is split into tokens according to predefined rules, in subword-based approaches **learn from data** and based on a given hyperparameter setting where to split a string into tokens. This means that 2 subword-based tokenizers may yield quite different token lists after processing an input text. As such, to make our GPT implementation work properly, we have to tokenize any input text the same way as was done for the pretrained model. The file `bpe.py` therefore contains the class `BPETokenizer` that implements the required tokenizer. Its source code is also adopted from Andrej Karpathy's [minGPT](https://github.com/karpathy/minGPT) and [nanoGPT](https://github.com/karpathy/nanoGPT) repositories.\n",
    "\n",
    "A deeper discussion of how this tokenizer works is beyond the scope of this notebook. Here, we merely create an instance for later use. Of course, you are very welcome to check out the code and Andrej's repository if you are interested in the inner workings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1151b7ab-cdfc-4a88-8c4e-0a572dee0abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BPETokenizer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18a5fe03-f2c6-4306-8344-a542e32c0680",
   "metadata": {},
   "source": [
    "Lastly, we define the method `generate()` below for actually using the tokenizer and the GPT model to generate responses. This method takes the following input parameters:\n",
    "\n",
    "* `model`: the instance of the GPT model.\n",
    "\n",
    "* `prompt`: the input text (or *prompt*) for the model; if the prompt is empty, the model will generate a response not conditioned on anything.\n",
    "\n",
    "* `num_samples`: the number of responses to be generated.\n",
    "\n",
    "* `steps`: the maximum number of tokens to be generated.\n",
    "\n",
    "* `do_sample`: if `True` the next token is sampled from the `top_k=10` (see method below) most likely tokens; if `False`, not surprisingly, all `num_sample` responses will be the same since always the most likely token will be used as the next generated token.\n",
    "\n",
    "Notice that this model calls the `generate()` method of the `GPT` class, which actually predicts/generates this next token. This method contains a loop that predicts the next token in each iteration. The input for an iteration is formed by the prompt and all so far generated tokens. If this input is too long, i.e., the input contains more than `block_size` tokens, the input is cut to the last `block_size` tokens.\n",
    "\n",
    "Of course, the return tokens are merely token indices w.r.t. to the underlying vocabulary. Therefore, as a last step, we have to use the tokenizer to convert the token indices to actual tokens/ words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2c6f2fa-105f-4efa-a2d5-8630c2545fca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(model, prompt='', num_samples=10, steps=20, do_sample=True):\n",
    "        \n",
    "    if prompt == '':\n",
    "        # to create unconditional samples...\n",
    "        # manually create a tensor with only the special <|endoftext|> token\n",
    "        # similar to what openai's code does here https://github.com/openai/gpt-2/blob/master/src/generate_unconditional_samples.py\n",
    "        x = torch.tensor([[tokenizer.encoder.encoder['<|endoftext|>']]], dtype=torch.long).to(device)\n",
    "    else:\n",
    "        x = tokenizer(prompt).to(device)\n",
    "    \n",
    "    # we'll process all desired num_samples in a batch, so expand out the batch dim\n",
    "    x = x.expand(num_samples, -1)\n",
    "\n",
    "    # forward the model `steps` times to get samples, in a batch\n",
    "    y = model.generate(x, max_new_tokens=steps, do_sample=do_sample, top_k=10)\n",
    "    \n",
    "    for i in range(num_samples):\n",
    "        out = tokenizer.decode(y[i].cpu().squeeze())\n",
    "        print('-'*80)\n",
    "        print(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62623b32-9373-4de9-8b52-f34f66993cd0",
   "metadata": {},
   "source": [
    "### Playing with GPT-2\n",
    "\n",
    "Now we are finally ready to use our GPT-2 model. Try different prompts to see which prompt seems to yield good responses, and which prompts do arguably completely fail.\n",
    "\n",
    "**\"Warning\":** Don't expect responses you are used to from using ChatGPT. A lot of improvements have been done from GPT-2 to ChatGPT (incl. various versions) that not only refers to the number of trainable parameters. For example, while GPT-2 XL has 1.5 Billion parameters, GPT-3 is said to have 175 billion parameters, although smaller versions seem to exist. In fact, OpenAI has only made the GPT-2 architecture public (so far), so you can find conflicting information about the more recent models online."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e1a74d-873f-4c4c-be8b-9ea58758bcd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate(model, prompt='What is Machine Learning?', num_samples=5, steps=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f7e6ab-9249-4903-90d8-7500528e5d8b",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0cdc232-2eaa-41b9-8612-c15a48e8d1f6",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "GPT-2, short for \"Generative Pre-trained Transformer 2,\" is a large language model developed by OpenAI. It represents a significant advancement in the field of natural language processing (NLP) and deep learning. GPT-2 builds upon the Transformer architecture, which has become a standard model for sequence-to-sequence tasks, such as language translation and text generation.\n",
    "\n",
    "One of the core characteristics of GPT-2 is its sheer size and complexity. It consists of a vast neural network with 1.5 billion parameters, making it one of the largest language models available at the time of its release. This extensive scale allows GPT-2 to capture intricate patterns and dependencies in natural language text, enabling it to generate coherent and contextually relevant responses across a wide range of tasks. Another key feature of GPT-2 is its ability to generate human-like text. By training on a diverse corpus of text data, GPT-2 learns to mimic the style, tone, and structure of the input text, producing outputs that often resemble human-authored content. This capability has significant implications for various applications, including text generation, conversational agents, content summarization, and creative writing.\n",
    "\n",
    "GPT-2's importance in the field of large language models stems from its performance and versatility. It has demonstrated state-of-the-art performance on various benchmark NLP tasks, including language modeling, text completion, and question answering. Additionally, its open-source release and widespread availability have facilitated research and development in the NLP community, sparking innovations in model architectures, training techniques, and downstream applications. Overall, GPT-2 represents a significant milestone in the advancement of large language models, paving the way for further progress in understanding and generating natural language text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca8e9c2-54a6-41ee-a80f-154117f915bc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py310]",
   "language": "python",
   "name": "conda-env-py310-py"
  },
  "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": 5
}