Knowledge_Distillation_using_T5_Online_Teacher.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ambideXtrous9/MTP-FlanT5-SBERT-Model-for-NewsQA/blob/main/Knowledge_Distillation_using_T5_Online_Teacher.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --quiet transformers\n",
        "!pip install --quiet pytorch-lightning\n",
        "!pip install --quiet lightning"
      ],
      "metadata": {
        "id": "8Z-el9G5cvmu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eSbQEF5vJS-V"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import numpy as np\n",
        "import pytorch_lightning as pl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torchmetrics.functional import accuracy\n",
        "from torchmetrics.classification import Accuracy\n",
        "from torch.optim import AdamW\n",
        "from transformers import (AutoTokenizer, AutoModelForSeq2SeqLM)\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from sklearn.model_selection import train_test_split\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from pytorch_lightning.callbacks import RichProgressBar,ModelCheckpoint\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hkrOBVTjcqwQ",
        "outputId": "e2e71636-d218-4477-a385-b30caeb26fe2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g1uZn0rrcZmg"
      },
      "outputs": [],
      "source": [
        "path = '/content/drive/MyDrive/MTP CODE/NewsQA_SPAN.feather'\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 195
        },
        "id": "C5hu4V57cZmg",
        "outputId": "318be020-fcf9-484b-935c-354508d42a01"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                            question        answer  \\\n",
              "0  Who is the managing director of Synergee Capital?  Vikram Dalal   \n",
              "1  What is the yield of 30- and 40-year governmen...            7%   \n",
              "2  What is the name of the ETF 2027 that a conser...           SDL   \n",
              "\n",
              "      ans_pos                                          paragraph  \\\n",
              "0  [133, 145]   \"Investors can use a combination of governmen...   \n",
              "1  [565, 567]   \"Investors can use a combination of governmen...   \n",
              "2  [209, 212]   According to financial planners, an example o...   \n",
              "\n",
              "   answer_start  answer_end  \n",
              "0           133         145  \n",
              "1           565         567  \n",
              "2           209         212  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-45afdb51-8ce1-467f-9d68-076f1ecbc44a\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>question</th>\n",
              "      <th>answer</th>\n",
              "      <th>ans_pos</th>\n",
              "      <th>paragraph</th>\n",
              "      <th>answer_start</th>\n",
              "      <th>answer_end</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>Who is the managing director of Synergee Capital?</td>\n",
              "      <td>Vikram Dalal</td>\n",
              "      <td>[133, 145]</td>\n",
              "      <td>\"Investors can use a combination of governmen...</td>\n",
              "      <td>133</td>\n",
              "      <td>145</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>What is the yield of 30- and 40-year governmen...</td>\n",
              "      <td>7%</td>\n",
              "      <td>[565, 567]</td>\n",
              "      <td>\"Investors can use a combination of governmen...</td>\n",
              "      <td>565</td>\n",
              "      <td>567</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>What is the name of the ETF 2027 that a conser...</td>\n",
              "      <td>SDL</td>\n",
              "      <td>[209, 212]</td>\n",
              "      <td>According to financial planners, an example o...</td>\n",
              "      <td>209</td>\n",
              "      <td>212</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-45afdb51-8ce1-467f-9d68-076f1ecbc44a')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-45afdb51-8ce1-467f-9d68-076f1ecbc44a button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-45afdb51-8ce1-467f-9d68-076f1ecbc44a');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-354a3549-93d7-424c-9032-db984301cee1\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-354a3549-93d7-424c-9032-db984301cee1')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-354a3549-93d7-424c-9032-db984301cee1 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "df = pd.read_feather(path)\n",
        "df.head(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fk91G_EzcZmh"
      },
      "outputs": [],
      "source": [
        "df = df.iloc[:5000]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t-l062pacZmh"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = 'google/flan-t5-small'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BBoQzvd8cZmh"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0bCUCVHVcZmi",
        "outputId": "8da83217-7778-4912-c3e3-77a5eb879c8a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Start Token: None\n",
            "End Token: None\n"
          ]
        }
      ],
      "source": [
        "start_token = tokenizer.cls_token\n",
        "end_token = tokenizer.sep_token\n",
        "\n",
        "print(\"Start Token:\", start_token)\n",
        "print(\"End Token:\", end_token)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JdsQRfiQcZmi",
        "outputId": "1a96ffb1-f8bb-4aa4-a1fd-9dc83b89725f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "32100\n"
          ]
        }
      ],
      "source": [
        "vocab_size = tokenizer.vocab_size\n",
        "print(vocab_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xFRuuamXcZmi"
      },
      "outputs": [],
      "source": [
        "class NQADataset(Dataset):\n",
        "    def __init__(self,data ,tokenizer ,source_max_token_len : int = 180,target_max_token_len : int = 5):\n",
        "\n",
        "        self.tokenizer = tokenizer\n",
        "        self.data = data\n",
        "        self.source_max_token_len = source_max_token_len\n",
        "        self.target_max_token_len = target_max_token_len\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self,index : int):\n",
        "        data_row = self.data.iloc[index]\n",
        "\n",
        "        source_encoding = self.tokenizer(\n",
        "            data_row['question'],\n",
        "            data_row['paragraph'],\n",
        "            max_length = self.source_max_token_len,\n",
        "            padding = \"max_length\",\n",
        "            truncation = \"only_second\",\n",
        "            return_attention_mask = True,\n",
        "            add_special_tokens = True,\n",
        "            return_tensors = \"pt\")\n",
        "\n",
        "        target_encoding = self.tokenizer(\n",
        "            data_row['answer'],\n",
        "            max_length = self.target_max_token_len,\n",
        "            padding = \"max_length\",\n",
        "            truncation = True,\n",
        "            return_attention_mask = True,\n",
        "            add_special_tokens = True,\n",
        "            return_tensors = \"pt\")\n",
        "\n",
        "        labels = target_encoding[\"input_ids\"]\n",
        "        # labels[labels == 0] = -100\n",
        "\n",
        "        return dict(\n",
        "            input_ids = source_encoding['input_ids'].flatten(),\n",
        "            attention_mask = source_encoding['attention_mask'].flatten(),\n",
        "            labels = labels.flatten())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q_9n8yuLcZmi"
      },
      "outputs": [],
      "source": [
        "\n",
        "class NQADataModule(pl.LightningDataModule):\n",
        "  def __init__(self,train_df,val_df,tokenizer,batch_size : int = 8,source_max_token_len : int = 180,target_max_token_len : int = 5):\n",
        "    super().__init__()\n",
        "    self.batch_size = batch_size\n",
        "    self.train_df = train_df\n",
        "    self.val_df = val_df\n",
        "    self.MODEL_NAME = MODEL_NAME\n",
        "    self.tokenizer = tokenizer\n",
        "    self.source_max_token_len = source_max_token_len\n",
        "    self.target_max_token_len = target_max_token_len\n",
        "\n",
        "  def setup(self,stage=None):\n",
        "    self.train_dataset = NQADataset(self.train_df,self.tokenizer,self.source_max_token_len,self.target_max_token_len)\n",
        "    self.val_dataset = NQADataset(self.val_df,self.tokenizer,self.source_max_token_len,self.target_max_token_len)\n",
        "\n",
        "\n",
        "  def train_dataloader(self):\n",
        "    return DataLoader(self.train_dataset,batch_size = self.batch_size,shuffle=True,num_workers=2)\n",
        "\n",
        "  def val_dataloader(self):\n",
        "    return DataLoader(self.val_dataset,batch_size = self.batch_size,num_workers=2)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Le9oNWlcZmi"
      },
      "outputs": [],
      "source": [
        "train_df, val_df = train_test_split(df,test_size=0.1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5WSme8T_cZmj",
        "outputId": "ceddea31-9860-4277-dffa-dd7bdfc1c3be"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(4500, 6)\n",
            "(500, 6)\n"
          ]
        }
      ],
      "source": [
        "print(train_df.shape)\n",
        "print(val_df.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hCT6Z4L3cZmj"
      },
      "outputs": [],
      "source": [
        "data_module = NQADataModule(train_df,val_df,tokenizer,batch_size = 8)\n",
        "data_module.setup()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n6ZOHv0DcZmj",
        "outputId": "ac5d4adc-9200-473a-acfc-17235bad279d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "input_ids: torch.Size([8, 180])\n",
            "attention_mask: torch.Size([8, 180])\n",
            "labels: torch.Size([8, 5])\n"
          ]
        }
      ],
      "source": [
        "sample_batch = next(iter(data_module.val_dataloader()))\n",
        "for key, value in sample_batch.items():\n",
        "    print(f\"{key}: {value.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wBb57WGjfijP"
      },
      "source": [
        "## Model\n",
        "\n",
        "![](https://media.arxiv-vanity.com/render-output/3715543/Figures/ModalNet-21.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jfd9TGWQcZmk"
      },
      "outputs": [],
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
        "        super(PositionalEncoding, self).__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0), :]\n",
        "        return self.dropout(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jDYV4CrjJiJw"
      },
      "outputs": [],
      "source": [
        "class Transformer(nn.Module):\n",
        "    \"\"\"\n",
        "    Classic Transformer that both encodes and decodes.\n",
        "    Prediction-time inference is done greedily.\n",
        "    NOTE: start token is hard-coded to be 0, end token to be 1. If changing, update predict() accordingly.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, num_classes: int, max_output_length: int, dim: int = 512):\n",
        "        super().__init__()\n",
        "\n",
        "        # Parameters\n",
        "        self.t5model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME,return_dict=True)\n",
        "\n",
        "        self.dim = dim\n",
        "        self.start_token = '[CLS]'\n",
        "        self.max_output_length = max_output_length\n",
        "        self.log_softmax = nn.LogSoftmax()\n",
        "\n",
        "        nhead = 8\n",
        "        num_layers = 8\n",
        "        dim_feedforward = dim\n",
        "\n",
        "        # Encoder part\n",
        "        self.embedding = nn.Embedding(num_classes, dim)\n",
        "        self.pos_encoder = PositionalEncoding(d_model=self.dim)\n",
        "        self.transformer_encoder = nn.TransformerEncoder(\n",
        "            encoder_layer=nn.TransformerEncoderLayer(d_model=self.dim,\n",
        "                                                     nhead=nhead,\n",
        "                                                     dim_feedforward=dim_feedforward),\n",
        "            num_layers=num_layers\n",
        "        )\n",
        "\n",
        "        # Decoder part\n",
        "        self.y_mask = self.generate_square_subsequent_mask(self.max_output_length)\n",
        "        self.transformer_decoder = nn.TransformerDecoder(\n",
        "            decoder_layer=nn.TransformerDecoderLayer(d_model=self.dim,\n",
        "                                                     nhead=nhead,\n",
        "                                                     dim_feedforward=dim_feedforward),\n",
        "            num_layers=num_layers\n",
        "        )\n",
        "        self.fc = nn.Linear(self.dim, num_classes)\n",
        "\n",
        "        # It is empirically important to initialize weights properly\n",
        "        self.init_weights()\n",
        "\n",
        "    def generate_square_subsequent_mask(self,size: int):\n",
        "        \"\"\"Generate a triangular (size, size) mask. From PyTorch docs.\"\"\"\n",
        "        mask = (torch.triu(torch.ones(size, size)) == 1).transpose(0, 1)\n",
        "        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
        "        return mask\n",
        "\n",
        "    def init_weights(self):\n",
        "        initrange = 0.1\n",
        "        self.embedding.weight.data.uniform_(-initrange, initrange)\n",
        "        self.fc.bias.data.zero_()\n",
        "        self.fc.weight.data.uniform_(-initrange, initrange)\n",
        "\n",
        "    def forward(self, x: torch.Tensor, y: torch.Tensor) :\n",
        "        \"\"\"\n",
        "        Input\n",
        "            x: (B, Sx) with elements in (0, C) where C is num_classes\n",
        "            y: (B, Sy) with elements in (0, C) where C is num_classes\n",
        "        Output\n",
        "            (B, C, Sy) logits\n",
        "        \"\"\"\n",
        "        encoded_x = self.encode(x)  # (Sx, B, E)\n",
        "        output = self.decode(y, encoded_x)  # (Sy, B, C)\n",
        "        return output.permute(1, 2, 0)  # (B, C, Sy)\n",
        "\n",
        "    def encode(self, x: torch.Tensor) :\n",
        "        \"\"\"\n",
        "        Input => x: (B, Sx) with elements in (0, C) where C is num_classes\n",
        "        Output => (Sx, B, E) embedding\n",
        "        \"\"\"\n",
        "        x = x.permute(1, 0)  # (Sx, B, E)\n",
        "        x = self.embedding(x) * math.sqrt(self.dim)  # (Sx, B, E)\n",
        "        x = self.pos_encoder(x)  # (Sx, B, E)\n",
        "        x = self.transformer_encoder(x)  # (Sx, B, E)\n",
        "        return x\n",
        "\n",
        "    def decode(self, y: torch.Tensor, encoded_x: torch.Tensor):\n",
        "        \"\"\"\n",
        "        Input\n",
        "            encoded_x: (Sx, B, E)\n",
        "            y: (B, Sy) with elements in (0, C) where C is num_classes\n",
        "        Output\n",
        "            (Sy, B, C) logits\n",
        "        \"\"\"\n",
        "        y = y.permute(1, 0)  # (Sy, B)\n",
        "        y = self.embedding(y) * math.sqrt(self.dim)  # (Sy, B, E)\n",
        "        y = self.pos_encoder(y)  # (Sy, B, E)\n",
        "        Sy = y.shape[0]\n",
        "        y_mask = self.y_mask[:Sy, :Sy].type_as(encoded_x)  # (Sy, Sy)\n",
        "        output = self.transformer_decoder(y, encoded_x, y_mask)  # (Sy, B, E)\n",
        "        output = self.fc(output)  # (Sy, B, C)\n",
        "        output = self.log_softmax(output)\n",
        "        return output\n",
        "\n",
        "    def teacher_flant5(self,input_ids,attention_mask,labels=None):\n",
        "        output = self.t5model(input_ids = input_ids,\n",
        "                            attention_mask = attention_mask,\n",
        "                            labels = labels)\n",
        "\n",
        "        return output.loss, output.logits.permute(0, 2, 1)\n",
        "\n",
        "    def predict(self, x: torch.Tensor) :\n",
        "        \"\"\"\n",
        "        Method to use at inference time. Predict y from x one token at a time. This method is greedy\n",
        "        decoding. Beam search can be used instead for a potential accuracy boost.\n",
        "\n",
        "        Input\n",
        "            x: (B, Sx) with elements in (0, C) where C is num_classes\n",
        "        Output\n",
        "            (B, C, Sy) logits\n",
        "        \"\"\"\n",
        "        encoded_x = self.encode(x)\n",
        "\n",
        "        output_tokens = (torch.ones((x.shape[0], self.max_output_length))).type_as(x).long() # (B, max_length)\n",
        "        output_tokens[:, 0] = 0  # Set start token\n",
        "        for Sy in range(1, self.max_output_length):\n",
        "            y = output_tokens[:, :Sy]  # (B, Sy)\n",
        "            output = self.decode(y, encoded_x)  # (Sy, B, C)\n",
        "            output = torch.argmax(output, dim=-1)  # (Sy, B)\n",
        "            output_tokens[:, Sy] = output[-1:]  # Set the last output token\n",
        "        return output_tokens\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4YX25k0JjU7"
      },
      "outputs": [],
      "source": [
        "class LitModel(pl.LightningModule):\n",
        "    \"\"\"Simple PyTorch-Lightning model to train our Transformer.\"\"\"\n",
        "\n",
        "    def __init__(self, model):\n",
        "        super().__init__()\n",
        "        self.model = model\n",
        "        self.t5model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME,return_dict=True)\n",
        "        self.loss = nn.CrossEntropyLoss(ignore_index=0)\n",
        "        self.criterion = nn.NLLLoss(ignore_index=0)\n",
        "        self.KLD = nn.KLDivLoss(reduction='batchmean')\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "\n",
        "        self.valloss = []\n",
        "        self.valacc = []\n",
        "        self.valt5acc = []\n",
        "\n",
        "        # for param in self.t5model.parameters():\n",
        "        #     param.requires_grad = False\n",
        "\n",
        "\n",
        "    def training_step(self, batch, batch_ind):\n",
        "        x = batch['input_ids']\n",
        "        y = batch['labels']\n",
        "        attention_mask = batch['attention_mask']\n",
        "\n",
        "\n",
        "        t5loss, t5logits = self.model.teacher_flant5(x,attention_mask,y)\n",
        "\n",
        "\n",
        "        # Teacher forcing: model gets input up to the last character,\n",
        "        # while ground truth (loss is calculated) is from the second character onward.\n",
        "        logits = self.model(x, y[:, :-1])\n",
        "        ce_loss = self.criterion(logits, y[:, 1:])\n",
        "\n",
        "        t5logits_reduced = t5logits[:, :32100, :4]\n",
        "\n",
        "\n",
        "        encod_probs = self.softmax(logits)\n",
        "        logits_prob = self.softmax(t5logits_reduced)\n",
        "\n",
        "        kld_loss = self.KLD(encod_probs.log(), logits_prob)\n",
        "\n",
        "\n",
        "        loss = kld_loss + ce_loss + t5loss\n",
        "\n",
        "        self.log_dict({\"KLD_loss\" : kld_loss,\n",
        "                      \"DEC_loss\" : ce_loss,\n",
        "                       \"T5_loss\" : t5loss,\n",
        "                      \"total_loss\" : loss},prog_bar=True,logger=True)\n",
        "        return kld_loss\n",
        "\n",
        "    def validation_step(self, batch, batch_ind):\n",
        "        x = batch['input_ids']\n",
        "        y = batch['labels']\n",
        "        attention_mask = batch['attention_mask']\n",
        "        # Teacher forcing: model gets input up to the last character,\n",
        "        # while ground truth (loss is calculated) is from the second character onward.\n",
        "        logits = self.model(x, y[:, :-1])\n",
        "        loss = self.criterion(logits, y[:, 1:])\n",
        "        pred = self.model.predict(x)\n",
        "\n",
        "\n",
        "        t5loss, t5logits = self.model.teacher_flant5(x,attention_mask,y)\n",
        "\n",
        "        t5probs = torch.nn.functional.softmax(t5logits, dim=1)\n",
        "\n",
        "        t5_pred = torch.argmax(t5probs, dim=1)\n",
        "\n",
        "\n",
        "        correct_predictions = (y == t5_pred).float()\n",
        "        row_accuracies = torch.mean(correct_predictions, dim=1)\n",
        "        t5accuracy = torch.mean(row_accuracies).item()\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "        correct_predictions = (y == pred).float()\n",
        "        row_accuracies = torch.mean(correct_predictions, dim=1)\n",
        "        accuracy = torch.mean(row_accuracies).item()\n",
        "\n",
        "        self.log_dict({\"val_acc\" : accuracy,\n",
        "                       \"val_t5acc\" : t5accuracy,\n",
        "                       \"val_loss\" : loss,\n",
        "                       },on_step=False, on_epoch=True, prog_bar=True)\n",
        "\n",
        "        self.valloss.append(loss.item())\n",
        "        self.valacc.append(accuracy)\n",
        "        self.valt5acc.append(t5accuracy)\n",
        "\n",
        "    def on_validation_epoch_end(self):\n",
        "\n",
        "        avg_val_loss = torch.mean(torch.tensor(self.valloss))\n",
        "        avg_val_acc = torch.mean(torch.tensor(self.valacc))\n",
        "        avg_valt5_acc = torch.mean(torch.tensor(self.valt5acc))\n",
        "\n",
        "        self.valacc.clear()\n",
        "        self.valloss.clear()\n",
        "        self.valt5acc.clear()\n",
        "\n",
        "        self.log_dict({\"avg_val_acc\" : avg_val_acc,\n",
        "                       \"avg_val_loss\" : avg_val_loss,\n",
        "                       \"avg_valt5_acc\" : avg_valt5_acc,\n",
        "                       },on_step=False, on_epoch=True, prog_bar=True)\n",
        "\n",
        "        print('-' * 90)\n",
        "        print(f'|  Epoch = {self.current_epoch + 1:3d} | '\n",
        "            f'Avg Val Loss: {avg_val_loss.item():5.2f} | '\n",
        "            f'Avg Val Accuracy: {avg_val_acc.item():5.2f} |'\n",
        "            f'Avg Val T5 Accuracy: {avg_valt5_acc.item():5.2f} |')\n",
        "        print('-' * 90)\n",
        "\n",
        "\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        return AdamW(self.parameters(),lr = 0.0001)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lF0vcnzVcZmm"
      },
      "outputs": [],
      "source": [
        "from lightning.pytorch.loggers import TensorBoardLogger\n",
        "\n",
        "logger = TensorBoardLogger(\"tb_logs\", name=\"my_model\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AwU4gsWtcZmm"
      },
      "outputs": [],
      "source": [
        "checkpoint_callback = ModelCheckpoint(\n",
        "    dirpath = 'checkpoints',\n",
        "    filename = 'Transformer_best_cp',\n",
        "    save_top_k = 1,\n",
        "    verbose = True,\n",
        "    monitor = 'avg_val_loss',\n",
        "    mode = 'min')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MtfnN2pEcZmm"
      },
      "outputs": [],
      "source": [
        "rpb = RichProgressBar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uZkz0fyvcZmn",
        "outputId": "e42795ce-632f-45f9-ddd7-efbcd9ebe89b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO: GPU available: True (cuda), used: True\n",
            "INFO:lightning.pytorch.utilities.rank_zero:GPU available: True (cuda), used: True\n",
            "INFO: TPU available: False, using: 0 TPU cores\n",
            "INFO:lightning.pytorch.utilities.rank_zero:TPU available: False, using: 0 TPU cores\n",
            "INFO: IPU available: False, using: 0 IPUs\n",
            "INFO:lightning.pytorch.utilities.rank_zero:IPU available: False, using: 0 IPUs\n",
            "INFO: HPU available: False, using: 0 HPUs\n",
            "INFO:lightning.pytorch.utilities.rank_zero:HPU available: False, using: 0 HPUs\n"
          ]
        }
      ],
      "source": [
        "trainer = pl.Trainer(accelerator=\"gpu\",\n",
        "                     devices=-1,\n",
        "                     logger=logger,\n",
        "                     callbacks=[checkpoint_callback,rpb],\n",
        "                     max_epochs = 20)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fs-hYIsUcZmn"
      },
      "outputs": [],
      "source": [
        "model = Transformer(num_classes=vocab_size, max_output_length=5)\n",
        "lit_model = LitModel(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 670,
          "referenced_widgets": [
            "dbb4976205aa47739d0a4fa6db7f6fe4",
            "b5253629ba31440c9144c8a36f9f91a2"
          ]
        },
        "id": "T-6_JRt5cZmn",
        "outputId": "ff4e2256-6d90-407d-d700-056b93fc660e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:pytorch_lightning.accelerators.cuda:LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┏━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n",
              "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName     \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType                      \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\n",
              "┡━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n",
              "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ model     │ Transformer                │  143 M │\n",
              "│\u001b[2m \u001b[0m\u001b[2m1\u001b[0m\u001b[2m \u001b[0m│ t5model   │ T5ForConditionalGeneration │ 77.0 M │\n",
              "│\u001b[2m \u001b[0m\u001b[2m2\u001b[0m\u001b[2m \u001b[0m│ loss      │ CrossEntropyLoss           │      0 │\n",
              "│\u001b[2m \u001b[0m\u001b[2m3\u001b[0m\u001b[2m \u001b[0m│ criterion │ NLLLoss                    │      0 │\n",
              "│\u001b[2m \u001b[0m\u001b[2m4\u001b[0m\u001b[2m \u001b[0m│ KLD       │ KLDivLoss                  │      0 │\n",
              "│\u001b[2m \u001b[0m\u001b[2m5\u001b[0m\u001b[2m \u001b[0m│ softmax   │ Softmax                    │      0 │\n",
              "└───┴───────────┴────────────────────────────┴────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n",
              "┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">   </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Name      </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Type                       </span>┃<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> Params </span>┃\n",
              "┡━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 0 </span>│ model     │ Transformer                │  143 M │\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 1 </span>│ t5model   │ T5ForConditionalGeneration │ 77.0 M │\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 2 </span>│ loss      │ CrossEntropyLoss           │      0 │\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 3 </span>│ criterion │ NLLLoss                    │      0 │\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 4 </span>│ KLD       │ KLDivLoss                  │      0 │\n",
              "│<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 5 </span>│ softmax   │ Softmax                    │      0 │\n",
              "└───┴───────────┴────────────────────────────┴────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mTrainable params\u001b[0m: 220 M                                                                                            \n",
              "\u001b[1mNon-trainable params\u001b[0m: 0                                                                                            \n",
              "\u001b[1mTotal params\u001b[0m: 220 M                                                                                                \n",
              "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 881                                                                        \n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Trainable params</span>: 220 M                                                                                            \n",
              "<span style=\"font-weight: bold\">Non-trainable params</span>: 0                                                                                            \n",
              "<span style=\"font-weight: bold\">Total params</span>: 220 M                                                                                                \n",
              "<span style=\"font-weight: bold\">Total estimated model params size (MB)</span>: 881                                                                        \n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Output()"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "dbb4976205aa47739d0a4fa6db7f6fe4"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "|  Epoch =   1 | Avg Val Loss:  1.38 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.46 |\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">|  Epoch =   1 | Avg Val Loss:  1.38 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.46 |\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "|  Epoch =   1 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.10 |\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">|  Epoch =   1 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.10 |\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO: Epoch 0, global step 563: 'avg_val_loss' reached 1.38702 (best 1.38702), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n",
            "INFO:lightning.pytorch.utilities.rank_zero:Epoch 0, global step 563: 'avg_val_loss' reached 1.38702 (best 1.38702), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "|  Epoch =   2 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.02 |\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">|  Epoch =   2 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.02 |\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO: Epoch 1, global step 1126: 'avg_val_loss' reached 1.38633 (best 1.38633), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n",
            "INFO:lightning.pytorch.utilities.rank_zero:Epoch 1, global step 1126: 'avg_val_loss' reached 1.38633 (best 1.38633), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "|  Epoch =   3 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.02 |\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">|  Epoch =   3 | Avg Val Loss:  1.39 | Avg Val Accuracy:  0.00 |Avg Val T5 Accuracy:  0.02 |\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "------------------------------------------------------------------------------------------\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">------------------------------------------------------------------------------------------\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO: Epoch 2, global step 1689: 'avg_val_loss' reached 1.38520 (best 1.38520), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n",
            "INFO:lightning.pytorch.utilities.rank_zero:Epoch 2, global step 1689: 'avg_val_loss' reached 1.38520 (best 1.38520), saving model to '/content/checkpoints/Transformer_best_cp-v2.ckpt' as top 1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "trainer.fit(lit_model, data_module)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JuHxC9KzcZmo"
      },
      "outputs": [],
      "source": [
        "cppath = 'checkpoints/Transformer_best_cp-v2.ckpt'\n",
        "trained_model = LitModel.load_from_checkpoint(cppath,\n",
        "                                              model=model)\n",
        "trained_model.freeze()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FsbXQYPacZmo",
        "outputId": "a82056f9-72f6-41b6-9e6b-03abc4f9d9ae"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "question        How much did Shiromani Akali Dal raise in 2019...\n",
              "answer                                                       16.4\n",
              "ans_pos                                                [676, 680]\n",
              "paragraph        TDP managed a total of Rs 91.5 crore in 2019-...\n",
              "answer_start                                                  676\n",
              "answer_end                                                    680\n",
              "Name: 2110, dtype: object"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ],
      "source": [
        "val_df.iloc[12]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dOIFwwqDcZmo"
      },
      "outputs": [],
      "source": [
        "source_encoding = tokenizer(\n",
        "            val_df.iloc[0]['question'],\n",
        "            val_df.iloc[0]['paragraph'],\n",
        "            max_length = 180,\n",
        "            padding = \"max_length\",\n",
        "            truncation = \"only_second\",\n",
        "            return_attention_mask = True,\n",
        "            add_special_tokens = True,\n",
        "            return_tensors = \"pt\")\n",
        "\n",
        "input_ids = source_encoding['input_ids'].flatten().to(\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DmhnKZv2Jk_J",
        "outputId": "1647d5dc-a72a-4c38-cb81-64e87826988c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pred:\n",
            "studenți experiments\n"
          ]
        }
      ],
      "source": [
        "# We can see that the decoding works correctly\n",
        "pred = trained_model.model.predict(input_ids.unsqueeze(0))\n",
        "print('Pred:')\n",
        "tk = tokenizer.decode(pred[0],skip_special_tokens=True)\n",
        "print(tk)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "dbb4976205aa47739d0a4fa6db7f6fe4": {
          "model_module": "@jupyter-widgets/output",
          "model_name": "OutputModel",
          "model_module_version": "1.0.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/output",
            "_model_module_version": "1.0.0",
            "_model_name": "OutputModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/output",
            "_view_module_version": "1.0.0",
            "_view_name": "OutputView",
            "layout": "IPY_MODEL_b5253629ba31440c9144c8a36f9f91a2",
            "msg_id": "",
            "outputs": [
              {
                "output_type": "display_data",
                "data": {
                  "text/plain": "Epoch 3/19 \u001b[38;2;98;6;224m━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[38;2;98;6;224m╸\u001b[0m\u001b[38;5;237m━━━━━━━━━━━━━━\u001b[0m \u001b[37m326/563\u001b[0m \u001b[38;5;245m0:01:22 • 0:01:00\u001b[0m \u001b[38;5;249m3.99it/s\u001b[0m \u001b[37mv_num: 20.000 KLD_loss: 448.503   \u001b[0m\n                                                                                 \u001b[37mDEC_loss: 1.392 T5_loss: 10.264   \u001b[0m\n                                                                                 \u001b[37mtotal_loss: 460.159 val_acc: 0.000\u001b[0m\n                                                                                 \u001b[37mval_t5acc: 0.017 val_loss: 1.385  \u001b[0m\n                                                                                 \u001b[37mavg_val_acc: 0.000 avg_val_loss:  \u001b[0m\n                                                                                 \u001b[37m1.385 avg_valt5_acc: 0.017        \u001b[0m\n",
                  "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Epoch 3/19 <span style=\"color: #6206e0; text-decoration-color: #6206e0\">━━━━━━━━━━━━━━━━━━━╸</span><span style=\"color: #3a3a3a; text-decoration-color: #3a3a3a\">━━━━━━━━━━━━━━</span> <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">326/563</span> <span style=\"color: #8a8a8a; text-decoration-color: #8a8a8a\">0:01:22 • 0:01:00</span> <span style=\"color: #b2b2b2; text-decoration-color: #b2b2b2\">3.99it/s</span> <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">v_num: 20.000 KLD_loss: 448.503   </span>\n                                                                                 <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">DEC_loss: 1.392 T5_loss: 10.264   </span>\n                                                                                 <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">total_loss: 460.159 val_acc: 0.000</span>\n                                                                                 <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">val_t5acc: 0.017 val_loss: 1.385  </span>\n                                                                                 <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">avg_val_acc: 0.000 avg_val_loss:  </span>\n                                                                                 <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">1.385 avg_valt5_acc: 0.017        </span>\n</pre>\n"
                },
                "metadata": {}
              }
            ]
          }
        },
        "b5253629ba31440c9144c8a36f9f91a2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}