decoder_transformers_with_pytorch_and_lightning_v2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "9f40c232-df6e-49df-9016-6459e4af2e1e",
      "metadata": {
        "tags": [],
        "id": "9f40c232-df6e-49df-9016-6459e4af2e1e"
      },
      "source": [
        "# StatQuest: Coding Transformers from Scratch!!!\n",
        "## Part 1: Decoder-Only Transformers\n",
        "\n",
        "Copyright 2024, Joshua Starmer"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f4226d63-8d76-40bc-a8e6-0f290a159418",
      "metadata": {
        "id": "f4226d63-8d76-40bc-a8e6-0f290a159418"
      },
      "source": [
        "----\n",
        "\n",
        "In this tutorial, we will use **[PyTorch](https://pytorch.org/) + [Lightning](https://www.lightning.ai/)** to create and optimize a **Decoder-Only Transformer**, like the one shown in the picture below. Decoder-Only Transformers are taking over AI right now, and quite possibly their most famous use is in ChatGPT.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/dec_transformer.png?raw=1\" alt=\"a decoder-only transformer neural network\" style=\"width: 800px;\">\n",
        "\n",
        "Although Decoder-Only Transformers look complicated and can do really cool things, the good news is that they don't actually require a lot of code. A lot of their power comes from simply making multiple copies of each component. So, with that said...\n",
        "\n",
        "In this tutorial, you will...\n",
        "\n",
        "- **[Code a Position Encoder Class From Scratch!!!](#position)** The position encoder gives a transformer a way to keep track of the order of the input tokens.\n",
        "\n",
        "- **[Code an Attention Class From Scratch!!!](#attention)** The attention class allows the transformer to keep track of the relationships among words in the input and the output.\n",
        "\n",
        "- **[Code a Decoder-Only Transformer Class From Scratch!!!](#decoder)** The Decoder-Only Transformer will combine the position encoder and attention classes that we wrote with built-in pytorch classes to process the user input and generate the output.\n",
        "\n",
        "- **[Train the Transformer!!!](#train)** We'll train the transformer to answer simple questions.\n",
        "\n",
        "- **[Use the Trained Transformer!!!](#use)** Finally, we'll use the transformer to answer simple questions.\n",
        "\n",
        "#### NOTE:\n",
        "This tutorial assumes that you already know the basics of coding in **Python** and are familiar with the theory behind **[Decoder-Only Transformers](https://youtu.be/bQ5BoolX9Ag)** and **[Backpropagation](https://youtu.be/IN2XmBhILt4)**. It also assumes that you are familiar with the **[Essential Matrix Algebra for Neural Networks](https://youtu.be/ZTt9gsGcdDo)** and how it applies to **[Transformers](https://youtu.be/KphmOJnLAdI)**. If not, check out the **StatQuests** by clicking on the links for each topic.\n",
        "\n",
        "#### ALSO NOTE:\n",
        "I strongly encourage you to play around with the code. Playing with the code is the best way to learn from it."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "855d3a52-a43e-44cf-b8b6-b2ce88bea382",
      "metadata": {
        "id": "855d3a52-a43e-44cf-b8b6-b2ce88bea382"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "63b86036-1369-441c-a14d-3ba7bdaa103b",
      "metadata": {
        "tags": [],
        "id": "63b86036-1369-441c-a14d-3ba7bdaa103b"
      },
      "source": [
        "# Import the modules that will do all the work\n",
        "\n",
        "The very first thing we need to do is load a bunch of Python modules. Python itself is just a basic programming language. These modules give us extra functionality to create and train a Neural Network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "5c520e0b-c6e4-43ce-93f5-c0f2b5e75438",
      "metadata": {
        "id": "5c520e0b-c6e4-43ce-93f5-c0f2b5e75438",
        "outputId": "5e665593-d73d-4272-83ce-caa806a880c8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting lightning\n",
            "  Downloading lightning-2.2.4-py3-none-any.whl (2.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m10.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: PyYAML<8.0,>=5.4 in /usr/local/lib/python3.10/dist-packages (from lightning) (6.0.1)\n",
            "Requirement already satisfied: fsspec[http]<2025.0,>=2022.5.0 in /usr/local/lib/python3.10/dist-packages (from lightning) (2023.6.0)\n",
            "Collecting lightning-utilities<2.0,>=0.8.0 (from lightning)\n",
            "  Downloading lightning_utilities-0.11.2-py3-none-any.whl (26 kB)\n",
            "Requirement already satisfied: numpy<3.0,>=1.17.2 in /usr/local/lib/python3.10/dist-packages (from lightning) (1.25.2)\n",
            "Requirement already satisfied: packaging<25.0,>=20.0 in /usr/local/lib/python3.10/dist-packages (from lightning) (24.0)\n",
            "Requirement already satisfied: torch<4.0,>=1.13.0 in /usr/local/lib/python3.10/dist-packages (from lightning) (2.2.1+cu121)\n",
            "Collecting torchmetrics<3.0,>=0.7.0 (from lightning)\n",
            "  Downloading torchmetrics-1.4.0.post0-py3-none-any.whl (868 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m868.8/868.8 kB\u001b[0m \u001b[31m16.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm<6.0,>=4.57.0 in /usr/local/lib/python3.10/dist-packages (from lightning) (4.66.4)\n",
            "Requirement already satisfied: typing-extensions<6.0,>=4.4.0 in /usr/local/lib/python3.10/dist-packages (from lightning) (4.11.0)\n",
            "Collecting pytorch-lightning (from lightning)\n",
            "  Downloading pytorch_lightning-2.2.4-py3-none-any.whl (802 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m802.2/802.2 kB\u001b[0m \u001b[31m18.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<2025.0,>=2022.5.0->lightning) (2.31.0)\n",
            "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<2025.0,>=2022.5.0->lightning) (3.9.5)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from lightning-utilities<2.0,>=0.8.0->lightning) (67.7.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch<4.0,>=1.13.0->lightning) (3.14.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch<4.0,>=1.13.0->lightning) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch<4.0,>=1.13.0->lightning) (3.3)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch<4.0,>=1.13.0->lightning) (3.1.4)\n",
            "Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB)\n",
            "Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB)\n",
            "Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB)\n",
            "Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB)\n",
            "Collecting nvidia-cublas-cu12==12.1.3.1 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB)\n",
            "Collecting nvidia-cufft-cu12==11.0.2.54 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB)\n",
            "Collecting nvidia-curand-cu12==10.3.2.106 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB)\n",
            "Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB)\n",
            "Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB)\n",
            "Collecting nvidia-nccl-cu12==2.19.3 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl (166.0 MB)\n",
            "Collecting nvidia-nvtx-cu12==12.1.105 (from torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB)\n",
            "Requirement already satisfied: triton==2.2.0 in /usr/local/lib/python3.10/dist-packages (from torch<4.0,>=1.13.0->lightning) (2.2.0)\n",
            "Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch<4.0,>=1.13.0->lightning)\n",
            "  Using cached nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (23.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>=2022.5.0->lightning) (4.0.3)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch<4.0,>=1.13.0->lightning) (2.1.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]<2025.0,>=2022.5.0->lightning) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]<2025.0,>=2022.5.0->lightning) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]<2025.0,>=2022.5.0->lightning) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->fsspec[http]<2025.0,>=2022.5.0->lightning) (2024.2.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch<4.0,>=1.13.0->lightning) (1.3.0)\n",
            "Installing collected packages: nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, lightning-utilities, nvidia-cusparse-cu12, nvidia-cudnn-cu12, nvidia-cusolver-cu12, torchmetrics, pytorch-lightning, lightning\n",
            "Successfully installed lightning-2.2.4 lightning-utilities-0.11.2 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.19.3 nvidia-nvjitlink-cu12-12.4.127 nvidia-nvtx-cu12-12.1.105 pytorch-lightning-2.2.4 torchmetrics-1.4.0.post0\n"
          ]
        }
      ],
      "source": [
        "!pip install lightning\n",
        "\n",
        "import torch ## torch let's us create tensors and also provides helper functions\n",
        "import torch.nn as nn ## torch.nn gives us nn.Module(), nn.Embedding() and nn.Linear()\n",
        "import torch.nn.functional as F # This gives us the softmax() and argmax()\n",
        "from torch.optim import Adam ## We will use the Adam optimizer, which is, essentially,\n",
        "                             ## a slightly less stochastic version of stochastic gradient descent.\n",
        "\n",
        "import lightning as L ## Lightning makes it easier to write, optimize and scale our code\n",
        "from torch.utils.data import TensorDataset, DataLoader ## We'll store our data in DataLoaders"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6e58ebb2-1798-41c3-9ff8-1b4f50605964",
      "metadata": {
        "id": "6e58ebb2-1798-41c3-9ff8-1b4f50605964"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fa9ddca3-53b6-451b-8fe1-6116e1f7f473",
      "metadata": {
        "id": "fa9ddca3-53b6-451b-8fe1-6116e1f7f473"
      },
      "source": [
        "# Create the input and output and data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "67d17fd2-ac78-4664-93bb-4dd64c448e1e",
      "metadata": {
        "id": "67d17fd2-ac78-4664-93bb-4dd64c448e1e"
      },
      "source": [
        "In this tutorial we will build a simple Decoder-Only Transformer that can answer two super simple questions, **What is StatQuest?** and **StatQuest is what?**, and give them both the same answer, **Awesome!!!**\n",
        "\n",
        "In order to keep track of our simple dataset,\n",
        "we'll create a dictionary that maps the words and tokens to ID numbers. This is because the class we will use to do word embedding for us, `nn.Embedding()`, only accepts ID numbers as input, rather than words or tokens. Then we will use the dictionary to create a **Dataloader** that contains the questions and the desired answers encoded as ID numbers. Ultimately we'll use the **Dataloader** to train the transformer. **NOTE:** Dataloaders are designed to scale to very large datasets, so this simple example should be useful even when you have a terabyte of text.\n",
        "\n",
        "**ALSO NOTE:** The **inputs** and **labels** for the training data used with a Decoder-Only Transformer can seem a little strange at first. This is because a Decoder-Only Transformer generates a lot of the user input in addition to the response. To get a sense of what this means, let's pretend we want to train our Decoder-Only Transformer to answer the question **What is StatQuest?** with the response **Awesome**. In the figure below, on the left side, we see that the first token in the input **What** generates the output **is**. During training, we can compare this output to the known second token in the input, and if it is different, use that difference to modify the weights and biases in the model. Thus, even though **is** is part of the **input**, it is also part of the **label** that we use to evaluate how well the Decoder-Only Transformer is performing and whether or not the weights and biases should be changed. Likewise, **StatQuest**, **<\\EOS>**, and **awesome** can also be in both the **input** and in the **label** because we know we want the Decoder-Only Transformer to use them as inputs and generate them as outputs."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "752cda95-7e7e-430a-8c1a-4afc9670379b",
      "metadata": {
        "id": "752cda95-7e7e-430a-8c1a-4afc9670379b"
      },
      "source": [
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/expected_input_output_1.png?raw=1\" alt=\"The expected input and outputs during training\" style=\"width: 800px;\">"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "63631347-8db6-4812-8198-9169b2df0b24",
      "metadata": {
        "id": "63631347-8db6-4812-8198-9169b2df0b24",
        "outputId": "fd5745b6-56d1-400f-cb05-9058e3d96b83",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[1, 2, 4, 3, 4],\n",
            "        [1, 0, 4, 3, 4]])\n"
          ]
        }
      ],
      "source": [
        "## first, we create a dictionary that maps vocabulary tokens to id numbers...\n",
        "token_to_id = {'what' : 0,\n",
        "               'is' : 1,\n",
        "               'statquest' : 2,\n",
        "               'awesome': 3,\n",
        "               '<EOS>' : 4, ## <EOS> = end of sequence\n",
        "              }\n",
        "## ...then we create a dictionary that maps the ids to tokens. This will help us interpret the output.\n",
        "## We use the \"map()\" function to apply the \"reversed()\" function to each tuple (i.e. ('what', 0)) stored\n",
        "## in the token_to_id dictionary. We then use dict() to make a new dictionary from the\n",
        "## reversed tuples.\n",
        "id_to_token = dict(map(reversed, token_to_id.items()))\n",
        "\n",
        "## NOTE: Because we are using a Decoder-Only Transformer, the inputs contain\n",
        "##       the questions (\"what is statquest?\" and \"statquest is what?\") followed\n",
        "##       by an <EOS> token followed by the response, \"awesome\".\n",
        "##       This is because all of those tokens will be used as inputs to the Decoder-Only\n",
        "##       Transformer during Training. (See the illustration above for more details)\n",
        "## ALSO NOTE: When we train this way, it's called \"teacher forcing\".\n",
        "##       Teacher forcing helps us train the neural network faster.\n",
        "inputs = torch.tensor([[token_to_id[\"what\"], ## input #1: what is statquest <EOS> awesome\n",
        "                        token_to_id[\"is\"],\n",
        "                        token_to_id[\"statquest\"],\n",
        "                        token_to_id[\"<EOS>\"],\n",
        "                        token_to_id[\"awesome\"]],\n",
        "\n",
        "                       [token_to_id[\"statquest\"], # input #2: statquest is what <EOS> awesome\n",
        "                        token_to_id[\"is\"],\n",
        "                        token_to_id[\"what\"],\n",
        "                        token_to_id[\"<EOS>\"],\n",
        "                        token_to_id[\"awesome\"]]])\n",
        "\n",
        "## NOTE: Because we are using a Decoder-Only Transformer the outputs, or\n",
        "##       the predictions, are the input questions (minus the first word) followed by\n",
        "##       <EOS> awesome <EOS>.  The first <EOS> means we're done processing the input question\n",
        "##       and the second <EOS> means we are done generating the output.\n",
        "##       See the illustration above for more details.\n",
        "# labels = torch.tensor([[token_to_id[\"is\"],\n",
        "#                        token_to_id[\"statquest\"],\n",
        "#                        token_to_id[\"<EOS>\"],\n",
        "#                        token_to_id[\"awesome\"],\n",
        "#                        token_to_id[\"<EOS>\"]],\n",
        "\n",
        "#                       [token_to_id[\"is\"],\n",
        "#                        token_to_id[\"what\"],\n",
        "#                        token_to_id[\"<EOS>\"],\n",
        "#                        token_to_id[\"awesome\"],\n",
        "#                        token_to_id[\"<EOS>\"]]])\n",
        "\n",
        "labels = torch.roll(inputs,-1,1)\n",
        "labels[:,-1] = token_to_id[\"<EOS>\"]\n",
        "\n",
        "## Now let's package everything up into a DataLoader...\n",
        "dataset = TensorDataset(inputs, labels)\n",
        "print(labels)\n",
        "dataloader = DataLoader(dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3de72068-c081-4868-9e60-1f5cc6cf45df",
      "metadata": {
        "id": "3de72068-c081-4868-9e60-1f5cc6cf45df"
      },
      "source": [
        "Now that we have created the input and output datasets and the **Dataloader** to train the model, let's start building it."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3cb1335c-7bd0-4c60-80c1-6043ee11accd",
      "metadata": {
        "id": "3cb1335c-7bd0-4c60-80c1-6043ee11accd"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "42efbdf7-b3de-4b12-b2ec-a54651b9df79",
      "metadata": {
        "id": "42efbdf7-b3de-4b12-b2ec-a54651b9df79"
      },
      "source": [
        "<a id=\"position\"></a>\n",
        "# Position Encoding\n",
        "\n",
        "Position Encoding helps the transformer keep track of the order of the words in the input and the output. For example, in the picture below, we see that the two phrases **Squatch eats pizza** and **Pizza eats Squatch** both have the exact same words, but, due to differences in the word order, have very different meanings. Thus, keeping track of word order is very important.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/squatch_eats_pizza.png?raw=1\" alt=\"Squatch eats pizza is very different from Pizza eats Squatch\" style=\"width: 800px;\">\n",
        "\n",
        "There are a bunch of ways for a transformer to keep track of word order, but one popular method is to use a series of alternating sine and cosine curves (seen below). The number of sine and cosine curves depends on how many numbers, or word embedding values, we use to represent each token. In the context of Transformers, the number of numbers, or word embedding values, we use to represent each token is the **dimension** of the transformer. So, if the transformer's dimension is 2, meaning that it uses 2 numbers to represent each token, then we only need one sine and one cosine to generate two position encoding values.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/pos_encoding_1.png?raw=1\" alt=\"Sine and cosine squiggles for position encoding\" style=\"width: 800px;\">\n",
        "\n",
        "In contrast, as we see in the illustration below, if the transformer's dimension is 4, then we'll need 2 sine curves alternating with 2 cosine curves, for a total of 4 curves.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/pos_encoding_2.png?raw=1\" alt=\"More sine and cosine squiggles for position encoding\" style=\"width: 800px;\">\n",
        "\n",
        "As we see in the illustration above, the additional pair of sine and cosine curves have a wider period (they repeat less frequently) than the first pair. Increasing the period for each additional pair of curves ensures that each position is represented by a unique combination of values.\n",
        "\n",
        "**NOTE:** The reason why we are bothering to create a class to do positional encoding, instead of just adding this code directly to the transformer, is that we can easily re-use it in an Encoder-Only Transformer or an Encoder-Decoder Transformer. So, by creating a class that does positional encoding, we can code it once, and then just create instances when and where we need it.\n",
        "\n",
        "**ALSO NOTE:** Since the position encoding values never change, meaning that the first token always uses the same position encoding values regardless of what that token is, we precompute them and save them in a lookup table. This makes adding position encoding values super fast.\n",
        "\n",
        "Now that we understand the ideas that we want to implement in the Position Encoding class, let's code it!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "b8b62789-ee84-49bf-b736-12c1b47b34ae",
      "metadata": {
        "id": "b8b62789-ee84-49bf-b736-12c1b47b34ae"
      },
      "outputs": [],
      "source": [
        "class PositionEncoding(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model=2, max_len=3):\n",
        "        ## d_model = The dimension of the transformer, which is also the number of embedding values per token.\n",
        "        ##           In the transformer I used in the StatQuest: Transformer Neural Networks Clearly Explained!!!\n",
        "        ##           d_model=2, so that's what we'll use as a default for now.\n",
        "        ##           However, in \"Attention Is All You Need\" d_model=512\n",
        "        ## max_len = maximum number of tokens we allow as input.\n",
        "        ##           Since we are precomputing the position encoding values and storing them in a lookup table\n",
        "        ##           we can use d_model and max_len to determine the number of rows and columns in that\n",
        "        ##           lookup table.\n",
        "        ##\n",
        "        ##           In this simple example, we are only using 2 word phrases + <SOS>, so we are using\n",
        "        ##           max_len=3 as the default setting.\n",
        "        ##           However, in The Annotated Transformer, they set the default value for max_len to 5000\n",
        "\n",
        "        super().__init__()\n",
        "        ## We call the super's init because by creating our own __init__() method, we overwrite the one\n",
        "        ## we inherited from nn.Module. So we have to explicity call nn.Module's __init__(), otherwise it\n",
        "        ## won't get initialized. NOTE: If we didn't write our own __init__(), then we would not have\n",
        "        ## to call super().__init__(). Alternatively, if we didn't want to access any of nn.Module's methods,\n",
        "        ## we wouldn't have to call it then either.\n",
        "\n",
        "        ## Now we create a lookup table, pe, of position encoding values and initialize all of them to 0.\n",
        "        ## To do this, we will make a matrix of 0s that has max_len rows and d_model columns.\n",
        "        ## for example...\n",
        "        ## torch.zeros(3, 2)\n",
        "        ## ...returns a matrix of 0s with 3 rows and 2 columns...\n",
        "        ## tensor([[0., 0.],\n",
        "        ##         [0., 0.],\n",
        "        ##         [0., 0.]])\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "\n",
        "        ## Now we create a sequence of numbers for each position that a token can have in the input (or output).\n",
        "        ## For example, if the input tokens where \"I'm happy today!\", then \"I'm\" would get the first\n",
        "        ## position, 0, \"happy\" would get the second position, 1, and \"today!\" would get the third position, 2.\n",
        "        ## NOTE: Since we are going to be doing math with these position indices to create the\n",
        "        ## positional encoding for each one, we need them to be floats rather than ints.\n",
        "        ##\n",
        "        ## NOTE: Two ways to create floats are...\n",
        "        ##\n",
        "        ## torch.arange(start=0, end=3, step=1, dtype=torch.float)\n",
        "        ##\n",
        "        ## ...and...\n",
        "        ##\n",
        "        ## torch.arange(start=0, end=3, step=1).float()\n",
        "        ##\n",
        "        ## ...but the latter is just as clear and requires less typing.\n",
        "        ##\n",
        "        ## Lastly, .unsqueeze(1) converts the single list of numbers that torch.arange creates into a matrix with\n",
        "        ## one row for each index, and all of the indices in a single column. So if \"max_len\" = 3, then we\n",
        "        ## would create a matrix with 3 rows and 1 column like this...\n",
        "        ##\n",
        "        ## torch.arange(start=0, end=3, step=1, dtype=torch.float).unsqueeze(1)\n",
        "        ##\n",
        "        ## ...returns...\n",
        "        ##\n",
        "        ## tensor([[0.],\n",
        "        ##         [1.],\n",
        "        ##         [2.]])\n",
        "        position = torch.arange(start=0, end=max_len, step=1).float().unsqueeze(1)\n",
        "\n",
        "\n",
        "        ## Here is where we start doing the math to determine the y-axis coordinates on the\n",
        "        ## sine and cosine curves.\n",
        "        ##\n",
        "        ## The positional encoding equations used in \"Attention is all you need\" are...\n",
        "        ##\n",
        "        ## PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))\n",
        "        ## PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))\n",
        "        ##\n",
        "        ## ...and we see, within the sin() and cos() functions, we divide \"pos\" by some number that depends\n",
        "        ## on the index (i) and total number of PE values we want per token (d_model).\n",
        "        ##\n",
        "        ## NOTE: When the index, i, is 0 then we are calculating the y-axis coordinates on the **first pair**\n",
        "        ##       of sine and cosine curves. When i=1, then we are calculating the y-axis coordiantes on the\n",
        "        ##       **second pair** of sine and cosine curves. etc. etc.\n",
        "        ##\n",
        "        ## Now, pretty much everyone calculates the term we use to divide \"pos\" by first, and they do it with\n",
        "        ## code that looks like this...\n",
        "        ##\n",
        "        ## div_term = torch.exp(torch.arange(start=0, end=d_model, step=2).float() * -(math.log(10000.0) / d_model))\n",
        "        ##\n",
        "        ## Now, at least to me, it's not obvious that div_term = 1/(10000^(2i/d_model)) for a few reasons:\n",
        "        ##\n",
        "        ##    1) div_term wraps everything in a call to torch.exp()\n",
        "        ##    2) It uses log()\n",
        "        ##    2) The order of the terms is different\n",
        "        ##\n",
        "        ## The reason for these differences is, presumably, trying to prevent underflow (getting too close to 0).\n",
        "        ## So, to show that div_term = 1/(10000^(2i/d_model))...\n",
        "        ##\n",
        "        ## 1) Swap out math.log() for torch.log() (doing this requires converting 10000.0 to a tensor, which is my\n",
        "        ##    guess for why they used math.log() instead of torch.log())...\n",
        "        ##\n",
        "        ## torch.exp(torch.arange(start=0, end=d_model, step=2).float() * -(torch.log(torch.tensor(10000.0)) / d_model))\n",
        "        ##\n",
        "        ## 2) Rearrange the terms...\n",
        "        ##\n",
        "        ## torch.exp(-1 * (torch.log(torch.tensor(10000.0)) * torch.arange(start=0, end=d_model, step=2).float() / d_model))\n",
        "        ##\n",
        "        ## 3) Pull out the -1 with exp(-1 * x) = 1/exp(x)\n",
        "        ##\n",
        "        ## 1/torch.exp(torch.log(torch.tensor(10000.0)) * torch.arange(start=0, end=d_model, step=2).float() / d_model)\n",
        "        ##\n",
        "        ## 4) Use exp(a * b) = exp(a)^b to pull out the 2i/d_model term...\n",
        "        ##\n",
        "        ## 1/torch.exp(torch.log(torch.tensor(10000.0)))^(torch.arange(start=0, end=d_model, step=2).float() / d_model)\n",
        "        ##\n",
        "        ## 5) Use exp(log(x)) = x to get the original form of the denominator...\n",
        "        ##\n",
        "        ## 1/(torch.tensor(10000.0)^(torch.arange(start=0, end=d_model, step=2).float() / d_model))\n",
        "        ##\n",
        "        ## 6) Bam.\n",
        "        ##\n",
        "        ## So, that being said, I don't think underflow is actually that big an issue. In fact, some coder at Hugging Face\n",
        "        ## also doesn't think so, and their code for positional encoding in DistilBERT (a streamlined version of BERT, which\n",
        "        ## is a transformer model)\n",
        "        ## calculates the values directly - using the form of the equation found in original Attention is all you need\n",
        "        ## manuscript. See...\n",
        "        ## https://github.com/huggingface/transformers/blob/455c6390938a5c737fa63e78396cedae41e4e87e/src/transformers/modeling_distilbert.py#L53\n",
        "        ## So I think we can simplify the code, but I'm also writing all these comments to show that it is equivalent to what\n",
        "        ## you'll see in the wild...\n",
        "        div_term = 1/torch.tensor(10000.0)**(torch.arange(start=0, end=d_model, step=2).float() / d_model)\n",
        "        ## ALSO NOTE: Setting step=2 results in the same sequence numbers that we would get if we multiplied i by 2.\n",
        "        ##            So we can save ourselves a little math by just setting step=2.\n",
        "\n",
        "        ## Now we calculate the actual positional encoding values. Remember 'pe' was initialized as a matrix of 0s\n",
        "        ## with max_len (max number of input tokens) rows and d_model (number of embedding values per token) columns.\n",
        "        pe[:, 0::2] = torch.sin(position * div_term) ## every other column, starting with the 1st, has sin() values\n",
        "        pe[:, 1::2] = torch.cos(position * div_term) ## every other column, starting with the 2nd, has cos() values\n",
        "        ## NOTE: If the notation for indexing 'pe[]' looks cryptic to you, read on...\n",
        "        ##\n",
        "        ## First, let's look at the general indexing notation:\n",
        "        ##\n",
        "        ## For each row or column in matrix we can select elements in that\n",
        "        ## row or column with the following indexs...\n",
        "        ##\n",
        "        ## i:j:k = select elements between i and j with stepsize = k.\n",
        "        ##\n",
        "        ## ...where...\n",
        "        ##\n",
        "        ## i defaults to 0\n",
        "        ## j defaults to the number of elements in the row, column or whatever.\n",
        "        ## k defaults to 1\n",
        "        ##\n",
        "        ## Now that we have looked at the general notation, let's look at specific\n",
        "        ## examples so that we can understand it.\n",
        "        ##\n",
        "        ## We'll start with: pe[:, 0::2]\n",
        "        ##\n",
        "        ## The stuff that comes before the comma (in this case ':') refers to the rows we want to select.\n",
        "        ## The ':' before the comma means \"select all rows\" because we are not providing specific\n",
        "        ## values for i, j and k and, instead, just using the default values.\n",
        "        ##\n",
        "        ## The stuff after the comma refers to the columns we want to select.\n",
        "        ## In this case, we have '0::2', and that means we start with\n",
        "        ## the first column (column =  0) and go to the end (using the default value for j)\n",
        "        ## and we set the stepsize to 2, which means we skip every other column.\n",
        "        ##\n",
        "        ## Now to understand pe[:, 1::2]\n",
        "        ##\n",
        "        ## Again, the stuff before the comma refers to the rows, and, just like before\n",
        "        ## we use default values for i,j and k, so we select all rows.\n",
        "        ##\n",
        "        ## The stuff that comes after the comma refers to the columns.\n",
        "        ## In this case, we start with the 2nd column (column = 1), and go to the end\n",
        "        ## (using the default value for 'j') and we set the stepsize to 2, which\n",
        "        ## means we skip every other column.\n",
        "        ##\n",
        "        ## NOTE: using this ':' based notation is called \"indexing\" and also called \"slicing\"\n",
        "\n",
        "        ## Now we \"register 'pe'.\n",
        "        self.register_buffer('pe', pe) ## \"register_buffer()\" ensures that\n",
        "                                       ## 'pe' will be moved to wherever the model gets\n",
        "                                       ## moved to. So if the model is moved to a GPU, then,\n",
        "                                       ## even though we don't need to optimize 'pe', it will\n",
        "                                       ## also be moved to that GPU. This, in turn, means\n",
        "                                       ## that accessing 'pe' will be relatively fast copared\n",
        "                                       ## to having a GPU have to get the data from a CPU.\n",
        "\n",
        "    ## Because this class, PositionEncoding, inherits from nn.Module, the forward() method\n",
        "    ## is called by default when we use a PositionEncoding() object.\n",
        "    ## In other words, after we create a PositionEncoding() object, pe = PositionEncoding(),\n",
        "    ## then pe(word_embeddings) will call forward() and so this is where\n",
        "    ## we will add the position encoding values to the word embedding values\n",
        "    def forward(self, word_embeddings):\n",
        "\n",
        "        return word_embeddings + self.pe[:word_embeddings.size(0), :] ## word_embeddings.size(0) = number of embeddings\n",
        "                                                                      ## NOTE: That second ':' is optional and\n",
        "                                                                      ## we could re-write it like this:\n",
        "                                                                      ## self.pe[:word_embeddings.size(0)]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f4fe202d-6995-4322-ad5c-119ef7dc4b28",
      "metadata": {
        "id": "f4fe202d-6995-4322-ad5c-119ef7dc4b28"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c20eccdc-3188-4c6d-94da-d28db057dc24",
      "metadata": {
        "id": "c20eccdc-3188-4c6d-94da-d28db057dc24"
      },
      "source": [
        "<a id=\"attention\"></a>\n",
        "# Attention\n",
        "We're going to code an `Attention` class to do all of the types of attention that a transformer might need: **Self-Attention**, **Masked Self-Attention** (which is used by the Decoder during training), and **Encoder-Decoder Attention**.\n",
        "\n",
        "**Self-Attention** is a type of attention used in Encoder-Decoder and Encoder-Only transformers. It allows every word in a phrase to define a relationship with any other word in the phrase, regardless of the order of the words. In other words, if the phrase is **The pizza came out of the oven and it tasted good!**, then the word **it** can define it's relationship with every word in that phrase, including words that came after it, like **tasted** and **good**, as illustrated by the blue arrows in the figure below.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/self_attention_1.png?raw=1\" alt=\"An illustration of self-attention\" style=\"width: 800px;\">\n",
        "\n",
        "**Masked Self-Attention** is used by Encoder-Decoder and Decoder-Only transformers and it allows each word in a phrase to define a relationship with itself and the words that came before it. In other words, **Masked Self-Attention** prevents the transformer from \"looking ahead\". This is illustrated below where the word **it** can define relationships with itself and everything that came earlier in the input. In Encoder-Decoder transformers, **Masked Self-Attention** is used during training, when we know what the output should be, but we still force the decoder to generate it one token at a time, thus, limiting attention to only output words that came earlier. In contrast, Decoder-Only transformers use **Masked Self-Attention** all the time, on the input and the output, during training and inference. Thus, even though the Decoder-Only transformer can see all of the input during training and inference, it still only allows the attention values for each word to depend on words that came before it.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/masked_attention_1.png?raw=1\" alt=\"An illustration of Masked Self-Attention\" style=\"width: 800px;\">\n",
        "\n",
        "**Encoder-Decoder Attention** is only used in Encoder-Decoder transformers, where there is a distinct separation of the part of the transformer that processes the input (the encoder) from the part that generates the output (the decoder). **Encoder-Decoder Attention** lets each word in the output (in the decoder) define relationships with all the words in the input (in the encoder), as illustrated in the figure below.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/enc_dec_attention_1.png?raw=1\" alt=\"An illustration of Encoder-Decoder Attention\" style=\"width: 800px;\">\n",
        "\n",
        "Now that we have a general sense of the three types of attention used in transformers, we can talk about how it's calculated."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "749d3c97-616d-4af1-a463-960fc36fdcf6",
      "metadata": {
        "id": "749d3c97-616d-4af1-a463-960fc36fdcf6"
      },
      "source": [
        "First, the general equations for the different types of attention are almost identical as seen in the figure below. In the equations, **Q** is for the **Query** matrix, **K** is for the **Key** matrix and **V** is for the **Value** matrix. On the left, we have the equation for Self-Attention and Encoder-Decoder Attention. The differences in these types of attention are not from the equation we use, but from how the **Q**, **K**, and **V** matrices are computed. On the right, we see the equation for Masked Self-Attention and the only difference it has from the equation on the left is the addition of a **Mask** matrix, **M**, that prevents words that come after a specific **Query** from being included in the final attention scores.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_equations.png?raw=1\" alt=\"Equations for computing attention\" style=\"width: 800px;\">\n",
        "\n",
        "**NOTE:** Since both equations are very similar, we'll go through one example and point out the key differences when we get to them.\n",
        "\n",
        "First, given word embedding values for each word/token in the input phrase **\\<SOS> let's go** in matrix form, we multiply them by matrices of weights to create **Queries**, **Keys**, and **Values**\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_compute_q.png?raw=1\" alt=\"Computing the Q matrix\" style=\"width: 800px;\">\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_compute_k.png?raw=1\" alt=\"Computing the K matrix\" style=\"width: 800px;\">\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_compute_k.png?raw=1\" alt=\"Computing the V matrix\" style=\"width: 800px;\">\n",
        "\n",
        "We then multiply the **Queries** by the transpose of the **Keys** so that the query for each word calculates a similarity value with the keys for all of the words. **NOTE:** As seen in the illustration below, Masked Self-Attention calculates the values for all **Query/Key** pairs, but, ultimately, ignores values for when a token's **Query** comes before other token's **Keys**. For example, if the **Query** is for the first token **\\<SOS>**, then Masked Self-Attention will ignore the values calculated with **Keys** for **Let's** and **go** because those tokens come after **\\<SOS>**.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_q_times_kt.png?raw=1\" alt=\"Calculating the similarities between Queries and Keys\" style=\"width: 800px;\">\n",
        "\n",
        "The next step is to scale the similarity scores by the square root of the number of columns in the **Key** matrix, which represents the number of values used to represent each token. In this case, we scale by the square root of 2.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_scaling_scores.png?raw=1\" alt=\"Scaling the similarities\" style=\"width: 800px;\">\n",
        "\n",
        "Now, if we were doing Masked Self-Attention, we would mask out the values we want to ignore by adding -infinity to them, as seen below. This step is the only difference between Self-Attention and Masked Self-Attention.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_masking.png?raw=1\" alt=\"Masking out scaled similarities for Masked Self-Attention\" style=\"width: 800px;\">\n",
        "\n",
        "The next step is to apply the **SoftMax()** function to each row in the scaled similarities. We'll do this first for the Self-Attention without a mask (below)...\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_softmax.png?raw=1\" alt=\"Applying the SoftMax() function to each row in the scaled similarity matrix\" style=\"width: 800px;\">\n",
        "\n",
        "...and we'll also do it for Masked Self-Attention (below).\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_softmax_masked.png?raw=1\" alt=\"Applying the SoftMax() function to each row in the masked scaled similarity matrix\" style=\"width: 800px;\">\n",
        "\n",
        "The `SoftMax()` function gives us percentages that the **Values** for each token should contribute to the attention score for a specific token. Thus, we can get the final attention scores by multiplying the percentages with the **Values** in matrix **V**. First, we'll do this with the unmasked percentages...\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_final_scores.png?raw=1\" alt=\"Calculating the final attention scores\" style=\"width: 800px;\">\n",
        "\n",
        "...and then we'll calculate the final Masked Self-Attention scores.\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/attention_final_scores_masked.png?raw=1\" alt=\"Calculating the final attention scores\" style=\"width: 800px;\">\n",
        "\n",
        "# BAM!\n",
        "\n",
        "Now that we know how to calculate the different types of attention, let's code the `Attention()` class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e3392130-cd25-4000-97bb-9612764c83a8",
      "metadata": {
        "id": "e3392130-cd25-4000-97bb-9612764c83a8"
      },
      "outputs": [],
      "source": [
        "class Attention(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model=2):\n",
        "        ## d_model = the number of embedding values per token.\n",
        "        ##           In the transformer I used in the StatQuest: Transformer Neural Networks Clearly Explained!!!\n",
        "        ##           d_model=2, so that's what we'll use as a default for now.\n",
        "        ##           However, in \"Attention Is All You Need\" d_model=512\n",
        "\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        self.d_model=d_model\n",
        "\n",
        "        ## Initialize the Weights (W) that we'll use to create the\n",
        "        ## query (q), key (k) and value (v) numbers for each token\n",
        "        ## NOTE: Most implementations that I looked at include the bias terms\n",
        "        ##       but I didn't use them in my video (since they are not in the\n",
        "        ##       original Attention is All You Need paper).\n",
        "        self.W_q = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "        self.W_k = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "        self.W_v = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "\n",
        "        ## NOTE: In this simple example, we are not training on the data in \"batches\"\n",
        "        ##       However, by defining variables for row_dim and col_dim, we could\n",
        "        ##       allow for batches by setting row_dim to 1 and col_com to 2.\n",
        "        self.row_dim = 0\n",
        "        self.col_dim = 1\n",
        "\n",
        "\n",
        "    def forward(self, encodings_for_q, encodings_for_k, encodings_for_v, mask=None):\n",
        "        ## Create the query, key and values using the encodings\n",
        "        ## associated with each token (token encodings)\n",
        "        ##\n",
        "        ## NOTE: For Encoder-Decoder Attention, the encodings for q come from\n",
        "        ##       the decoder and the encodings for k and v come from the output\n",
        "        ##       from the encoder.\n",
        "        ##       In all of the other types of attention, the encodings all\n",
        "        ##       come from the same source.\n",
        "        q = self.W_q(encodings_for_q)\n",
        "        k = self.W_k(encodings_for_k)\n",
        "        v = self.W_v(encodings_for_v)\n",
        "\n",
        "        ## Compute attention scores\n",
        "        ## the equation is (q * k^T)/sqrt(d_model)\n",
        "        ## NOTE: It seems most people use \"reverse indexing\" for the dimensions when transposing k\n",
        "        ##       k.transpose(dim0, dim1) will transpose k by swapping dim0 and dim1\n",
        "        ##       In standard matrix notation, we would want to swap rows (dim=0) with columns (dim=1)\n",
        "        ##       If we have 3 dimensions, because of batching, and the batch was the first dimension\n",
        "        ##       And thus dims are defined batch = 0, rows = 1, columns = 2\n",
        "        ##       then dim0=-2 = 3 - 2 = 1. dim1=-1 = 3 - 1 = 2.\n",
        "        ##       Alternatively, we could put the batches in dim 3, and thus, dim 0 would still be rows\n",
        "        ##       and dim 1 would still be columns. I'm not sure why batches are put in dim 0...\n",
        "        ##\n",
        "        ##       Likewise, the q.size(-1) uses negative indexing to rever to the number of columns in the query\n",
        "        ##       which tells us d_model. Alternatively, we could ust q.size(2) if we have batches in the first\n",
        "        ##       dimension or q.size(1) if we have batches in the 3rd dimension.\n",
        "        ##\n",
        "        ##       Since there are a bunch of ways to index things, I think the best thing to do is use\n",
        "        ##       variables \"row_dim\" and \"col_dim\" instead of numbers...\n",
        "        sims = torch.matmul(q, k.transpose(dim0=self.row_dim, dim1=self.col_dim))\n",
        "\n",
        "        scaled_sims = sims / torch.tensor(k.size(self.col_dim)**0.5)\n",
        "\n",
        "        if mask is not None:\n",
        "            ## Here we are masking out things we don't want to pay attention to,\n",
        "            ## like tokens that come after the current token.\n",
        "            ## We can also use masking to block out the <PAD> token,\n",
        "            ## which is used when we have a batch of inputs sequences\n",
        "            ## and they are not all the exact same length. Because the batch is passed\n",
        "            ## in as a matrix, each input sequence has to have the same length, so we\n",
        "            ## add <PAD> to the shorter sequences so that they are all as long ast the\n",
        "            ## longest sequence.\n",
        "            ##\n",
        "            ## We replace <PAD>, or tokens that come after the current token\n",
        "            ## with a very large negative number so that the SoftMax() function\n",
        "            ## will give all masked elements an output value (or \"probability\") of 0.\n",
        "            scaled_sims = scaled_sims.masked_fill(mask=mask, value=-1e9) # I've also seen -1e20 and -9e15 used in masking\n",
        "\n",
        "        ## Apply softmax to determine what percent of each token's value to\n",
        "        ## use in the final attention values.\n",
        "        attention_percents = F.softmax(scaled_sims, dim=self.col_dim)\n",
        "\n",
        "        ## Scale the values by their associated percentages and add them up.\n",
        "        attention_scores = torch.matmul(attention_percents, v)\n",
        "\n",
        "        return attention_scores"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d92c2f59-b64e-4dbc-9a33-74391ff84683",
      "metadata": {
        "id": "d92c2f59-b64e-4dbc-9a33-74391ff84683"
      },
      "source": [
        "# BAM!\n",
        "\n",
        "Now that we have coded the `Attention()` class, we can build a Decoder-Only Transformer."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "291359b2-d509-43d3-aacc-648721690246",
      "metadata": {
        "id": "291359b2-d509-43d3-aacc-648721690246"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "791ffd7e-6671-4153-8d47-4c3b74e61341",
      "metadata": {
        "id": "791ffd7e-6671-4153-8d47-4c3b74e61341"
      },
      "source": [
        "<a id=\"decoder\"></a>\n",
        "# The Decoder-Only Transformer"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "abf2d3f7-41cd-4645-a105-6b9a924df22d",
      "metadata": {
        "id": "abf2d3f7-41cd-4645-a105-6b9a924df22d"
      },
      "source": [
        "\n",
        "\n",
        "<img src=\"https://github.com/StatQuest/decoder_transformer_from_scratch/blob/main/images/dec_transformer.png?raw=1\" alt=\"A diagram of an decoder\" style=\"width: 800px;\">\n",
        "\n",
        "A Decoder-Only Transformer simply brings together...\n",
        "\n",
        "- Word Embedding\n",
        "- Position Encoding\n",
        "- Masked Self-Attention\n",
        "- Residual Connections\n",
        "- A fully connected layer\n",
        "- SoftMax - However, the loss function we are using `nn.CrossEntropyLoss()`, applies the SoftMax for us."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1f286589-e933-46d2-aeda-600c74799357",
      "metadata": {
        "id": "1f286589-e933-46d2-aeda-600c74799357"
      },
      "outputs": [],
      "source": [
        "class DecoderOnlyTransformer(L.LightningModule):\n",
        "\n",
        "    def __init__(self, num_tokens=4, d_model=2, max_len=3):\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        ## We are set the seed so that you can get the same results as me.\n",
        "        L.seed_everything(seed=42)\n",
        "\n",
        "\n",
        "        ## NOTE: In this simple example, we are just using a \"single layer\" decoder.\n",
        "        ##       If we wanted to have multiple layers of decoder, then we would\n",
        "        ##       take the output of one decoder module and use it as input to\n",
        "        ##       the next module.\n",
        "\n",
        "        self.we = nn.Embedding(num_embeddings=num_tokens,\n",
        "                               embedding_dim=d_model)\n",
        "\n",
        "        self.pe = PositionEncoding(d_model=d_model,\n",
        "                                   max_len=max_len)\n",
        "\n",
        "        self.self_attention = Attention(d_model=d_model)\n",
        "        ## NOTE: In this simple example, we are not doing multi-head attention\n",
        "        ## If we wanted to do multi-head attention, we could\n",
        "        ## initailize more Attention objects like this...\n",
        "        ##\n",
        "        ## self.self_attention_2 = Attention(d_model=d_model)\n",
        "        ## self.self_attention_3 = Attention(d_model=d_model)\n",
        "        ##\n",
        "        ## If d_model=2, then using 3 self_attention objects would\n",
        "        ## result in d_model*3 = 6 self-attention values per token,\n",
        "        ## so we would need to initialize\n",
        "        ## a fully connected layer to reduce the dimension of the\n",
        "        ## self attention values back down to d_model like this:\n",
        "        ##\n",
        "        ## self.reduce_attention_dim = nn.Linear(in_features=(num_attention_heads*d_model), out_features=d_model)\n",
        "\n",
        "        self.fc_layer = nn.Linear(in_features=d_model, out_features=num_tokens)\n",
        "\n",
        "        self.loss = nn.CrossEntropyLoss()\n",
        "\n",
        "        self.row_dim = 0\n",
        "        self.col_dim = 1\n",
        "\n",
        "\n",
        "    def forward(self, token_ids):\n",
        "\n",
        "        word_embeddings = self.we(token_ids)\n",
        "        position_encoded = self.pe(word_embeddings)\n",
        "\n",
        "        ## For the decoder-only transformer, we need to use \"masked self-attention\" so that\n",
        "        ## when we are training we can't cheat and look ahead at\n",
        "        ## what words come after the current word.\n",
        "        ## To create the mask we are creating a matrix where the lower triangle\n",
        "        ## is filled with 0, and everything above the diagonal is filled with 0s.\n",
        "        mask = torch.tril(torch.ones((token_ids.size(dim=0), token_ids.size(dim=0)), device=self.device))\n",
        "        ## NOTE: The device=self.device is needed because we are creating a new\n",
        "        ##       tensor, mask, in the forward() method, which, by default, goes\n",
        "        ##       to the CPU. If all we have is a CPU, then we don't need it, but\n",
        "        ##       if we want to train on a GPU, we need to make sure mask goes\n",
        "        ##       there too. Using self.device allows us tyo not worry about whether\n",
        "        ##       or not we are using a GPU or CPU or whatever, it will make sure\n",
        "        ##       mask is where it needs to go.\n",
        "\n",
        "        ## We then replace the 0s above the digaonal, which represent the words\n",
        "        ## we want to be masked out, with \"True\", and replace the 1s in the lower\n",
        "        ## triangle, which represent the words we want to include when we calcualte\n",
        "        ## self-attention for a specific word in the output, with \"False\".\n",
        "        mask = mask == 0\n",
        "\n",
        "        self_attention_values = self.self_attention(position_encoded,\n",
        "                                                    position_encoded,\n",
        "                                                    position_encoded,\n",
        "                                                    mask=mask)\n",
        "        ## NOTE: If we were doing multi-head attention, we would\n",
        "        ## calculate the self-attention values with the other attention objects\n",
        "        ## like this...\n",
        "        ##\n",
        "        ## self_attention_values_2 = self.self_attention_2(...)\n",
        "        ## self_attention_values 3 = self.self_attention_3(...)\n",
        "        ##\n",
        "        ## ...then we would concatenate all the self attention values...\n",
        "        ##\n",
        "        ## all_self_attention_values = torch.cat(self_attention_values_1, ...)\n",
        "        ##\n",
        "        ## ...and then run them through reduce_dim to get back to d_model values per token\n",
        "        ##\n",
        "        ## final_self_attention_values = self.reduce_attention_dim(all_self_attention_values)\n",
        "\n",
        "        residual_connection_values = position_encoded + self_attention_values\n",
        "\n",
        "        fc_layer_output = self.fc_layer(residual_connection_values)\n",
        "\n",
        "        return fc_layer_output\n",
        "\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        ## configure_optimizers() simply passes the parameters we want to\n",
        "        ## optimize to the optimzes and sets the learning rate\n",
        "        return Adam(self.parameters(), lr=0.1)\n",
        "\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        ## training_step() is called by Lightning trainer when\n",
        "        ## we want to train the model.\n",
        "        input_tokens, labels = batch # collect input\n",
        "        output = self.forward(input_tokens[0])\n",
        "        loss = self.loss(output, labels[0])\n",
        "\n",
        "        return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dae5288f-80dc-4ae9-96d5-55b951feeec7",
      "metadata": {
        "id": "dae5288f-80dc-4ae9-96d5-55b951feeec7"
      },
      "source": [
        "# BAM!\n",
        "\n",
        "Now that we have coded up the `DecoderOnlyTransformer()` class, let's see if it works correctly without training. You never know, it might just work!\n",
        "\n",
        "To use the transformer, we run an input phrase, either **what is statquest \\<EOS>** or **statquest is what \\<EOS>**, through the transformer to get the next predicted token. If the next predicted token is not **\\<EOS>**, then we add the predicted token to the input tokens and run that through the transformer and repeat until we get the **\\<EOS>** token or reach the maximum sequence length."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0c7d3c9f-17c0-41e7-b654-c03540713b90",
      "metadata": {
        "id": "0c7d3c9f-17c0-41e7-b654-c03540713b90",
        "outputId": "f44b6d2e-2565-43c3-9eea-75484373cbc7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Global seed set to 42\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted Tokens:\n",
            "\n",
            "\t <EOS>\n"
          ]
        }
      ],
      "source": [
        "## First, create a model from DecoderOnlyTransformer()\n",
        "model = DecoderOnlyTransformer(num_tokens=len(token_to_id), d_model=2, max_len=6)\n",
        "\n",
        "## Now create the input for the transformer...\n",
        "model_input = torch.tensor([token_to_id[\"what\"],\n",
        "                            token_to_id[\"is\"],\n",
        "                            token_to_id[\"statquest\"],\n",
        "                            token_to_id[\"<EOS>\"]])\n",
        "input_length = model_input.size(dim=0)\n",
        "\n",
        "## Now get get predictions from the model\n",
        "predictions = model(model_input)\n",
        "## NOTE: \"predictions\" is the output from the fully connected layer,\n",
        "##      not a softmax() function. We could, if we wanted to,\n",
        "##      Run \"predictions\" through a softmax() function, but\n",
        "##      since we're going to select the item with the largest value\n",
        "##      we can just use argmax instead...\n",
        "## ALSO NOTE: \"predictions\" is a matrix, with one row of predicted values\n",
        "##      per input token. Since we only want the prediction from the\n",
        "##      last row (the most recent prediction) we use reverse index for the\n",
        "##      row, -1.\n",
        "predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "## We'll store predicted_id in an array, predicted_ids, that\n",
        "## we'll add to each time we predict a new output token.\n",
        "predicted_ids = predicted_id\n",
        "\n",
        "## Now use a loop to predict output tokens until we get an\n",
        "## <EOS> token.\n",
        "max_length = 6\n",
        "for i in range(input_length, max_length):\n",
        "    if (predicted_id == token_to_id[\"<EOS>\"]): # if the prediction is <EOS>, then we are done\n",
        "        break\n",
        "\n",
        "    model_input = torch.cat((model_input, predicted_id))\n",
        "\n",
        "    predictions = model(model_input)\n",
        "    predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "    predicted_ids = torch.cat((predicted_ids, predicted_id))\n",
        "\n",
        "## Now printout the predicted output phrase.\n",
        "print(\"Predicted Tokens:\\n\")\n",
        "for id in predicted_ids:\n",
        "    print(\"\\t\", id_to_token[id.item()])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5c451f90-3382-4931-92fc-5cd3670be7e3",
      "metadata": {
        "id": "5c451f90-3382-4931-92fc-5cd3670be7e3"
      },
      "source": [
        "And, without training, the transformer predicts **\\<EOS>**, but we wanted it to predict **awesome \\<EOS>** So, since the transformer didn't correctly respond to the prompt, we'll have to train it."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "39d51f3f-02c7-4d30-8e18-43be37996454",
      "metadata": {
        "id": "39d51f3f-02c7-4d30-8e18-43be37996454"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "982ca23f-0c46-4aa9-aebf-32e5bcedece6",
      "metadata": {
        "id": "982ca23f-0c46-4aa9-aebf-32e5bcedece6"
      },
      "source": [
        "<a id=\"train\"></a>\n",
        "# Train the Decoder-Only Transformer!!!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ddc78dd9-ca22-4c6a-bd0f-26540716620e",
      "metadata": {
        "id": "ddc78dd9-ca22-4c6a-bd0f-26540716620e"
      },
      "source": [
        "To train a decoder-only transformer, we simply create a Lightning `Trainer()` and train the transformer with the `dataloader` that we created earlier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a96bd93c-31f4-4ef8-8c4f-fd33c3c20344",
      "metadata": {
        "tags": [],
        "id": "a96bd93c-31f4-4ef8-8c4f-fd33c3c20344",
        "outputId": "41c762c1-2444-42ea-cb31-61374350fc77",
        "colab": {
          "referenced_widgets": [
            "bfab0a7d7a6349759e7764f2a77c51eb"
          ]
        }
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "GPU available: False, used: False\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "HPU available: False, using: 0 HPUs\n",
            "\n",
            "  | Name           | Type             | Params\n",
            "----------------------------------------------------\n",
            "0 | we             | Embedding        | 10    \n",
            "1 | pe             | PositionEncoding | 0     \n",
            "2 | self_attention | Attention        | 12    \n",
            "3 | fc_layer       | Linear           | 15    \n",
            "4 | loss           | CrossEntropyLoss | 0     \n",
            "----------------------------------------------------\n",
            "37        Trainable params\n",
            "0         Non-trainable params\n",
            "37        Total params\n",
            "0.000     Total estimated model params size (MB)\n",
            "/Users/joshstarmer/opt/anaconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:430: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 6 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
            "  rank_zero_warn(\n",
            "/Users/joshstarmer/opt/anaconda3/lib/python3.9/site-packages/lightning/pytorch/loops/fit_loop.py:280: PossibleUserWarning: The number of training batches (2) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
            "  rank_zero_warn(\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "bfab0a7d7a6349759e7764f2a77c51eb",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Training: 0it [00:00, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "`Trainer.fit` stopped: `max_epochs=30` reached.\n"
          ]
        }
      ],
      "source": [
        "trainer = L.Trainer(max_epochs=30)\n",
        "trainer.fit(model, train_dataloaders=dataloader)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0acf9d1f-ac29-4316-88bb-5c6d0f3e04b6",
      "metadata": {
        "id": "0acf9d1f-ac29-4316-88bb-5c6d0f3e04b6"
      },
      "source": [
        "# Double BAM!!!\n",
        "\n",
        "Now that we've trained the transformer, let's use it!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1b80a32f-ded6-4299-bf89-12c6bbd8ced2",
      "metadata": {
        "id": "1b80a32f-ded6-4299-bf89-12c6bbd8ced2"
      },
      "source": [
        "----"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c040285c-c6de-4aac-b5b9-925e84f5e83f",
      "metadata": {
        "id": "c040285c-c6de-4aac-b5b9-925e84f5e83f"
      },
      "source": [
        "<a id=\"use\"></a>\n",
        "# Use the Trained Transformer!!!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0666f428-eb0d-4c43-91d5-48c67cdcc5f8",
      "metadata": {
        "id": "0666f428-eb0d-4c43-91d5-48c67cdcc5f8"
      },
      "source": [
        "To use the transformer that we just trained, we just repeat what we did earlier, only this time we use the trained transformer instead of an untrained transformer. First, we'll see if it correctly responds to the prompt **What is StatQuest?**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "717eaa4e-2249-46a1-b01f-9ed76beed0de",
      "metadata": {
        "id": "717eaa4e-2249-46a1-b01f-9ed76beed0de",
        "outputId": "f9374e22-3500-4faf-84e8-d5a887a5e30e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted Tokens:\n",
            "\n",
            "\t awesome\n",
            "\t <EOS>\n"
          ]
        }
      ],
      "source": [
        "model_input = torch.tensor([token_to_id[\"what\"],\n",
        "                            token_to_id[\"is\"],\n",
        "                            token_to_id[\"statquest\"],\n",
        "                            token_to_id[\"<EOS>\"]])\n",
        "input_length = model_input.size(dim=0)\n",
        "\n",
        "predictions = model(model_input)\n",
        "predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "predicted_ids = predicted_id\n",
        "\n",
        "for i in range(input_length, max_length):\n",
        "    if (predicted_id == token_to_id[\"<EOS>\"]): # if the prediction is <EOS>, then we are done\n",
        "        break\n",
        "\n",
        "    model_input = torch.cat((model_input, predicted_id))\n",
        "\n",
        "    predictions = model(model_input)\n",
        "    predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "    predicted_ids = torch.cat((predicted_ids, predicted_id))\n",
        "\n",
        "print(\"Predicted Tokens:\\n\")\n",
        "for id in predicted_ids:\n",
        "    print(\"\\t\", id_to_token[id.item()])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "40d52e9f-81dc-4f51-a699-5a19127b4a19",
      "metadata": {
        "id": "40d52e9f-81dc-4f51-a699-5a19127b4a19"
      },
      "source": [
        "Hooray!!! We got the correct output! Now let's see if it correctly responds to the prompt **StatQuest is what?**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "89881ada-c49d-4805-a989-21020d9ca4c9",
      "metadata": {
        "id": "89881ada-c49d-4805-a989-21020d9ca4c9",
        "outputId": "cd5b68b7-29d7-4dbf-fdb3-7f242e910f18"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Predicted Tokens:\n",
            "\n",
            "\t awesome\n",
            "\t <EOS>\n"
          ]
        }
      ],
      "source": [
        "## Now let's ask the other question...\n",
        "model_input = torch.tensor([token_to_id[\"statquest\"],\n",
        "                            token_to_id[\"is\"],\n",
        "                            token_to_id[\"what\"],\n",
        "                            token_to_id[\"<EOS>\"]])\n",
        "input_length = model_input.size(dim=0)\n",
        "\n",
        "predictions = model(model_input)\n",
        "predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "predicted_ids = predicted_id\n",
        "\n",
        "for i in range(input_length, max_length):\n",
        "    if (predicted_id == token_to_id[\"<EOS>\"]): # if the prediction is <EOS>, then we are done\n",
        "        break\n",
        "\n",
        "    model_input = torch.cat((model_input, predicted_id))\n",
        "\n",
        "    predictions = model(model_input)\n",
        "    predicted_id = torch.tensor([torch.argmax(predictions[-1,:])])\n",
        "    predicted_ids = torch.cat((predicted_ids, predicted_id))\n",
        "\n",
        "print(\"Predicted Tokens:\\n\")\n",
        "for id in predicted_ids:\n",
        "    print(\"\\t\", id_to_token[id.item()])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e3a6dda5-f868-410e-8038-b1bd8d4024df",
      "metadata": {
        "id": "e3a6dda5-f868-410e-8038-b1bd8d4024df"
      },
      "source": [
        "And the output for both questions is **awesome \\<EOS>**, which is exactly what we want.\n",
        "\n",
        "# TRIPLE BAM!!!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c5e297c9-9735-44fd-83ef-fd267cba327c",
      "metadata": {
        "id": "c5e297c9-9735-44fd-83ef-fd267cba327c"
      },
      "source": [
        "**NOTE:** With all the comments in the `PositionEncoding()`, `Attention()`, and `DecoderOnlyTransformer` classes, it may seem like we had to write a lot of code to create a Decoder-Only Transformer. Not so. Below we see that all three classes are only a handful of lines of code. Isn't that bonkers? The most state-of-the-art model isn't that big of a deal!\n",
        "\n",
        "```python\n",
        "class PositionEncoding(nn.Module):\n",
        "    \n",
        "    def __init__(self, d_model=2, max_len=3):\n",
        "        \n",
        "        super().__init__()\n",
        "        \n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(start=0, end=max_len, step=1).float().unsqueeze(1)\n",
        "        div_term = 1/torch.tensor(10000.0)**(torch.arange(start=0, end=d_model, step=2).float() / d_model)\n",
        "\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "        \n",
        "    def forward(self, word_embeddings):\n",
        "        \n",
        "        return word_embeddings + self.pe[:word_embeddings.size(0), :]\n",
        "\n",
        "class Attention(nn.Module):\n",
        "    \n",
        "    def __init__(self, d_model=2):\n",
        "        \n",
        "        super().__init__()\n",
        "        \n",
        "        self.W_q = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "        self.W_k = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "        self.W_v = nn.Linear(in_features=d_model, out_features=d_model, bias=False)\n",
        "        \n",
        "        self.row_dim = 0\n",
        "        self.col_dim = 1\n",
        "\n",
        "        \n",
        "    def forward(self, encodings_for_q, encodings_for_k, encodings_for_v, mask=None):\n",
        "\n",
        "        q = self.W_q(encodings_for_q)\n",
        "        k = self.W_k(encodings_for_k)\n",
        "        v = self.W_v(encodings_for_v)\n",
        "\n",
        "        sims = torch.matmul(q, k.transpose(dim0=self.row_dim, dim1=self.col_dim))\n",
        "\n",
        "        scaled_sims = sims / torch.tensor(k.size(self.col_dim)**0.5)\n",
        "\n",
        "        if mask is not None:\n",
        "            scaled_sims = scaled_sims.masked_fill(mask=mask, value=-1e9)\n",
        "        \n",
        "        attention_percents = F.softmax(scaled_sims, dim=self.col_dim)\n",
        "        attention_scores = torch.matmul(attention_percents, v)\n",
        "        \n",
        "        return attention_scores\n",
        "\n",
        "class DecoderOnlyTransformer(L.LightningModule):\n",
        "    \n",
        "    def __init__(self, num_tokens=4, d_model=2, max_len=3):\n",
        "        \n",
        "        super().__init__()\n",
        "        \n",
        "        L.seed_everything(seed=42)\n",
        "        \n",
        "        self.we = nn.Embedding(num_embeddings=num_tokens,\n",
        "                               embedding_dim=d_model)     \n",
        "        self.pe = PositionEncoding(d_model=d_model,\n",
        "                                   max_len=max_len)\n",
        "        self.self_attention = Attention(d_model=d_model)\n",
        "        self.fc_layer = nn.Linear(in_features=d_model, out_features=num_tokens)\n",
        "        \n",
        "        self.loss = nn.CrossEntropyLoss()\n",
        "        \n",
        "        self.row_dim = 0\n",
        "        self.col_dim = 1\n",
        "        \n",
        "        \n",
        "    def forward(self, token_ids):\n",
        "                \n",
        "        word_embeddings = self.we(token_ids)        \n",
        "        position_encoded = self.pe(word_embeddings)\n",
        "        \n",
        "        mask = torch.tril(torch.ones((token_ids.size(dim=0), token_ids.size(dim=0))))\n",
        "        mask = mask == 0\n",
        "        \n",
        "        self_attention_values = self.self_attention(position_encoded,\n",
        "                                                    position_encoded,\n",
        "                                                    position_encoded,\n",
        "                                                    mask=mask)\n",
        "                \n",
        "        residual_connection_values = position_encoded + self_attention_values        \n",
        "        fc_layer_output = self.fc_layer(residual_connection_values)\n",
        "        \n",
        "        return fc_layer_output\n",
        "    \n",
        "  \n",
        "    def configure_optimizers(self):\n",
        "        return Adam(self.parameters(), lr=0.1)\n",
        "    \n",
        "    \n",
        "    def training_step(self, batch, batch_idx):\n",
        "        input_tokens, labels = batch # collect input\n",
        "        output = self.forward(input_tokens[0])\n",
        "        loss = self.loss(output, labels[0])\n",
        "                    \n",
        "        return loss\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3c21bd5a-7a4a-4ec6-adb2-d3fa38718fe2",
      "metadata": {
        "id": "3c21bd5a-7a4a-4ec6-adb2-d3fa38718fe2"
      },
      "source": [
        "# BONUS BAM!!!\n",
        "\n",
        "Now that we can code our own transformer, let's do what people actually do in practice with transformers - **[fine tune a pre-trained LLM](https://lightning.ai/lightning-ai/studios/instruction-finetuning-tinyllama-1-1b-llm)**."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.12"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}