Exercise2a.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/biggojo/Assignments-DL/blob/main/Exercise2a.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### This submission is for:\n",
        "- Abraham Gassama (2285843)\n",
        "- Zhuo Le Lee (2317240)"
      ],
      "metadata": {
        "cell_id": "e6929dd6290a45bbbeba2dbd64adda8b",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 153.171875,
        "id": "fW922_J_pOO-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Exercise 2A - Transformers"
      ],
      "metadata": {
        "cell_id": "9339153b4e8f400894090f4cdec8212f",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 82,
        "id": "xPQt80WhpOPs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this exercise, you'll implement a basic encoder-only Transformer architecture with PyTorch. We will start with building the basic building blocks and then integrate them into a fully-fleged Transformer model. We train the model to solve a POS-Tagging problem (more on that later). In the previous exercise, you implemented your work in numpy. Now, we will switch to PyTorch, which will track the gradients for us and allows us to focus more on the network itself."
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "eef422cc65524297ab7b1220162873b0",
        "tags": [],
        "is_collapsed": false,
        "owner_user_id": "175a6e67-5f66-4c81-960a-2a75fbd3d9af",
        "deepnote_cell_type": "text-cell-p",
        "id": "YDMaj5jWpOPu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can receive up to three points for your implementation of Exercise 2A. Together with Exercise 1, you can get up to six bonus points for the exam."
      ],
      "metadata": {
        "cell_id": "64bcb0b3c2bb43a6b1620d47ecea1de6",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "m6cSfTtzpOPw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Important Notice**: Throughout the notebook, basic structures are provided such as functions and classes without bodies or partial bodies, and variables that you need to assign to. **Don't change the names of functions, variables, and classes - and make sure that you are using them!** You're allowed to introduce helper variables and functions. Occasionally, we use **type annotations** that you should follow. They are not enforced by Python. Whenenver you see an ellipsis `...` you're supposed to insert code."
      ],
      "metadata": {
        "cell_id": "c66a0519697243859f4aa30b0a0e06c0",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 119.5625,
        "id": "-nTKihOkpOPy"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "1ed55a8f3a8d4f8ea7f58929e119f5ed",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "aef22c3f",
        "execution_start": 1657123526340,
        "execution_millis": 40884,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 696,
        "id": "o3VfJAIepOP0",
        "outputId": "9e626129-9024-4baf-f1c5-91e9b09514ea",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "!pip install torchtext torchdata torchmetrics"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: torchtext in /usr/local/lib/python3.7/dist-packages (0.13.1)\n",
            "Requirement already satisfied: torchdata in /usr/local/lib/python3.7/dist-packages (0.4.1)\n",
            "Requirement already satisfied: torchmetrics in /usr/local/lib/python3.7/dist-packages (0.9.3)\n",
            "Requirement already satisfied: torch==1.12.1 in /usr/local/lib/python3.7/dist-packages (from torchtext) (1.12.1+cu113)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchtext) (1.21.6)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torchtext) (4.64.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torchtext) (2.23.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.12.1->torchtext) (4.1.1)\n",
            "Requirement already satisfied: urllib3>=1.25 in /usr/local/lib/python3.7/dist-packages (from torchdata) (1.25.11)\n",
            "Requirement already satisfied: portalocker>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from torchdata) (2.5.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from torchmetrics) (21.3)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->torchmetrics) (3.0.9)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torchtext) (2022.6.15)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->torchtext) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->torchtext) (2.10)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "20b2dbfb9550468d85d409a5a67f64c0",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "2c538688",
        "execution_start": 1657123567236,
        "execution_millis": 810,
        "owner_user_id": "06b28ca6-80fe-4ecd-a509-50438de77bba",
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 149.375,
        "id": "0mY3MFbfpOQQ"
      },
      "source": [
        "import torch\n",
        "from torch import nn, Tensor\n",
        "\"\"\"\n",
        "Shifted this definition to the beginning, since some of the implemented functions at the beginning\n",
        "require this variable.\n",
        "\"\"\"\n",
        "DEVICE = 'cuda' # later replace with 'cuda' for GPU"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's actually start with a few basic functions that we will need throughout the exercise, namely **Softmax** and **ReLu**.\n",
        "\n",
        "$\\text{Softmax}(x_{i}) = \\frac{\\exp(x_i)}{\\sum_j \\exp(x_j)}$\n",
        "\n",
        "$\\text{ReLU}(x) = \\max(0, x)$"
      ],
      "metadata": {
        "cell_id": "925794d055334d71a1874bd0928f2358",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 131.1875,
        "id": "xqp7JL3PpOQT"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "89cb4d8c8ca04395a96eb2deedea1989",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 133,
        "id": "1ZlGwRNGpOQW"
      },
      "source": [
        "def softmax(input: Tensor) -> Tensor:\n",
        "      \n",
        "    \"\"\"\n",
        "    Note 1: This implementation of softmax does the softmax operation only on the last dimension\n",
        "    of the given input. To perform the operation on other axis/dimensions of the input, \n",
        "    use nn.Softmax\n",
        "\n",
        "    Note 2: This implementation assumes that the input tensor has a dimension of 2. This means that \n",
        "    the input has to be reshaped first, if its dimension does not correspond to 2.\n",
        "    \"\"\"\n",
        "    # Clone the input tensor and save as output tensor to be overwritten \n",
        "    output_tensor = torch.clone(input)\n",
        "\n",
        "    # Perform the exponent operation over all elements in the input tensor\n",
        "    exp_input = torch.exp(input)\n",
        "    \n",
        "    # Sum over the columns \n",
        "    sum_exp = torch.sum(exp_input, dim=1)\n",
        "\n",
        "    # Overwrite the output tensor \n",
        "    for i in range(input.size()[0]):\n",
        "      output_tensor[i, :] = exp_input[i, :] / sum_exp[i]\n",
        "    \n",
        "    return output_tensor\n",
        "\n",
        "class relu(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(relu, self).__init__()\n",
        "\n",
        "  def forward(self, input: Tensor) -> Tensor:\n",
        "    return torch.maximum(input.to(DEVICE), torch.zeros(input.size()).to(DEVICE))\n"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sanity checks for the ReLU and softmax function\n",
        "\n",
        "# Create a tensor of random values with 15 columns and 1 row to match the assumption made\n",
        "sample_input = torch.rand(4, 10)\n",
        "\n",
        "# Note: The output tensors are rounded to 4 decimals places to allow for easier comparison\n",
        "# Output from our implementation\n",
        "SM_out1 = torch.round(softmax(sample_input), decimals=4).to(DEVICE)\n",
        "ReLU_ori = relu()\n",
        "ReLU_out1 = torch.round(ReLU_ori(sample_input), decimals=4).to(DEVICE)\n",
        "\n",
        "# Output from PyTorch Implementation\n",
        "SM_pytorch = nn.Softmax(dim=-1)\n",
        "SM_out2 = torch.round(SM_pytorch(sample_input), decimals=4).to(DEVICE)\n",
        "ReLU_pytorch = nn.ReLU()\n",
        "ReLU_out2 = torch.round(ReLU_pytorch(sample_input), decimals=4).to(DEVICE)\n",
        "\n",
        "\n",
        "softmax_compare = torch.eq(SM_out1, SM_out2)\n",
        "relu_compare = torch.eq(ReLU_out1, ReLU_out2)\n",
        "\n",
        "\n",
        "print(softmax_compare)\n",
        "print(relu_compare)\n",
        "\n",
        "# Both functions passed the sanity checks "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nLAXx1DRnW8Q",
        "outputId": "674ee240-4f45-423d-f8c5-05cc2437b02d"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True]],\n",
            "       device='cuda:0')\n",
            "tensor([[True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True],\n",
            "        [True, True, True, True, True, True, True, True, True, True]],\n",
            "       device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Transformer Block"
      ],
      "metadata": {
        "cell_id": "65d60a38ea0d41ca934e8a0f782b1672",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 70,
        "id": "_R6L7FsMpOQZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "A typical transformer block consists of the following \n",
        "- Multi-Head Attention\n",
        "- Layer Normalization\n",
        "- Linear Layer\n",
        "- Residual Connections"
      ],
      "metadata": {
        "cell_id": "090be921347b4f65a3dd0fca98b29c54",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 178.953125,
        "id": "tGlQFefdpOQb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center><img src=\"https://i.imgur.com/ZKgcoe4.png\" alt=\"transformer block visualization\" width=\"200\">"
      ],
      "metadata": {
        "cell_id": "f87bf87fea1a44bcbc2d5c044843e007",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 353,
        "id": "E4l0bYkUpOQf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the next few subsections, we will build these basic building blocks."
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "1c8bd5f72db0415bbca4cdd48b0ecd52",
        "tags": [],
        "is_collapsed": false,
        "deepnote_cell_type": "text-cell-p",
        "id": "Xm8WA_q-pOQj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multi-Head Attention"
      ],
      "metadata": {
        "cell_id": "59b1b38243a84b5c9571ddb814cd1a49",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "kD74-tq4pOQk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Multi-Head Attention concatenates the outputs of several so called **attention heads**.\n",
        "\n",
        "$\\textrm{MHA}(Q,K,V) = \\textrm{Concat}(H_1,...,H_h)$"
      ],
      "metadata": {
        "cell_id": "a2a24844affe40a89f26ef6ca0346276",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 88.78125,
        "id": "HUWCLswzpOQl"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center><img src=\"https://www.tensorflow.org/images/tutorials/transformer/multi_head_attention.png\" width=300>"
      ],
      "metadata": {
        "cell_id": "ef2e9b876ee44b69a993f15856b206e9",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 376.15625,
        "id": "toF2B-c3pOQm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "One attention head consists of linear projections for each of $Q, K$ and $V$ and an attention mechanism called **Scaled Dot-Product Attention**. The attention mechanism scales down the dot products by $\\sqrt{d_k}$.\n",
        "\n",
        "$\\textrm{Attention}(Q,K,V)=\\textrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V$\n",
        "\n",
        "\n",
        "\n",
        "If we assume that $q$ and $v$ are $d_k$-dimensional vectors and its components are independent random variables with mean $0$ and a variance of $d_k$, then their dot product has a mean of $0$ and variance of $d_k$. It is preferred to have a variance of $1$ and that's why we scale them down by $\\sqrt{d_k}$.\n",
        "\n",
        "The dot product $q \\cdot v$ resembles a measure of similarity.\n"
      ],
      "metadata": {
        "cell_id": "a045d2caec3e4020a7bb3fee5680442b",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 233.90625,
        "id": "QUHITGgupOQm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center><img src=\"https://www.tensorflow.org/images/tutorials/transformer/scaled_attention.png\" width=\"350\">"
      ],
      "metadata": {
        "cell_id": "9495e574d00e4d3798368384c6b3825b",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 421.609375,
        "id": "4BoPKh5TpOQo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's start implementing these components. Note that our classes inherit from PyTorch's `nn.Module`. These modules allow us to hold our parameters and easily move them to the GPU (with `.to(...)`). It also let's us define the computation that is performed at every call, in the `forward()` method. For example, when we have an `Attention` module, initialize it like `attention = Attention(...)`, we are able to call it with `attention(Q, K, V)` (it'll execute the `forward` function in an optimized way)."
      ],
      "metadata": {
        "cell_id": "e3ed8097791f458a85f3aabfd0756d28",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 119.5625,
        "id": "KDf7KvfupOQq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Assumptions: \n",
        "\n",
        "\n",
        "1.   Q, K and V are **batches** of matrices, each with shape *(batch_size, seq_length, num_features)*. \n",
        "2.   The attention here is self-attention, this means that the sequence length of Q and K are the same. \n",
        "3.   The number of hidden dimensions for all Q, K and V are the same.\n",
        "4.   For added simplicity, we first assume that the batch size is one. We then introduce batched attention matrix after the simple attention has passed all of the sanity checks\n",
        "\n",
        "### Operations in Attention and their outputs: \n",
        "\n",
        "\n",
        "\n",
        "1.   Multiplying the query (Q) and key (K) arrays in a linear matrix multiplication. This is the attention of this layer and determines how important each element in the key sequences is with regard to the query sequence. \n",
        "\n",
        "> Output 1: *(batch_size, seq_length, seq_length)*\n",
        "\n",
        "2.   Scaling the attention to have variance of 1, regardless of the size of seq_length\n",
        "\n",
        "> Output 2: *(batch_size, seq_length, seq_length)*\n",
        "\n",
        "3.   Normalizing the attention array across the keys dimension with softmax, so that all the attention weights sum to one. (Because we can't pay more than 100% attention!)\n",
        "\n",
        "> Output 3: *(batch_size, seq_length, seq_length)*\n",
        "\n",
        "4.   Multiplying the attention matrix with the value (V) array using matrix multiplication. \n",
        "\n",
        "> Output 4: *(batch_size, seq_length, num_features)*\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "7569pcqgQRC9"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "2fc20cdd2c4c4e95aa1cb044b1258d63",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 187,
        "id": "-jjFTztxpOQs"
      },
      "source": [
        "# Implementation of a single attention head\n",
        "\n",
        "class Attention(nn.Module):\n",
        "    def __init__(self, hidden_n: int):\n",
        "        \"\"\"\n",
        "        hidden_n : Number of hidden dimensions\n",
        "        It is here assumed that the number of hidden dimensions corresponds\n",
        "        to the number of features in each input timestep.\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.hidden_n = hidden_n\n",
        "        \n",
        "        \"\"\"\n",
        "        nn.Linear intializes and defines a transformation matrix, where all but \n",
        "        the last dimension of the output is different than the input \n",
        "        \"\"\"\n",
        "        self.Q_linear = nn.Linear(hidden_n, hidden_n)\n",
        "        self.V_linear = nn.Linear(hidden_n, hidden_n)\n",
        "        self.K_linear = nn.Linear(hidden_n, hidden_n)\n",
        "\n",
        "\n",
        "    def forward(self, Q: Tensor, K: Tensor, V: Tensor, mask: Tensor = None) -> Tensor:\n",
        "        \"\"\"\n",
        "        Note 1: If it was a batched attention, then all of the following matmul \n",
        "        has to be changed to bmm (batched matrix multiplication, supported by\n",
        "        pytorch)\n",
        "\n",
        "        Note 2: The above softmax function would not work here, if the number of \n",
        "        time steps (i.e. sequence length) and the number of features are both\n",
        "        more than one. The above softmax function would have to be revised to \n",
        "        include the number of dimensions.\n",
        "        \n",
        "        **Correction**: The above softmax function would work, since we are applying\n",
        "        softmax only to the hidden dimension, i.e. the last dimension of the given\n",
        "        embedding tensor as input. \n",
        "\n",
        "        Note 3: The dimensions of the transpose has to be double-checked against\n",
        "        the size of the inputs to the attention layer later, when the inputs\n",
        "        are given.\n",
        "\n",
        "        Note 4: A naive implementation of the masking operation is given here, \n",
        "        should the a mask be defined as an input. This might have to be revised\n",
        "        in later steps.`\n",
        "\n",
        "        Args:\n",
        "            Q (:class:`torch.Tensor` [batch_size, output_length, hidden_n]): \n",
        "            Sequence of queries to query the context.\n",
        "            K (:class:`torch.Tensor` [batch_size, query_length, hidden_n]): \n",
        "            Data overwhich to apply the attention mechanism.\n",
        "            V (:class:`torch.Tensor` [batch_size, output_length, hidden_n]): \n",
        "            Sequence of values to be scaled by the attended features\n",
        "\n",
        "            here: output_length == query_length == seq_length \n",
        "\n",
        "        Returns:\n",
        "            :class:`torch.Tensor`:\n",
        "\n",
        "            output (:class:`torch.Tensor` [batch size, output_length, dimensions]):\n",
        "            Tensor containing the attended features.\n",
        "        \"\"\"\n",
        "\n",
        "        # Sanity Checks\n",
        "        assert Q.size()[-1] == K.size()[-1] == V.size()[-1], \"The hidden dimensions of Q, K, and V must match up!\"\n",
        "        assert Q.size()[-1] == self.hidden_n, \"The hidden dimension of the given inputs must match that defined during the init!\"\n",
        "\n",
        "        batch_size, seq_length, hidden_n = Q.size() \n",
        "        \n",
        "       \n",
        "        \"\"\"\n",
        "        Step 1: Transforming Q, K, and V via linear transformation (with bias). This improves the versatility of the \n",
        "        network in adapting the given input arrays to its specific function. \n",
        "\n",
        "        (batch_size * seq_length, hidden_n) * (hidden_n * hidden_n) ->\n",
        "        (batch_size * seq_length, hidden_n)\n",
        "        \"\"\"\n",
        "\n",
        "        Q = Q.reshape(batch_size * seq_length, hidden_n)\n",
        "        K = K.reshape(batch_size * seq_length, hidden_n)\n",
        "        V = V.reshape(batch_size * seq_length, hidden_n)\n",
        "\n",
        "        Q_arr = self.Q_linear(Q).reshape(batch_size, seq_length, hidden_n)\n",
        "        K_arr = self.K_linear(K).reshape(batch_size, seq_length, hidden_n)\n",
        "        V_arr = self.V_linear(V).reshape(batch_size, seq_length, hidden_n)\n",
        "\n",
        "\n",
        "        \"\"\"\n",
        "        weights (:class:`torch.FloatTensor` [batch size, seq_length, seq_length]):\n",
        "        Tensor containing attention weights.\n",
        "        \"\"\"\n",
        "        scores = torch.bmm(Q_arr, K_arr.transpose(-1,1).contiguous()) / (self.hidden_n ** 0.5)\n",
        "        \n",
        "        \"\"\"\n",
        "        For the case of masked attention, a masking operation has to be implemented\n",
        "        \"\"\"\n",
        "        if mask is not None:\n",
        "          mask = mask.unsqueeze(-1)\n",
        "          scores = scores.masked_fill(mask == 0, float(\"-1e20\"))\n",
        "          scores = softmax(scores)\n",
        "\n",
        "\n",
        "        scores = scores.view(batch_size * seq_length, seq_length)\n",
        "        weights = softmax(scores)\n",
        "        weights = weights.view(batch_size, seq_length, seq_length)\n",
        "\n",
        "        \"\"\"\n",
        "        (batch_size, seq_length, seq_length) * (batch_size, seq_length, hidden_n) ->\n",
        "        (batch_size, seq_length, hidden_n)\n",
        "        \"\"\"\n",
        "\n",
        "        output = torch.bmm(weights, V_arr)\n",
        "        return output \n"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "4ad1acddce234ba98f163b58e34c04d1",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 241,
        "id": "fdfZ-iMSpOQt"
      },
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, hidden_n: int, h: int = 2):\n",
        "        \"\"\"\n",
        "        hidden_n: hidden dimension\n",
        "        h: number of heads\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.hidden_n = hidden_n\n",
        "        self.h = h\n",
        "        # self.d_k = hidden_n // h\n",
        "\n",
        "        \"\"\"\n",
        "        Intialize an attention head for each element in h\n",
        "        \"\"\"\n",
        "        self.heads = nn.ModuleList(\n",
        "            [Attention(hidden_n) for x in range(h)]\n",
        "        )\n",
        "        \n",
        "        \"\"\"\n",
        "        After concatenating, the output size would be (batch_size, seq_length, num_features * h)\n",
        "        This can be reverted back to the original output size via a linear layer\n",
        "\n",
        "        Note 1: There are other implementation variations regarding this step.\n",
        "        E.g. we could split the input into n heads, so that each will have the dimensions\n",
        "        (batch_size, seq_length, num_features, num_features / n), with the last dimension denoting\n",
        "        which head each input is subjected to. Each input is then passed through the according\n",
        "        attention head and then concatenated. \n",
        "\n",
        "        This variation was not chosen for simplicity, e.g. self.d_k has to be a factor of \n",
        "        self.hidden_n\n",
        "        \"\"\"\n",
        "        self.out = nn.Linear(h * hidden_n, hidden_n)\n",
        "\n",
        "   \n",
        "    def forward(self, Q: Tensor, K: Tensor, V: Tensor, mask=None):\n",
        "        return self.out(\n",
        "            torch.cat([h(Q,K,V) for h in self.heads], dim=-1)\n",
        "        )"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Layer Normalization"
      ],
      "metadata": {
        "cell_id": "96fa5effdd954c498f1efa8cd0580bc3",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "yUZX_55zpOQu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "From the lecture, remember layer normalization where the values are normalized across the feature dimension, independently for each sample in the batch. For that, first calculate mean and standard-deviation across the feature dimension and then scale them appropriately such that the mean is 0 and the standard deviation is 1. Introduce **two sets of learnable parameters**, one for shifting the mean (addition) and one for scaling the variance (multiplication) the normalized features (i.e., two parameters for each feature). Tip: Use `nn.Parameter` for that."
      ],
      "metadata": {
        "cell_id": "fba0740f863c4787aa23494388944b49",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 141.953125,
        "id": "_Ufh5KKxpOQv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "$y_{\\textrm{norm}}=\\frac{x-\\mu}{\\sqrt{\\sigma+\\epsilon}}$\n",
        "\n",
        "$y=y_{\\textrm{norm}}\\cdot\\beta+\\alpha$"
      ],
      "metadata": {
        "cell_id": "5b0e5f8055a945929dd74b990ecd2a89",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 91.5,
        "id": "7emhwjKJpOQw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center>\n",
        "<img src=\"https://i.stack.imgur.com/E3104.png\" alt=\"visualization of layer norm vs. batch norm\" width=\"420\">"
      ],
      "metadata": {
        "cell_id": "f96e811cce6d457f85eead1edae4692f",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 92.390625,
        "id": "DYLNkoNNpOQx"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "6d5af9330ebd43c5be334c4f3922031c",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "9c4ce5c0",
        "execution_start": 1656077679263,
        "execution_millis": 142,
        "output_cleared": true,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 346,
        "id": "oLOucdf_pOQx"
      },
      "source": [
        "class LayerNorm(nn.Module):\n",
        "    def __init__(self, norm_shape):\n",
        "        \"\"\"\n",
        "        Applies Layer Normalization over a mini-batch of inputs as described in \n",
        "        [Layer Normalization](https://arxiv.org/pdf/1607.06450.pdf)\n",
        "        \n",
        "        DISCLAIMER: This layer normalization currently works only when the dimension of \n",
        "        the input to be normalized is an integer. \n",
        "\n",
        "        The mean and standard-deviation are calculated over the last 'D' dimensions, \n",
        "        where 'D' is the dimension of :attr:'norm_shape'. If a single integer is used, \n",
        "        it is treated as a singleton list, and this module will normalize over the \n",
        "        last dimension which is expected to be of that specific size. \n",
        "\n",
        "        :param norm_shape: The dimension of the layer to be normalized, i.e. the \n",
        "        shape of the input tensor or the last dimension of the input tensor.\n",
        "\n",
        "        \"\"\"\n",
        "\n",
        "        super().__init__()\n",
        "\n",
        "        # If the given dimension is an integral, make it into a tuple \n",
        "        if isinstance(norm_shape, int):\n",
        "          norm_shape = (norm_shape, )\n",
        "        \n",
        "        self.norm_shape = tuple(norm_shape)\n",
        "\n",
        "\n",
        "        \"\"\"\n",
        "        :param alpha: Scale parameter (Initialized to zeros)\n",
        "        :param beta: Offset parameter (Initialized to ones)\n",
        "        :param epsilon: A value added to the demoniator for numerical stability\n",
        "\n",
        "        The initialized value of alpha and beta should be the same as the dimension\n",
        "        of the layer to be normalized, assuming that the multiplication between\n",
        "        beta and y_norm is element-wise.\n",
        "\n",
        "        Note 1: It is often said that the parameters initialized with nn.Parameter()\n",
        "        often have dimunitive values like 1.4013e-45. Other alternatives such as\n",
        "        nn.Linear() might be worth considering, since it does initial processing to the \n",
        "        input tensor such as uniformization. \n",
        "     \n",
        "        Note 3: The paramters can also be intialized with torch.Tensor(*norm_shape)\n",
        "        or torch.empty(norm_shape)\n",
        "        \"\"\"\n",
        "        self.alpha = torch.nn.Parameter(torch.empty(self.norm_shape))\n",
        "        self.beta = torch.nn.Parameter(torch.rand(self.norm_shape))\n",
        "        self.epsilon = 1e-10\n",
        "\n",
        "        # Initializing the weights with ones and biases with zeros\n",
        "        self.alpha.data.zero_()\n",
        "        self.beta.data.fill_(1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Note 1: It is assumed that the feature dimension is the last dimension, which \n",
        "        matches our assumption from above.\n",
        "        \"\"\"\n",
        "        mean = x.mean(dim=-1, keepdim=True)\n",
        "        var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)\n",
        "        std = (var + self.epsilon).sqrt()\n",
        "        y_norm = (x - mean) / std\n",
        "\n",
        "        y = (y_norm * self.beta) + self.alpha\n",
        "        return y"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# LayerNorm Sanity Check \n",
        "\n",
        "# NLP Example\n",
        "batch, sentence_length, embedding_dim = 20,5,10\n",
        "embedding = torch.randn(batch, sentence_length, embedding_dim)\n",
        "\n",
        "# Layer normalization as implemented by the module nn\n",
        "layer_norm_nn = nn.LayerNorm(embedding_dim)\n",
        "out_nn = layer_norm_nn(embedding)\n",
        "\n",
        "# Our version of layer normalization\n",
        "layer_norm_ori = LayerNorm(embedding_dim)\n",
        "out_ori = layer_norm_ori(embedding)\n",
        "\n",
        "# Check 1: Compare both tensors to see if they match\n",
        "layernorm_compare = torch.eq(torch.round(out_nn, decimals=3), torch.round(out_ori, decimals=3))\n",
        "print(layernorm_compare)\n",
        "# Results: They match, if the values of rounded up to three decimals, which is acceptable\n",
        "\n",
        "# Check 2: See if the output dimension matches the input dimension\n",
        "assert embedding.size() == out_ori.size(), \"Output size does not match with input size\"\n",
        "# Results: Passed.\n",
        "\n",
        "\"\"\"\n",
        "DISCLAIMER: The sanity checks did not check if the learnable affine transform paramters\n",
        "are intialized correctly and can be learned stimultaneously with the other weights of\n",
        "the architecture. \n",
        "\"\"\""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "tAgVGPvqtixi",
        "outputId": "9d84b872-6afd-4943-e5e5-ad26525de3cf"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[[ True,  True,  True,  True,  True,  True,  True, False,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True, False,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True, False,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True, False,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]],\n",
            "\n",
            "        [[ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True],\n",
            "         [ True,  True,  True,  True,  True,  True,  True,  True,  True,  True]]])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\nDISCLAIMER: The sanity checks did not check if the learnable affine transform paramters\\nare intialized correctly and can be learned stimultaneously with the other weights of\\nthe architecture. \\n'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Transformer Block"
      ],
      "metadata": {
        "cell_id": "246bb7a8b64f4596b26778aa9ce5bc85",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "8M3jDuK4pOQy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we bring all ingredients together into a single module. Don't forget to add the residual connections."
      ],
      "metadata": {
        "cell_id": "93c32b55767c46049ffb13bc012bb500",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "mdAHryyXpOQz"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "4e040c8ebd5d42ddb8b00090dffdd960",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 241,
        "id": "GXb1LEVzpOQ0"
      },
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, hidden_n: int, h: int = 2):\n",
        "        \"\"\"\n",
        "        hidden_n: hidden dimension\n",
        "        h: number of heads\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.hidden_n = hidden_n\n",
        "        self.h = h\n",
        "        self.attention = MultiHeadAttention(hidden_n, h)\n",
        "        self.norm1 = LayerNorm(hidden_n)\n",
        "        self.norm2 = LayerNorm(hidden_n)\n",
        "\n",
        "        # This following variable can be eventually defined via input \n",
        "        ff_dim = 2048\n",
        "        self.feed_forward = nn.Sequential(\n",
        "            nn.Linear(hidden_n, ff_dim*hidden_n),\n",
        "            # Here we are using our implemented version of ReLU, which might not be\n",
        "            # compatible in this specific context\n",
        "            relu(),\n",
        "            nn.Linear(ff_dim*hidden_n, hidden_n)\n",
        "        )\n",
        "\n",
        "    # Assuming that the query, key and value arrays are the same and given as input\n",
        "    def forward(self, x: Tensor):\n",
        "        \"\"\"\n",
        "        Alternatively, the inputs could be normalized first before using them as\n",
        "        input for the attention layer, while the residual connections are not normalised\n",
        "\n",
        "        x_norm1 = self.norm1(x)\n",
        "        x_out1 = x + self.attention(x_norm1, x_norm1, x_norm1)\n",
        "        x_norm2 = self.norm_2(x_out1)\n",
        "        x_out2 = x_out1 + self.feed_forward(x_norm2)\n",
        "\n",
        "        return x_out2\n",
        "\n",
        "        \"\"\"\n",
        "        \n",
        "        x_out1 = x + self.attention(x, x, x)\n",
        "        x_attn = self.norm1(x)\n",
        "\n",
        "        x_out2 = x_attn + self.feed_forward(x_attn)\n",
        "        x_ff = self.norm2(x_out2)\n",
        "        return x_ff\n"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A Simple Transformer Architecture"
      ],
      "metadata": {
        "cell_id": "add8b0cdf6ef4a79bf2a71788029d617",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 70,
        "id": "V9e1Z3q6pOQ8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's stack our transformer blocks and add an embedding layer for a simple transformer architecture. You are allowed to use `nn.Embedding` here."
      ],
      "metadata": {
        "cell_id": "0215525d6258493485ff71f3496128a2",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "K12QDX7wpOQ-"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "613b3d3f1e7145ba8ea535c0a34a66a5",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 277,
        "id": "yi8_268IpOQ_"
      },
      "source": [
        "class Transformer(nn.Module):\n",
        "    def __init__(self, emb_n: int, hidden_n: int, n:int =3, h:int =2):\n",
        "        \"\"\"\n",
        "        emb_n: number of token embeddings\n",
        "        hidden_n: hidden dimension\n",
        "        n: number of layers\n",
        "        h: number of heads per layer\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.emb_n = emb_n\n",
        "        self.hidden_n = hidden_n\n",
        "        self.n = n\n",
        "        self.h = h\n",
        "\n",
        "        self.embed = nn.Embedding(emb_n, hidden_n, padding_idx=0)\n",
        "        self.layers = nn.ModuleList(\n",
        "            [\n",
        "                TransformerBlock(hidden_n, h)\n",
        "                for i in range(n)\n",
        "            ]\n",
        "        )\n",
        "\n",
        "    def forward(self, src: Tensor) -> Tensor:\n",
        "      src = self.embed(src)\n",
        "      for layer in self.layers:\n",
        "        src = layer(src)\n",
        "\n",
        "      return src"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## POS-Tagging"
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "24ce63a017054111a0e7606078f244eb",
        "tags": [],
        "is_collapsed": false,
        "deepnote_cell_type": "text-cell-h2",
        "id": "SxwJIhL9pORA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Part-Of-Speech-Tagging (**POS-Tagging**) is a **sequence labeling problem** where we categorize words in a text in correspondence with a particular part of speech (e.g., \"noun\" or \"adjective\"). A few examples and classes are shown in the following table:"
      ],
      "metadata": {
        "cell_id": "46e6832e31434597932051c2d7d833db",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "OAdzksGepORB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "|  POS Tag  |  Description  |  Examples  |\n",
        "|-----------|------------|------------|\n",
        "|  NN | Noun (singular, common) | mass, wind, ...  |\n",
        "|  NNP | Noun (singular, proper) | Obama, Liverpool, ...  |\n",
        "| CD  | Numeral (cardinal)  | 1890, 0.5, ...  |\n",
        "|  DT | Determiner  | all, any, ... |\n",
        "| JJ | Adjective (ordinal) | oiled, third, ... |\n",
        "... many more"
      ],
      "metadata": {
        "cell_id": "0b933b07b931423dae80af31abe46b69",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 200.734375,
        "id": "lQtj6490pORC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### CoNLL2000 Dataset"
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "c1ea426748d34a2891270fb4c9189387",
        "tags": [],
        "is_collapsed": false,
        "deepnote_cell_type": "text-cell-h3",
        "id": "GEVJAksHpORD"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's load our dataset which is the **CoNLL2000 dataset** and look at an example."
      ],
      "metadata": {
        "cell_id": "ec19498428f042a09867d13bfcbbeaf1",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "Fccv3ZmjpORF"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "56f8d97328eb45808c642c232c18ddeb",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "c7ccc53b",
        "execution_start": 1657123575378,
        "execution_millis": 3449,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 329.375,
        "id": "D0jZVbb4pORG",
        "outputId": "214aa423-b5b9-403b-bd9e-addc13cf2266",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchtext.datasets import CoNLL2000Chunking\n",
        "import pandas as pd\n",
        "\n",
        "train_df = pd.DataFrame(CoNLL2000Chunking()[0], columns=['words', 'pos_tags', 'chunk'])\n",
        "test_df = pd.DataFrame(CoNLL2000Chunking()[1], columns=['words', 'pos_tags', 'chunk'])\n",
        "\n",
        "train_src, train_tgt = train_df['words'].tolist(), train_df['pos_tags'].tolist()\n",
        "test_src, test_tgt = test_df['words'].tolist(), test_df['pos_tags'].tolist()\n",
        "\n",
        "print(train_src[0])\n",
        "print(train_tgt[0])"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['Confidence', 'in', 'the', 'pound', 'is', 'widely', 'expected', 'to', 'take', 'another', 'sharp', 'dive', 'if', 'trade', 'figures', 'for', 'September', ',', 'due', 'for', 'release', 'tomorrow', ',', 'fail', 'to', 'show', 'a', 'substantial', 'improvement', 'from', 'July', 'and', 'August', \"'s\", 'near-record', 'deficits', '.']\n",
            "['NN', 'IN', 'DT', 'NN', 'VBZ', 'RB', 'VBN', 'TO', 'VB', 'DT', 'JJ', 'NN', 'IN', 'NN', 'NNS', 'IN', 'NNP', ',', 'JJ', 'IN', 'NN', 'NN', ',', 'VB', 'TO', 'VB', 'DT', 'JJ', 'NN', 'IN', 'NNP', 'CC', 'NNP', 'POS', 'JJ', 'NNS', '.']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, we need to create a vocabulary. Our dataset is already tokenized. However, we need to assign ids to them in order to input them to the embedding layer. We also need the number of embeddings (`num_embeddings`) for the size of our lookup table of `nn.Embedding`.\n",
        "\n",
        "Thus, we will iterate over all sentences replace them with ids and the mapping to our vocabulary. It'll be handy to have two different mappings, from id to token, as well as, from token to id. Note that we will add a special token `<unk>` with id `0` for words that are unknown (that are not in the training dataset but could possibly be in the test dataset)."
      ],
      "metadata": {
        "cell_id": "555caa09a6714a63a9cdee39264291f7",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 178.34375,
        "id": "fvO0ZgRjpORH"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "9b665796e64d4aa694d1718b5156294d",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "2276b8b9",
        "execution_start": 1656084179768,
        "execution_millis": 1512,
        "output_cleared": true,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 148,
        "id": "eSzrCN_YpORJ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ba6d4210-fe85-4abf-d5c6-6f48edfb97e6"
      },
      "source": [
        "# Shift the first entry to the key of one, so that key of zero is saved for the padding token\n",
        "vocabulary_id2token : dict = {1: '<unk>'}\n",
        "vocabulary_token2id : dict = {'<unk>': 1}\n",
        "\n",
        "# Combining all the sentences in the dataset into one big list\n",
        "sentence_list = []\n",
        "for sentence in (train_src + test_src): sentence_list += sentence\n",
        "\n",
        "# Removing duplicates\n",
        "sentence_list = list(set(sentence_list))\n",
        "\n",
        "# Create two enumerate objects, one for each list\n",
        "enum1 = enumerate(sentence_list)\n",
        "enum2 = enumerate(sentence_list)\n",
        "\n",
        "# Create two dictionaries from the two enum objects\n",
        "vocab_id2token = dict((i+2,j) for i,j in enum1)\n",
        "vocab_token2id = dict((j,i+2) for i,j in enum2)\n",
        "\n",
        "# Append the existing dictionaries\n",
        "vocabulary_id2token.update(vocab_id2token)\n",
        "vocabulary_token2id.update(vocab_token2id)\n",
        "\n",
        "# Size of each dictionary\n",
        "voc_id2token_size = len(vocabulary_id2token)\n",
        "voc_token2id_size = len(vocabulary_token2id)\n",
        "\n",
        "# Sanity Check\n",
        "assert len(vocabulary_id2token) == len(vocabulary_token2id), \"The length of the dictionaries should be the same, since they are identical but in different order\"\n",
        "\n",
        "# First first elements of each dictionary\n",
        "l5_vocabid2token = {k: vocabulary_id2token[k] for k in list(vocabulary_id2token)[-5:]}\n",
        "l5_vocab_token2id = {k: vocabulary_token2id[k] for k in list(vocabulary_token2id)[-5:]}\n",
        "\n",
        "print(\"Here are the first five elements from each dictionary:\")\n",
        "print(f\"vocabulary_id2token: {l5_vocabid2token}.\")\n",
        "print(f\"vocabulary_token2id: {l5_vocab_token2id}.\")\n",
        "print()\n",
        "print(f\"Size of the dictionary vocabulary_i2token: {voc_id2token_size}\")\n",
        "print(f\"Size of the dictionary vocabulary_token2id: {voc_token2id_size}\")"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Here are the first five elements from each dictionary:\n",
            "vocabulary_id2token: {21586: 'grandkids', 21587: 'underreacting', 21588: 'absence', 21589: 'Ernesto', 21590: 'Time'}.\n",
            "vocabulary_token2id: {'grandkids': 21586, 'underreacting': 21587, 'absence': 21588, 'Ernesto': 21589, 'Time': 21590}.\n",
            "\n",
            "Size of the dictionary vocabulary_i2token: 21590\n",
            "Size of the dictionary vocabulary_token2id: 21590\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's do the same for our classes:"
      ],
      "metadata": {
        "cell_id": "eed144592b0a4aa9a752fab5230b27e6",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "8Aj_VDDDpORM"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "a4401fdd42024ef9b0af6c95af4ff963",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 133,
        "id": "YSr0BudxpORM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e8fa60d5-3797-4a5b-a6a7-ed2a97fc7981"
      },
      "source": [
        "classes_id2name : dict = {}\n",
        "classes_name2id : dict = {}\n",
        "\n",
        "# Combining all the post-tags in the dataset into one big list\n",
        "pos_tag_list = []\n",
        "for posttag in (train_tgt + test_tgt): pos_tag_list += posttag\n",
        "\n",
        "# Removing duplicates\n",
        "pos_tag_list = list(set(pos_tag_list))\n",
        "\n",
        "# Create two enumerate objects, one for each list\n",
        "enum1 = enumerate(pos_tag_list)\n",
        "enum2 = enumerate(pos_tag_list)\n",
        "\n",
        "# Create two dictionaries from the two enum objects\n",
        "\"\"\"\n",
        "Note 1: We have to skip the tag_id of 0, since we are using that to pad\n",
        "the targets later in both the train and test dataset.\n",
        "\"\"\"\n",
        "posttag_id2name = dict((i+1,j) for i,j in enum1)\n",
        "posttag_name2id = dict((j,i+1) for i,j in enum2)\n",
        "\n",
        "# Append the existing dictionaries\n",
        "classes_id2name.update(posttag_id2name)\n",
        "classes_name2id.update(posttag_name2id)\n",
        "\n",
        "# Size of each dictionary\n",
        "posttag_id2name_size = len(classes_id2name)\n",
        "posttag_name2id_size = len(classes_name2id)\n",
        "\n",
        "# Sanity Check\n",
        "assert posttag_id2name_size == posttag_name2id_size, \"The length of the dictionaries should be the same, since they are identical but in different order\"\n",
        "\n",
        "# First first elements of each dictionary\n",
        "l5_classes_id2name = {k: classes_id2name[k] for k in list(classes_id2name)[-5:]}\n",
        "l5_classes_name2id = {k: classes_name2id[k] for k in list(classes_name2id)[-5:]}\n",
        "\n",
        "print(\"Here are the first five elements from each dictionary:\")\n",
        "print(f\"classes_id2name: {l5_classes_id2name}.\")\n",
        "print(f\"classes_name2id: {l5_classes_name2id}.\")\n",
        "print()\n",
        "print(f\"Size of the dictionary classes_id2name: {posttag_id2name_size}\")\n",
        "print(f\"Size of the dictionary classes_name2id: {posttag_name2id_size}\")\n"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Here are the first five elements from each dictionary:\n",
            "classes_id2name: {40: ',', 41: 'DT', 42: 'NNS', 43: 'PDT', 44: 'WDT'}.\n",
            "classes_name2id: {',': 40, 'DT': 41, 'NNS': 42, 'PDT': 43, 'WDT': 44}.\n",
            "\n",
            "Size of the dictionary classes_id2name: 44\n",
            "Size of the dictionary classes_name2id: 44\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's use PyTorch's `Dataset` and `DataLoader` to help us batch our data. Let's also replace tokens and classes with our ids. For that, complete `get_token_ids` and `get_class_ids`."
      ],
      "metadata": {
        "cell_id": "ef26b718c8654c7cb211c7542430ee79",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "51EuviJDpORO"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "33964355330b45ad98f64b88799e8672",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "a6bcb220",
        "execution_start": 1656333952181,
        "execution_millis": 2,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 508,
        "id": "1bHCQHSHpORP"
      },
      "source": [
        "from typing import List\n",
        "\n",
        "def get_token_ids(src: List[str]) -> List[int]:\n",
        "    return [vocabulary_token2id[word] for word in src]\n",
        "\n",
        "def get_class_ids(tgt: List[str]) -> List[int]:\n",
        "    return [classes_name2id[class_name] for class_name in tgt]\n",
        "\n",
        "class ConllDataset(Dataset):\n",
        "  def __init__(self, src, tgt):\n",
        "        self.src = src\n",
        "        self.tgt = tgt\n",
        "\n",
        "  def __len__(self):\n",
        "        return len(self.src)\n",
        "\n",
        "  def __getitem__(self, index):\n",
        "        src = self.src[index]\n",
        "        tgt = self.tgt[index]\n",
        "        \n",
        "        return {\n",
        "            'src': get_token_ids(src),\n",
        "            'tgt': get_class_ids(tgt),\n",
        "        }\n",
        "\n",
        "train_dataset = ConllDataset(train_src, train_tgt)\n",
        "test_dataset = ConllDataset(test_src, test_tgt)"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will use a **batch size of 32**."
      ],
      "metadata": {
        "cell_id": "ce970c07a9ca4b6ab5a4ee941f3492b9",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "fsxtB4-ApORQ"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "e66f66e38e4343a8ac9237e4e9bf74e5",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "f1c68216",
        "execution_start": 1656086903728,
        "execution_millis": 49,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 76,
        "id": "9vUIXD-mpORa"
      },
      "source": [
        "BATCH_SIZE = 4"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "However, since our examples are of different length, we need to pad shorter examples to the length of the example with the maximum length in our batch. So, let's define a special **padding token** in our vocabulary:"
      ],
      "metadata": {
        "cell_id": "a4c1832e408347d89d304e7870522878",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "qnUZjljzpORb"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "00825cb3c6194a91b8cc03eee1a61e7c",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 115,
        "id": "8kGeA4IjpORc"
      },
      "source": [
        "padding_token = \"PAD\"\n",
        "padding_token_id = 0\n",
        "\n",
        "vocabulary_id2token.update({padding_token_id: padding_token})\n",
        "vocabulary_token2id.update({padding_token: padding_token_id})\n",
        "\n",
        "classes_id2name.update({padding_token_id: padding_token})\n",
        "classes_name2id.update({padding_token: padding_token_id})\n",
        "\n",
        "# Update the sizes of the dictionaries\n",
        "voc_id2token_size = len(vocabulary_id2token)\n",
        "voc_token2id_size = len(vocabulary_token2id)\n",
        "posttag_id2name_size = len(classes_id2name)\n",
        "posttag_name2id_size = len(classes_name2id)\n"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `collate_fn` is the function that actually receives a batch and needs to add the padding tokens, then returns `src` and `tgt` as `Tensor`s of size `[B, S]` where `B` is our batch size and `S` our maximum sequence length. This function should additionally return a `mask`, a `Tensor` with binary values to indicate whether the specific element is a padding token or not (0 if it's a padding token, 1 if not), such that we can ignore padding tokens in our attention mechanism and loss calculation. "
      ],
      "metadata": {
        "cell_id": "c7b65a2d6b654c8bbe59ad7432758a8a",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 119.5625,
        "id": "CswtGa6lpORe"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "9266354695a646ebbd776e4290baa002",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "d3107fb6",
        "execution_start": 1656086903777,
        "execution_millis": 0,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 238,
        "id": "z1Bfar5ApORf"
      },
      "source": [
        "from typing import Dict\n",
        "\n",
        "def collate_fn(batch: List[Dict]) -> Dict[str, Tensor]:\n",
        "    \"\"\"\n",
        "    batch: list of dictionaries with keys src and tgt (as defined in ConllDataset)\n",
        "    e.g. [dict1, dict2, dict3, dict4, ..., dict 32]\n",
        "    \n",
        "    Each dictionary has the following key and values:\n",
        "    \"src\": a list of token ids based on the tokens of the word sequence\n",
        "    \"tgt\": a list of class ids based on the POS-tag of each token\n",
        "\n",
        "    Note 1: We are padding every input of the batch based on the maximum sequence \n",
        "    length of our batch B. Another alternative would be padding every input based\n",
        "    on the maximum sequence length of the entire dataset. \n",
        "    \"\"\"\n",
        "    \n",
        "    #1: Unpack the given batch of data\n",
        "    \"\"\"\n",
        "    token_ids: A list of lists for the token_ids of each sequence within this sampled batch\n",
        "    class_ids: A list of lists for the class_ids of each sequence within this sampled batch\n",
        "    \"\"\"\n",
        "    \n",
        "    token_ids = []           \n",
        "    class_ids = []            \n",
        "    \n",
        "    for x in batch:\n",
        "      token_ids.append(x[\"src\"])\n",
        "      class_ids.append(x[\"tgt\"])\n",
        "      assert len(x[\"src\"]) == len(x[\"tgt\"]), \"The length of the sequences token_ids and class_ids should be the same, since each token should be assigned to a class\"\n",
        "    \n",
        "    assert len(token_ids) == len(class_ids) == BATCH_SIZE, \"The amount of sequences within this batch should correspond with the batch size, and the size of both lists should be the same\"\n",
        "\n",
        "    #2: Get maximum sequence length within this batch \n",
        "    list_len = [len(i) for i in token_ids]\n",
        "    max_seq_len = max(list_len)\n",
        "\n",
        "    #3: Create tensors from the unpacked data and intialize a mask for each sequence in batch\n",
        "    mask = torch.full((BATCH_SIZE, max_seq_len), -1)\n",
        "\n",
        "    #4: Pad sequences that are shorter than the maximum sequence length \n",
        "    #4.1: Iterate over every item in the list, i.e. every word sequence in the batch \n",
        "    for i in range(len(token_ids)):\n",
        "\n",
        "      #4.2: Get each sequence length (to be compared with the maximum sequence length)\n",
        "      seq_length = len(token_ids[i])\n",
        "      length_diff = max_seq_len - seq_length\n",
        "\n",
        "      #4.3: Pad the input sequences if it is shorter than the maximum sequence length (in-place) \n",
        "      padding = [padding_token_id] * length_diff\n",
        "      token_ids[i] = token_ids[i] + padding\n",
        "      class_ids[i] = class_ids[i] + padding\n",
        "\n",
        "      #4.4: Create the attention mask for later use\n",
        "      ones = torch.ones(seq_length)\n",
        "      zeros = torch.zeros(length_diff)\n",
        "      mask[i] = torch.cat((ones, zeros))\n",
        "\n",
        "    #5: Create tensors for src and tgt from nested lists\n",
        "    src = torch.tensor(token_ids)\n",
        "    tgt = torch.tensor(class_ids)\n",
        "\n",
        "    # Sanity checks as per the given requirements\n",
        "    assert len(mask) == BATCH_SIZE, \"The number of masks in the list of masks should correspond to the number of word sequences, i.e. the batch size\"\n",
        "    assert src.size(dim=0) == BATCH_SIZE, \"The size of the first dimension of the src tensor should match the batch size\"\n",
        "    assert src.size(dim=1) == max_seq_len, \"The size of the second dimension of the src tensor should match the maximum sequence length\"\n",
        "    assert tgt.size(dim=0) == BATCH_SIZE, \"The size of the first dimension of the tgt tensor should match the batch size\"\n",
        "    assert tgt.size(dim=1) == max_seq_len, \"The size of the second dimension of the src tensor should match the maximum sequence length\"\n",
        "\n",
        "\n",
        "    return {\n",
        "        'src': src,\n",
        "        'tgt': tgt,\n",
        "        'mask': mask,\n",
        "    }\n",
        "    "
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "With that, we can use PyTorch's `DataLoader` which will shuffle and batch our data automatically."
      ],
      "metadata": {
        "cell_id": "27c262f705004bac83e8306a99ba8bcb",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "knrTyfFOpORh"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "110b64825faf4982943ed9eff63ed0ff",
        "tags": [],
        "deepnote_to_be_reexecuted": false,
        "source_hash": "50d4e2b2",
        "execution_start": 1656086903778,
        "execution_millis": 0,
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 94,
        "id": "iNloZCbKpORi"
      },
      "source": [
        "train_data_loader = DataLoader(train_dataset, collate_fn=collate_fn, batch_size=BATCH_SIZE, shuffle=True)\n",
        "test_data_loader = DataLoader(test_dataset, collate_fn=collate_fn, batch_size=BATCH_SIZE, shuffle=True)"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Architecture"
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "803ef4d85f0a4e2fbcb2ee182520c7f5",
        "tags": [],
        "is_collapsed": false,
        "deepnote_cell_type": "text-cell-h3",
        "id": "TPzC618spORj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's build a transformer model with three layers, three attention heads and an embedding dimension of 128. Also, let's not forget to add a classification head to our model."
      ],
      "metadata": {
        "cell_id": "06f55de40ea248099cbf94bacf9b1268",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "BXp6bahApORk"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "5f1f177f8cf343f1bd763f05a8c8703d",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 223,
        "id": "dJ8QzN22pORl"
      },
      "source": [
        "EMB_DIMENSION = 128\n",
        "NUM_LAYERS = 3\n",
        "NUM_ATT_HEADS = 3\n",
        "\n",
        "class CoNLL2000Transformer(nn.Module):\n",
        "    def __init__(self, transformer, hidden_n:int):\n",
        "        super().__init__()\n",
        "        self.transformer = transformer\n",
        "        self.classification_layer = nn.Linear(hidden_n, posttag_id2name_size)\n",
        "        self.hidden_n = hidden_n\n",
        "\n",
        "    def forward(self, src: Tensor) -> Tensor:\n",
        "      src = self.transformer(src)\n",
        "      output = self.classification_layer(src)\n",
        "      output = softmax(output)\n",
        "\n",
        "      return output\n",
        "\n",
        "\n",
        "\n",
        "model = CoNLL2000Transformer(Transformer(voc_id2token_size, EMB_DIMENSION, NUM_LAYERS, NUM_ATT_HEADS), EMB_DIMENSION)"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training"
      ],
      "metadata": {
        "cell_id": "f5aafb3762974b07be3a8a89899db999",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "UkLzvXY4pORo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Initialize the **AdamW** optimizer from the `torch.optim` module and choose the most appropriate loss function for our task."
      ],
      "metadata": {
        "cell_id": "965bfa0ea57d4052aabfdba6904f65de",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "7e8aN7JspORp"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "5101fa497d074916b870e1e4fe9147d4",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 79,
        "id": "8SggjZW9pORq"
      },
      "source": [
        "optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-8)\n",
        "\n",
        "# We are using cross entropy loss here, since the task is a classification task\n",
        "\"\"\"\n",
        "Note 1: The softmax operation might be included in the cross entropy loss defined below.\n",
        "If this is the case (assuming that it is noticeable from the training loss), then the softmax\n",
        "operation before the output must be removed. \n",
        "\"\"\" \n",
        "criterion = nn.CrossEntropyLoss()"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Build a basic training loop and train the network for three epochs.\n",
        "- Use everything we've built to far, including `train_data_loader`, `model`, `optimizer` and `criterion`.\n",
        "- At every 50th step print the average loss of the last 50 steps. \n",
        "- It is suggested to make a basic training procedure to work on the CPU first. Once it successfully runs on the CPU, you can switch to the GPU (click on change runtime and add an hardware accelerator if you use Colab) and run for the whole three epochs. Note: For this to work, you need to transfer the `model` and the input tensors to the GPU memory. This simply works by calling `.to(device)` on the model and tensors, where `device` and either be `cpu` or `cuda` (for the GPU)."
      ],
      "metadata": {
        "cell_id": "de1bb8eda67644db93779c72899e588c",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 220.734375,
        "id": "WKMcNCFUpORt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "lBcxYjhFwqBn"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "19af38a10aa64403b865dde3603d43e3",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 169,
        "id": "8iRz6ZC7pORu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "16613531-5794-43d1-9779-454440d7a7af"
      },
      "source": [
        "# The variable DEVICE has been moved to the beginning of this notebook\n",
        "EPOCHS = 3\n",
        "\n",
        "\"\"\"\n",
        "Note 1: I am implementing gradient accumulation here, since the number of parameters\n",
        "is too high and the batch_size is only limited to 4, or maybe even 8. Other alternatives\n",
        "to this problem might exist, however were not considered due to the limited time frame\n",
        "given to work on this assignment. \n",
        "\n",
        "Here, the batch size was reduced to 4, while a gradient accumulation technique allows\n",
        "the effective batch size to be 4*8 = 32, i.e. the original pre-set batch size \n",
        "\"\"\"\n",
        "GRAD_ACC = 8\n",
        "\n",
        "model = model.to(DEVICE)\n",
        "\n",
        "\"\"\"\n",
        "DISCLAIMER: The library tqdm and time is installed here to implement a progress bar,\n",
        "so that the training progress can be better checked. Since this has little relevance to\n",
        "the actual exercise, we are hoping this does not affect the grading. \n",
        "\n",
        "    # Setup loop with TQDM and dataloader, can be removed if not allowed\n",
        "    loop = tqdm(train_data_loader, leave=True)\n",
        "    for batch in loop:\n",
        "      \n",
        "      #1: Initialize calculated gradients (from previous step)\n",
        "      optimizer.zero_grad()\n",
        "\n",
        "      #2: Pull all tensor batches required for training\n",
        "      input_sentence = batch[]\n",
        "\n",
        "from tqdm import tqdm\n",
        "import time\n",
        "\"\"\"\n",
        "\n",
        "# Activate training mode \n",
        "model.train()\n",
        "\n",
        "# Initiate losses:\n",
        "running_loss = 0\n",
        "avg_loss_in_50 = 0 \n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "  for i, data in enumerate(train_data_loader):\n",
        "\n",
        "    #1: Unpack the data from the dataloader\n",
        "    \"\"\"\n",
        "    Every data instance is a dictionary of tensors, with each tensor having \n",
        "    the shape (batch_size, maximum_seq_length)\n",
        "    \"\"\"\n",
        "    inputs = data[\"src\"].to(DEVICE)\n",
        "    labels = data[\"tgt\"].to(DEVICE)\n",
        "    masks = data[\"mask\"].to(DEVICE)\n",
        "\n",
        "    #2: Make predictions for this batch \n",
        "    masked_outputs = model(inputs)\n",
        "\n",
        "    #3: Multiply the outputs with the mask to obtain original sequence length \n",
        "    \"\"\"\n",
        "    Note 1: We need to tranpose the outputs so that the shape is (batch_size, posttag_name2id_size, max_seq_length), \n",
        "    instead of (batch_size, max_seq_length, posttag_name2id_size), so that the cross-entropy loss can \n",
        "    be operated based on the given labels and outputs \n",
        "    \"\"\"\n",
        "    outputs = torch.transpose(torch.mul(masked_outputs, torch.unsqueeze(masks, -1)), 1, 2)\n",
        "\n",
        "\n",
        "    #4: Compute the loss and its gradients (w/ gradient accumulation)\n",
        "    loss = criterion(outputs, labels) \n",
        "    (loss / GRAD_ACC).backward()\n",
        "\n",
        "    #5: Adjust learning weights\n",
        "    if (i+1) % GRAD_ACC == 0:\n",
        "      optimizer.step()\n",
        "      model.zero_grad()\n",
        "\n",
        "    #6: Gather data and report \n",
        "    running_loss += loss.item()\n",
        "    if i % 50 == 49:\n",
        "      avg_loss_in_50 = running_loss / 50\n",
        "      print('  batch {} loss: {}'.format(i + 1, avg_loss_in_50))\n",
        "      running_loss = 0\n",
        "      \n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  batch 50 loss: 3.741136088371277\n",
            "  batch 100 loss: 3.6963444757461548\n",
            "  batch 150 loss: 3.686091432571411\n",
            "  batch 200 loss: 3.680394377708435\n",
            "  batch 250 loss: 3.6688568782806397\n",
            "  batch 300 loss: 3.6674539613723756\n",
            "  batch 350 loss: 3.6468973112106324\n",
            "  batch 400 loss: 3.65030912399292\n",
            "  batch 450 loss: 3.6412076854705813\n",
            "  batch 500 loss: 3.623781099319458\n",
            "  batch 550 loss: 3.623956661224365\n",
            "  batch 600 loss: 3.62014262676239\n",
            "  batch 650 loss: 3.6128553104400636\n",
            "  batch 700 loss: 3.5921430444717406\n",
            "  batch 750 loss: 3.594812426567078\n",
            "  batch 800 loss: 3.595312180519104\n",
            "  batch 850 loss: 3.5713216972351076\n",
            "  batch 900 loss: 3.5712246227264406\n",
            "  batch 950 loss: 3.5826078510284423\n",
            "  batch 1000 loss: 3.569926528930664\n",
            "  batch 1050 loss: 3.584328837394714\n",
            "  batch 1100 loss: 3.5478425884246825\n",
            "  batch 1150 loss: 3.572008852958679\n",
            "  batch 1200 loss: 3.566558141708374\n",
            "  batch 1250 loss: 3.5537648677825926\n",
            "  batch 1300 loss: 3.557211012840271\n",
            "  batch 1350 loss: 3.551089506149292\n",
            "  batch 1400 loss: 3.5465862131118775\n",
            "  batch 1450 loss: 3.5623220586776734\n",
            "  batch 1500 loss: 3.5649252986907958\n",
            "  batch 1550 loss: 3.5523539352416993\n",
            "  batch 1600 loss: 3.543512954711914\n",
            "  batch 1650 loss: 3.5497413444519044\n",
            "  batch 1700 loss: 3.5481896924972536\n",
            "  batch 1750 loss: 3.549347171783447\n",
            "  batch 1800 loss: 3.5497617053985597\n",
            "  batch 1850 loss: 3.539479079246521\n",
            "  batch 1900 loss: 3.5372597980499267\n",
            "  batch 1950 loss: 3.5410515451431275\n",
            "  batch 2000 loss: 3.541939125061035\n",
            "  batch 2050 loss: 3.5312809896469117\n",
            "  batch 2100 loss: 3.5495792865753173\n",
            "  batch 2150 loss: 3.537966780662537\n",
            "  batch 2200 loss: 3.5413275718688966\n",
            "  batch 50 loss: 5.955001673698425\n",
            "  batch 100 loss: 3.5409623479843138\n",
            "  batch 150 loss: 3.5359039545059203\n",
            "  batch 200 loss: 3.5307200288772584\n",
            "  batch 250 loss: 3.5280885028839113\n",
            "  batch 300 loss: 3.531585359573364\n",
            "  batch 350 loss: 3.526674957275391\n",
            "  batch 400 loss: 3.5380522775650025\n",
            "  batch 450 loss: 3.5185292243957518\n",
            "  batch 500 loss: 3.51458327293396\n",
            "  batch 550 loss: 3.5264933490753174\n",
            "  batch 600 loss: 3.5298707723617553\n",
            "  batch 650 loss: 3.5278522062301634\n",
            "  batch 700 loss: 3.5076611757278444\n",
            "  batch 750 loss: 3.5347572469711306\n",
            "  batch 800 loss: 3.5221783924102783\n",
            "  batch 850 loss: 3.509157056808472\n",
            "  batch 900 loss: 3.5281166076660155\n",
            "  batch 950 loss: 3.5227133941650393\n",
            "  batch 1000 loss: 3.530958065986633\n",
            "  batch 1050 loss: 3.519317140579224\n",
            "  batch 1100 loss: 3.5433973455429078\n",
            "  batch 1150 loss: 3.517849168777466\n",
            "  batch 1200 loss: 3.5107548570632936\n",
            "  batch 1250 loss: 3.5328478813171387\n",
            "  batch 1300 loss: 3.5199922275543214\n",
            "  batch 1350 loss: 3.507997360229492\n",
            "  batch 1400 loss: 3.508193893432617\n",
            "  batch 1450 loss: 3.525442752838135\n",
            "  batch 1500 loss: 3.517271695137024\n",
            "  batch 1550 loss: 3.5262843132019044\n",
            "  batch 1600 loss: 3.529618511199951\n",
            "  batch 1650 loss: 3.508557367324829\n",
            "  batch 1700 loss: 3.509939184188843\n",
            "  batch 1750 loss: 3.5238031435012815\n",
            "  batch 1800 loss: 3.51554573059082\n",
            "  batch 1850 loss: 3.505765428543091\n",
            "  batch 1900 loss: 3.5028772497177125\n",
            "  batch 1950 loss: 3.507873983383179\n",
            "  batch 2000 loss: 3.5350933027267457\n",
            "  batch 2050 loss: 3.504694037437439\n",
            "  batch 2100 loss: 3.508936085700989\n",
            "  batch 2150 loss: 3.5128893423080445\n",
            "  batch 2200 loss: 3.513604130744934\n",
            "  batch 50 loss: 5.894685635566711\n",
            "  batch 100 loss: 3.5147788619995115\n",
            "  batch 150 loss: 3.50993528842926\n",
            "  batch 200 loss: 3.522672576904297\n",
            "  batch 250 loss: 3.507266502380371\n",
            "  batch 300 loss: 3.510155873298645\n",
            "  batch 350 loss: 3.5050617170333864\n",
            "  batch 400 loss: 3.5103186321258546\n",
            "  batch 450 loss: 3.507813892364502\n",
            "  batch 500 loss: 3.5056025695800783\n",
            "  batch 550 loss: 3.5078724193573\n",
            "  batch 600 loss: 3.530418386459351\n",
            "  batch 650 loss: 3.50709014415741\n",
            "  batch 700 loss: 3.5002088737487793\n",
            "  batch 750 loss: 3.525472545623779\n",
            "  batch 800 loss: 3.5019213008880614\n",
            "  batch 850 loss: 3.5076691818237307\n",
            "  batch 900 loss: 3.500648760795593\n",
            "  batch 950 loss: 3.5048089933395388\n",
            "  batch 1000 loss: 3.496481900215149\n",
            "  batch 1050 loss: 3.4992879676818847\n",
            "  batch 1100 loss: 3.5001492404937746\n",
            "  batch 1150 loss: 3.5019566726684572\n",
            "  batch 1200 loss: 3.5053520727157594\n",
            "  batch 1250 loss: 3.5024382781982424\n",
            "  batch 1300 loss: 3.505931782722473\n",
            "  batch 1350 loss: 3.5054508876800536\n",
            "  batch 1400 loss: 3.509045286178589\n",
            "  batch 1450 loss: 3.5212679386138914\n",
            "  batch 1500 loss: 3.4858214235305787\n",
            "  batch 1550 loss: 3.5121602964401246\n",
            "  batch 1600 loss: 3.5040649604797363\n",
            "  batch 1650 loss: 3.5084221267700197\n",
            "  batch 1700 loss: 3.500363645553589\n",
            "  batch 1750 loss: 3.482015166282654\n",
            "  batch 1800 loss: 3.51060001373291\n",
            "  batch 1850 loss: 3.503650107383728\n",
            "  batch 1900 loss: 3.50427933216095\n",
            "  batch 1950 loss: 3.5053282833099364\n",
            "  batch 2000 loss: 3.5006749391555787\n",
            "  batch 2050 loss: 3.493293595314026\n",
            "  batch 2100 loss: 3.5028782367706297\n",
            "  batch 2150 loss: 3.4939092063903807\n",
            "  batch 2200 loss: 3.490356755256653\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluation"
      ],
      "metadata": {
        "cell_id": "2f12827ec5ad4a4ba18ae8d260857cd8",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "8kWWHxb1pORv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's see what's the accuracy is of our model. Since we already implemented accuracy in the previous exercise, we'll now let you use the torchmetrics package."
      ],
      "metadata": {
        "cell_id": "23af9a15c9e84bd3bea4a62c7c7e9459",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "7TWPAFvFpORv"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "18891986243d49afa9305716cbd96aa7",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 97,
        "id": "6ElKGbYmpORw"
      },
      "source": [
        "from torchmetrics import Accuracy\n",
        "\n",
        "\"\"\"\n",
        "Note 1: The parameter average='micro' sometimes bugged the usage of the accuracy metric, \n",
        "creating an error. If this is the case, removing the parameter would help solve the issue.\n",
        "This does not change anything, since this setting is also the default setting. \n",
        "\"\"\"\n",
        "accuracy = Accuracy(average='micro').to(DEVICE)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Calculate the average accuracy of all examples in the test dataset."
      ],
      "metadata": {
        "cell_id": "1cd5f69a8dcd41f8918bc9946a1b57e3",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "SKdTakYSpORw"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "f73337d4e25a494e9c1b138823bc1d22",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 61,
        "id": "YUd9H43EpORx",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3f3cbe12-e3c9-44cf-c131-0f855ca234b0"
      },
      "source": [
        "model.eval()\n",
        "\n",
        "with torch.no_grad():\n",
        "  for i, data in enumerate(test_data_loader):\n",
        "      \n",
        "      #1: Unpack the data from the dataloader\n",
        "      \"\"\"\n",
        "      Every data instance is a dictionary of tensors, with each tensor having \n",
        "      the shape (batch_size, maximum_seq_length)\n",
        "      \"\"\"\n",
        "      inputs = data[\"src\"].to(DEVICE)\n",
        "      labels = data[\"tgt\"].to(DEVICE)\n",
        "      masks = data[\"mask\"].to(DEVICE)\n",
        "\n",
        "      #2: Make predictions for this batch \n",
        "      masked_outputs = model(inputs)\n",
        "\n",
        "      #3: Multiply the outputs with the mask\n",
        "      \"\"\"\n",
        "      Note 1: Since we are in inference / evaluation mode, we take the argmax of the last dimension, i.e. posttag_name2id_size, \n",
        "      and reshape the shape of the predictions so that it batches the labels, i.e. (batch_size, max_seq_length)\n",
        "\n",
        "      Shapes:\n",
        "      unmasked_outputs: (batch_size, max_seq_length, posttag_name2id_size)\n",
        "      pred: (batch_size, max_seq_length)\n",
        "\n",
        "      \"\"\"\n",
        "      unmasked_outputs = torch.mul(masked_outputs, torch.unsqueeze(masks, -1)).to(DEVICE)\n",
        "      values, indices = torch.topk(unmasked_outputs, k=1, dim=-1)\n",
        "      pred = torch.squeeze(indices, -1).to(DEVICE)\n",
        "\n",
        "      \"\"\"\n",
        "      Why is the size of the second dimension here not BATCH_SIZE?                 !!!\n",
        "      \"\"\"\n",
        "\n",
        "      for x in range(BATCH_SIZE):\n",
        "        batch_acc = accuracy(labels[x], pred[x])\n",
        "      \n",
        "      #4: Print batch accuracy for each 25 batches \n",
        "      if i % 25 == 24:\n",
        "        print('  Batch {} Batch accuracy: {}'.format(i + 1, batch_acc))\n",
        "\n",
        "print('Overall accuracy:', accuracy.compute())"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Batch 25 Batch accuracy: 0.8275862336158752\n",
            "  Batch 50 Batch accuracy: 0.8888888955116272\n",
            "  Batch 75 Batch accuracy: 0.5151515007019043\n",
            "  Batch 100 Batch accuracy: 0.8928571343421936\n",
            "  Batch 125 Batch accuracy: 0.7142857313156128\n",
            "  Batch 150 Batch accuracy: 0.8235294222831726\n",
            "  Batch 175 Batch accuracy: 0.7435897588729858\n",
            "  Batch 200 Batch accuracy: 0.75\n",
            "  Batch 225 Batch accuracy: 0.9111111164093018\n",
            "  Batch 250 Batch accuracy: 0.9677419066429138\n",
            "  Batch 275 Batch accuracy: 0.7678571343421936\n",
            "  Batch 300 Batch accuracy: 0.8260869383811951\n",
            "  Batch 325 Batch accuracy: 0.6206896305084229\n",
            "  Batch 350 Batch accuracy: 0.8125\n",
            "  Batch 375 Batch accuracy: 0.6585366129875183\n",
            "  Batch 400 Batch accuracy: 0.7777777910232544\n",
            "  Batch 425 Batch accuracy: 0.8529411554336548\n",
            "  Batch 450 Batch accuracy: 0.800000011920929\n",
            "  Batch 475 Batch accuracy: 0.8085106611251831\n",
            "  Batch 500 Batch accuracy: 0.9047619104385376\n",
            "Overall accuracy: tensor(0.8107, device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's also look at the accuracy **for each class separately**:"
      ],
      "metadata": {
        "cell_id": "d8a4e62ebfb64dc5a6659df362eaeb50",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "JyOpnkgQpORx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "b-ryoAW4yV14"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class_acc = Accuracy(average='none', num_classes=posttag_id2name_size).to(DEVICE)\n",
        "for i, data in enumerate(test_data_loader):\n",
        "    \n",
        "    #1: Unpack the data from the dataloader\n",
        "    \"\"\"\n",
        "    Every data instance is a dictionary of tensors, with each tensor having \n",
        "    the shape (batch_size, maximum_seq_length)\n",
        "    \"\"\"\n",
        "    inputs = data[\"src\"].to(DEVICE)\n",
        "    labels = data[\"tgt\"].to(DEVICE)\n",
        "    masks = data[\"mask\"].to(DEVICE)\n",
        "\n",
        "    #2: Make predictions for this batch \n",
        "    masked_outputs = model(inputs)\n",
        "\n",
        "    #3: Multiply the outputs with the mask\n",
        "    \"\"\"\n",
        "    Note 1: Since we are in inference / evaluation mode, we take the argmax of the last dimension, i.e. posttag_name2id_size, \n",
        "    and reshape the shape of the predictions so that it batches the labels, i.e. (batch_size, max_seq_length)\n",
        "\n",
        "    Shapes:\n",
        "    unmasked_outputs: (batch_size, max_seq_length, posttag_name2id_size)\n",
        "    pred: (batch_size, max_seq_length)\n",
        "\n",
        "    \"\"\"\n",
        "    unmasked_outputs = torch.mul(masked_outputs, torch.unsqueeze(masks, -1)).to(DEVICE)\n",
        "    values, indices = torch.topk(unmasked_outputs, k=1, dim=-1)\n",
        "    pred = torch.squeeze(indices, -1).to(DEVICE)\n",
        "\n",
        "    for i in range(BATCH_SIZE):\n",
        "      batch_acc = class_acc(labels[i], pred[i])\n",
        "\n",
        "    \"\"\"\n",
        "    #4: Print batch accuracy for each 25 batches \n",
        "    if i % 25 == 24:\n",
        "      print('  Batch {} Batch accuracy: {}'.format(i + 1, batch_acc))\n",
        "    \"\"\"\n",
        "\n"
      ],
      "metadata": {
        "id": "-xkOhnwHM1bk",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 359
        },
        "outputId": "826052ea-4d29-46c0-eca0-3f3231522ac1"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-41-f37d722d602b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0;31m#2: Make predictions for this batch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0mmasked_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0;31m#3: Multiply the outputs with the mask\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-19-d566952a7e4d>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, src)\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msrc\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m       \u001b[0msrc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m       \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclassification_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m       \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-10-248847fdd296>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, src)\u001b[0m\n\u001b[1;32m     24\u001b[0m       \u001b[0msrc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m       \u001b[0;32mfor\u001b[0m \u001b[0mlayer\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayers\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m         \u001b[0msrc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlayer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msrc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0msrc\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-9-aa80ba28a0a4>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0mx_attn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 42\u001b[0;31m         \u001b[0mx_out2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx_attn\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeed_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_attn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m         \u001b[0mx_ff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_out2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mx_ff\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    137\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    138\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 139\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    140\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-3-2f813a2ac554>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaximum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mDEVICE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzeros\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mDEVICE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "per_class_acc = class_acc.compute().tolist()\n",
        "per_class_acc = [round(item, 4) for item in per_class_acc]\n",
        "class_ids = range(posttag_id2name_size)\n",
        "class_accuracies = {class_ids[i]: per_class_acc[i] for i in range(len(per_class_acc))}\n",
        "\n",
        "class_accuracies"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "76ZhX8OIyUZf",
        "outputId": "c5679fc7-175c-433e-c06c-039fc21bbff1"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{0: 0.9935,\n",
              " 1: 0.1871,\n",
              " 2: 0.3125,\n",
              " 3: 0.3065,\n",
              " 4: 0.0,\n",
              " 5: nan,\n",
              " 6: 0.144,\n",
              " 7: 0.743,\n",
              " 8: nan,\n",
              " 9: 0.0,\n",
              " 10: 0.0,\n",
              " 11: 0.3571,\n",
              " 12: 0.0,\n",
              " 13: 0.0588,\n",
              " 14: 1.0,\n",
              " 15: 0.3623,\n",
              " 16: 0.4,\n",
              " 17: 0.0,\n",
              " 18: 0.0,\n",
              " 19: 0.7939,\n",
              " 20: 0.0677,\n",
              " 21: 0.0,\n",
              " 22: 0.0811,\n",
              " 23: 0.0,\n",
              " 24: 0.0,\n",
              " 25: 0.0,\n",
              " 26: 0.0,\n",
              " 27: 0.0,\n",
              " 28: 0.0,\n",
              " 29: 0.0,\n",
              " 30: 0.0,\n",
              " 31: 0.0,\n",
              " 32: nan,\n",
              " 33: 0.9722,\n",
              " 34: 0.8856,\n",
              " 35: nan,\n",
              " 36: 0.2435,\n",
              " 37: 0.9825,\n",
              " 38: 0.0,\n",
              " 39: 0.0752,\n",
              " 40: 0.0682,\n",
              " 41: 0.2593,\n",
              " 42: 0.0,\n",
              " 43: 0.0}"
            ]
          },
          "metadata": {},
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Positional Embeddings"
      ],
      "metadata": {
        "formattedRanges": [],
        "cell_id": "c92bfb82034344bb9cd220de8e4e157c",
        "tags": [],
        "is_collapsed": false,
        "deepnote_cell_type": "text-cell-h2",
        "id": "TG0fqk_apORy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The attention mechanism does not consider the position of the tokens which hurts its performance for many problems. We can solve this issue in several ways. We can either add a positional encoding (via trigonometric functions) or we can learn positional embeddings along the way, in a similar way as BERT does. Here, we will add learnable positional embeddings to our exisisting model with another embedding layer."
      ],
      "metadata": {
        "cell_id": "070435ca57a94f1ea6e71050d57619aa",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 119.5625,
        "id": "Q7EM7zGdpORy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The longest sequence in our dataset has 78 tokens (you can trust us on that). So, let's set the number of embeddings for our positional embedding layer to that number. Again, you should use `nn.Embedding`."
      ],
      "metadata": {
        "cell_id": "5ff82560d6f24630a9bed80d6b38b7f1",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "W9HUerk7pOSJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Copy the inner parts of your `Transformer` class and add positional embeddings to it."
      ],
      "metadata": {
        "cell_id": "0363608d5a9a48f5be6f2e6ddbd7397c",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "RPqh5asxpOSK"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "129e0b2b46c8494c90491fe2d2e137f0",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 313,
        "id": "Zp6iSw4rpOSL"
      },
      "source": [
        "MAX_SEQ_LEN = 78\n",
        "\n",
        "class TransformerPos(nn.Module):\n",
        "    def __init__(self, emb_n: int, pos_emb_n: int, hidden_n: int, n:int =3, h:int =2):\n",
        "        \"\"\"\n",
        "        emb_n: number of token embeddings\n",
        "        pos_emb_n: number of position embeddings\n",
        "        hidden_n: hidden dimension\n",
        "        n: number of layers\n",
        "        h: number of heads per layer\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.positional_embeddings = pos_emb_n\n",
        "        self.emb_n = emb_n\n",
        "        self.hidden_n = hidden_n\n",
        "        self.n = n\n",
        "        self.h = h\n",
        "\n",
        "        self.embed = nn.Embedding(emb_n, hidden_n, padding_idx=0)\n",
        "        self.pos_emb = nn.Embedding(MAX_SEQ_LEN, hidden_n)\n",
        "        self.layers = nn.ModuleList(\n",
        "            [\n",
        "                TransformerBlock(hidden_n, h)\n",
        "                for i in range(n)\n",
        "            ]\n",
        "        )\n",
        "\n",
        "    def forward(self, src: Tensor) -> Tensor:\n",
        "      batch_size, seq_length = src.shape\n",
        "      positions = torch.arrange(0, seq_length).expand(N, seq_length).to(DEVICE)\n",
        "      src = self.embed(src) + self.pos_emb(positions)\n",
        "      for layer in self.layers:\n",
        "        src = layer(src)\n",
        "\n",
        "      return src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "b398b56387f54999966ee9d22a0150b1",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 61,
        "id": "bpR_ON4MpOSP"
      },
      "source": [
        "model_pos = CoNLL2000Transformer(TransformerPos(voc_id2token_size, MAX_SEQ_LEN, EMB_DIMENSION, NUM_LAYERS, NUM_ATT_HEADS), EMB_DIMENSION)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training"
      ],
      "metadata": {
        "cell_id": "d8f6b79fce144176832210f7a1494288",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "8dQriOaJpOSR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Same procedure as before. Let's reinitialize our optimizer and our loss function and run the same training loop with our new model `model_pos`."
      ],
      "metadata": {
        "cell_id": "507700c451e842c4b9dbdd02857b234d",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 74.78125,
        "id": "1U37j-KwpOSR"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "5101fa497d074916b870e1e4fe9147d4",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 79,
        "id": "jEUroeq3_Mb2"
      },
      "source": [
        "optimizer = torch.optim.AdamW(model_pos.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-8)\n",
        "\n",
        "# We are using cross entropy loss here, since the task is a classification task\n",
        "\"\"\"\n",
        "Note 1: The softmax operation might be included in the cross entropy loss defined below.\n",
        "If this is the case (assuming that it is noticeable from the training loss), then the softmax\n",
        "operation before the output must be removed. \n",
        "\"\"\" \n",
        "criterion = nn.CrossEntropyLoss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "vCHSPaXa_Mb2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "19af38a10aa64403b865dde3603d43e3",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 169,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "16613531-5794-43d1-9779-454440d7a7af",
        "id": "X5mK8W8r_Mb3"
      },
      "source": [
        "# The variable DEVICE has been moved to the beginning of this notebook\n",
        "EPOCHS = 3\n",
        "\n",
        "\"\"\"\n",
        "Note 1: I am implementing gradient accumulation here, since the number of parameters\n",
        "is too high and the batch_size is only limited to 4, or maybe even 8. Other alternatives\n",
        "to this problem might exist, however were not considered due to the limited time frame\n",
        "given to work on this assignment. \n",
        "\n",
        "Here, the batch size was reduced to 4, while a gradient accumulation technique allows\n",
        "the effective batch size to be 4*8 = 32, i.e. the original pre-set batch size \n",
        "\"\"\"\n",
        "GRAD_ACC = 8\n",
        "\n",
        "model_pos = model_pos.to(DEVICE)\n",
        "\n",
        "\"\"\"\n",
        "DISCLAIMER: The library tqdm and time is installed here to implement a progress bar,\n",
        "so that the training progress can be better checked. Since this has little relevance to\n",
        "the actual exercise, we are hoping this does not affect the grading. \n",
        "\n",
        "    # Setup loop with TQDM and dataloader, can be removed if not allowed\n",
        "    loop = tqdm(train_data_loader, leave=True)\n",
        "    for batch in loop:\n",
        "      \n",
        "      #1: Initialize calculated gradients (from previous step)\n",
        "      optimizer.zero_grad()\n",
        "\n",
        "      #2: Pull all tensor batches required for training\n",
        "      input_sentence = batch[]\n",
        "\n",
        "from tqdm import tqdm\n",
        "import time\n",
        "\"\"\"\n",
        "\n",
        "# Activate training mode \n",
        "model_pos.train()\n",
        "\n",
        "# Initiate losses:\n",
        "running_loss = 0\n",
        "avg_loss_in_50 = 0 \n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "  for i, data in enumerate(train_data_loader):\n",
        "\n",
        "    #1: Unpack the data from the dataloader\n",
        "    \"\"\"\n",
        "    Every data instance is a dictionary of tensors, with each tensor having \n",
        "    the shape (batch_size, maximum_seq_length)\n",
        "    \"\"\"\n",
        "    inputs = data[\"src\"].to(DEVICE)\n",
        "    labels = data[\"tgt\"].to(DEVICE)\n",
        "    masks = data[\"mask\"].to(DEVICE)\n",
        "\n",
        "    #2: Make predictions for this batch \n",
        "    masked_outputs = model_pos(inputs)\n",
        "\n",
        "    #3: Multiply the outputs with the mask to obtain original sequence length \n",
        "    \"\"\"\n",
        "    Note 1: We need to tranpose the outputs so that the shape is (batch_size, posttag_name2id_size, max_seq_length), \n",
        "    instead of (batch_size, max_seq_length, posttag_name2id_size), so that the cross-entropy loss can \n",
        "    be operated based on the given labels and outputs \n",
        "    \"\"\"\n",
        "    outputs = torch.transpose(torch.mul(masked_outputs, torch.unsqueeze(masks, -1)), 1, 2)\n",
        "\n",
        "\n",
        "    #4: Compute the loss and its gradients (w/ gradient accumulation)\n",
        "    loss = criterion(outputs, labels) \n",
        "    (loss / GRAD_ACC).backward()\n",
        "\n",
        "    #5: Adjust learning weights\n",
        "    if (i+1) % GRAD_ACC == 0:\n",
        "      optimizer.step()\n",
        "      model_pos.zero_grad()\n",
        "\n",
        "    #6: Gather data and report \n",
        "    running_loss += loss.item()\n",
        "    if i % 50 == 49:\n",
        "      avg_loss_in_50 = running_loss / 50\n",
        "      print('  batch {} loss: {}'.format(i + 1, avg_loss_in_50))\n",
        "      running_loss = 0\n",
        "      \n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  batch 50 loss: 3.741136088371277\n",
            "  batch 100 loss: 3.6963444757461548\n",
            "  batch 150 loss: 3.686091432571411\n",
            "  batch 200 loss: 3.680394377708435\n",
            "  batch 250 loss: 3.6688568782806397\n",
            "  batch 300 loss: 3.6674539613723756\n",
            "  batch 350 loss: 3.6468973112106324\n",
            "  batch 400 loss: 3.65030912399292\n",
            "  batch 450 loss: 3.6412076854705813\n",
            "  batch 500 loss: 3.623781099319458\n",
            "  batch 550 loss: 3.623956661224365\n",
            "  batch 600 loss: 3.62014262676239\n",
            "  batch 650 loss: 3.6128553104400636\n",
            "  batch 700 loss: 3.5921430444717406\n",
            "  batch 750 loss: 3.594812426567078\n",
            "  batch 800 loss: 3.595312180519104\n",
            "  batch 850 loss: 3.5713216972351076\n",
            "  batch 900 loss: 3.5712246227264406\n",
            "  batch 950 loss: 3.5826078510284423\n",
            "  batch 1000 loss: 3.569926528930664\n",
            "  batch 1050 loss: 3.584328837394714\n",
            "  batch 1100 loss: 3.5478425884246825\n",
            "  batch 1150 loss: 3.572008852958679\n",
            "  batch 1200 loss: 3.566558141708374\n",
            "  batch 1250 loss: 3.5537648677825926\n",
            "  batch 1300 loss: 3.557211012840271\n",
            "  batch 1350 loss: 3.551089506149292\n",
            "  batch 1400 loss: 3.5465862131118775\n",
            "  batch 1450 loss: 3.5623220586776734\n",
            "  batch 1500 loss: 3.5649252986907958\n",
            "  batch 1550 loss: 3.5523539352416993\n",
            "  batch 1600 loss: 3.543512954711914\n",
            "  batch 1650 loss: 3.5497413444519044\n",
            "  batch 1700 loss: 3.5481896924972536\n",
            "  batch 1750 loss: 3.549347171783447\n",
            "  batch 1800 loss: 3.5497617053985597\n",
            "  batch 1850 loss: 3.539479079246521\n",
            "  batch 1900 loss: 3.5372597980499267\n",
            "  batch 1950 loss: 3.5410515451431275\n",
            "  batch 2000 loss: 3.541939125061035\n",
            "  batch 2050 loss: 3.5312809896469117\n",
            "  batch 2100 loss: 3.5495792865753173\n",
            "  batch 2150 loss: 3.537966780662537\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluation"
      ],
      "metadata": {
        "cell_id": "2dfedd24910c45689ed2a3d4ed804af6",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 62,
        "id": "U2lnOjv3pOST"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's check if our performance on the accuracy got improved."
      ],
      "metadata": {
        "cell_id": "472e05fab0d149b8b269a12d8e363b85",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "GFFi_8DOpOST"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "38e0fe847e464f59b1861a749fc43811",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 61,
        "id": "ErAMnxfZpOSU"
      },
      "source": [
        "..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Again, let's also check each class. Which classes got improved the most by adding positional embeddings?"
      ],
      "metadata": {
        "cell_id": "a49fe34010574a0096a020ead994158d",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 52.390625,
        "id": "KkcUg_bFpOSU"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "cell_id": "d107d15cf5fe4ee396f6733e093873cc",
        "tags": [],
        "deepnote_cell_type": "code",
        "deepnote_cell_height": 61,
        "id": "AyfXP2GNpOSV"
      },
      "source": [
        "..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The last question in this assignment doesn't require you to code anything. Instead, you're asked to point out possible issues with our current approach and name potential improvements. \n",
        "* ...\n",
        "* ...\n",
        "* ..."
      ],
      "metadata": {
        "cell_id": "bfece568a8c1460fafd968adf1972f19",
        "tags": [],
        "deepnote_cell_type": "markdown",
        "deepnote_cell_height": 175.953125,
        "id": "Zv6hz5NqpOSV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a style='text-decoration:none;line-height:16px;display:flex;color:#5B5B62;padding:10px;justify-content:end;' href='https://deepnote.com?utm_source=created-in-deepnote-cell&projectId=faa4af3b-d086-4f42-8b7d-d29c91b1d0f6' target=\"_blank\">\n",
        "<img alt='Created in deepnote.com' style='display:inline;max-height:16px;margin:0px;margin-right:7.5px;' src='data:image/svg+xml;base64,PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz4KPHN2ZyB3aWR0aD0iODBweCIgaGVpZ2h0PSI4MHB4IiB2aWV3Qm94PSIwIDAgODAgODAiIHZlcnNpb249IjEuMSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIiB4bWxuczp4bGluaz0iaHR0cDovL3d3dy53My5vcmcvMTk5OS94bGluayI+CiAgICA8IS0tIEdlbmVyYXRvcjogU2tldGNoIDU0LjEgKDc2NDkwKSAtIGh0dHBzOi8vc2tldGNoYXBwLmNvbSAtLT4KICAgIDx0aXRsZT5Hcm91cCAzPC90aXRsZT4KICAgIDxkZXNjPkNyZWF0ZWQgd2l0aCBTa2V0Y2guPC9kZXNjPgogICAgPGcgaWQ9IkxhbmRpbmciIHN0cm9rZT0ibm9uZSIgc3Ryb2tlLXdpZHRoPSIxIiBmaWxsPSJub25lIiBmaWxsLXJ1bGU9ImV2ZW5vZGQiPgogICAgICAgIDxnIGlkPSJBcnRib2FyZCIgdHJhbnNmb3JtPSJ0cmFuc2xhdGUoLTEyMzUuMDAwMDAwLCAtNzkuMDAwMDAwKSI+CiAgICAgICAgICAgIDxnIGlkPSJHcm91cC0zIiB0cmFuc2Zvcm09InRyYW5zbGF0ZSgxMjM1LjAwMDAwMCwgNzkuMDAwMDAwKSI+CiAgICAgICAgICAgICAgICA8cG9seWdvbiBpZD0iUGF0aC0yMCIgZmlsbD0iIzAyNjVCNCIgcG9pbnRzPSIyLjM3NjIzNzYyIDgwIDM4LjA0NzY2NjcgODAgNTcuODIxNzgyMiA3My44MDU3NTkyIDU3LjgyMTc4MjIgMzIuNzU5MjczOSAzOS4xNDAyMjc4IDMxLjY4MzE2ODMiPjwvcG9seWdvbj4KICAgICAgICAgICAgICAgIDxwYXRoIGQ9Ik0zNS4wMDc3MTgsODAgQzQyLjkwNjIwMDcsNzYuNDU0OTM1OCA0Ny41NjQ5MTY3LDcxLjU0MjI2NzEgNDguOTgzODY2LDY1LjI2MTk5MzkgQzUxLjExMjI4OTksNTUuODQxNTg0MiA0MS42NzcxNzk1LDQ5LjIxMjIyODQgMjUuNjIzOTg0Niw0OS4yMTIyMjg0IEMyNS40ODQ5Mjg5LDQ5LjEyNjg0NDggMjkuODI2MTI5Niw0My4yODM4MjQ4IDM4LjY0NzU4NjksMzEuNjgzMTY4MyBMNzIuODcxMjg3MSwzMi41NTQ0MjUgTDY1LjI4MDk3Myw2Ny42NzYzNDIxIEw1MS4xMTIyODk5LDc3LjM3NjE0NCBMMzUuMDA3NzE4LDgwIFoiIGlkPSJQYXRoLTIyIiBmaWxsPSIjMDAyODY4Ij48L3BhdGg+CiAgICAgICAgICAgICAgICA8cGF0aCBkPSJNMCwzNy43MzA0NDA1IEwyNy4xMTQ1MzcsMC4yNTcxMTE0MzYgQzYyLjM3MTUxMjMsLTEuOTkwNzE3MDEgODAsMTAuNTAwMzkyNyA4MCwzNy43MzA0NDA1IEM4MCw2NC45NjA0ODgyIDY0Ljc3NjUwMzgsNzkuMDUwMzQxNCAzNC4zMjk1MTEzLDgwIEM0Ny4wNTUzNDg5LDc3LjU2NzA4MDggNTMuNDE4MjY3Nyw3MC4zMTM2MTAzIDUzLjQxODI2NzcsNTguMjM5NTg4NSBDNTMuNDE4MjY3Nyw0MC4xMjg1NTU3IDM2LjMwMzk1NDQsMzcuNzMwNDQwNSAyNS4yMjc0MTcsMzcuNzMwNDQwNSBDMTcuODQzMDU4NiwzNy43MzA0NDA1IDkuNDMzOTE5NjYsMzcuNzMwNDQwNSAwLDM3LjczMDQ0MDUgWiIgaWQ9IlBhdGgtMTkiIGZpbGw9IiMzNzkzRUYiPjwvcGF0aD4KICAgICAgICAgICAgPC9nPgogICAgICAgIDwvZz4KICAgIDwvZz4KPC9zdmc+' > </img>\n",
        "Created in <span style='font-weight:600;margin-left:4px;'>Deepnote</span></a>"
      ],
      "metadata": {
        "tags": [],
        "created_in_deepnote_cell": true,
        "deepnote_cell_type": "markdown",
        "id": "LwD3wZunpOSW"
      }
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "orig_nbformat": 2,
    "deepnote": {},
    "deepnote_notebook_id": "df3e6be7-b747-492e-86ed-19bc62a6bb4c",
    "deepnote_execution_queue": [],
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "language_info": {
      "name": "python"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  }
}