Deep_Learning_V0_9_3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I5wpYP2tFoP1"
      },
      "source": [
        "# Deep Learning Project\n",
        "How to run this Notebook:\n",
        "- Run every cell up to the Training Chapter\n",
        "\n",
        "For training:\n",
        "- Simply decide to run a Generative training approach or a Discriminative one, running only the cell.\n",
        "\n",
        "For testing:\n",
        "- Run the \"DSI Model with Foundation Model\" chapter, then run the \"Restore a Checkpoint and run a Test\" one."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VWJOpDVyK604"
      },
      "source": [
        "## Install required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TCueW7x8K818",
        "outputId": "ab4304e1-2476-41fe-b3a4-122c5ef28272"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using Colab. Downloading necessary libraries...\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m105.0/105.0 MB\u001b[0m \u001b[31m9.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m280.0/280.0 kB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m183.4/183.4 kB\u001b[0m \u001b[31m7.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m800.9/800.9 kB\u001b[0m \u001b[31m13.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m27.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m300.8/300.8 kB\u001b[0m \u001b[31m23.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m815.9/815.9 kB\u001b[0m \u001b[31m43.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m840.4/840.4 kB\u001b[0m \u001b[31m45.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m195.4/195.4 kB\u001b[0m \u001b[31m21.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m257.9/257.9 kB\u001b[0m \u001b[31m23.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m829.5/829.5 kB\u001b[0m \u001b[31m50.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m61.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m241.2/241.2 kB\u001b[0m \u001b[31m21.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.5/56.5 kB\u001b[0m \u001b[31m6.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m3.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.4/55.4 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.4/8.4 MB\u001b[0m \u001b[31m20.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "try:\n",
        "    from google.colab import drive\n",
        "    usingColab = True\n",
        "    print(\"Using Colab. Downloading necessary libraries...\")\n",
        "    # -- WINDOWS --\n",
        "    # !pip uninstall bitsandbytes bitsandbytes-windows\n",
        "    # !pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.2.post2-py3-none-win_amd64.whl # 0.41.2\n",
        "    # !pip install https://github.com/habibzadeh/bitsandbytes/releases/download/0.42.0_win_cuda_12.1/bitsandbytes-0.42.0-py39-none-any.whl # 0.42.0\n",
        "\n",
        "    # -- LINUX --\n",
        "    !pip install bitsandbytes -q\n",
        "\n",
        "    # Other simple libraries\n",
        "    !pip install accelerate sentencepiece peft pytorch_lightning scikit-learn wandb lightning-bolts langchain -q\n",
        "\n",
        "    # Update transformers to latest version\n",
        "    !pip install -U gdown transformers -q\n",
        "except:\n",
        "    usingColab = False\n",
        "    print(\"Not using Colab. Assuming all libraries have been already downloaded...\")\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mu86BnTiLIU_"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BUFzw8opew3M",
        "outputId": "eaea2ee9-0674-41a8-cd05-e88cb8c1c22e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/pl_bolts/__init__.py:11: FutureWarning: In the future `np.object` will be defined as the corresponding NumPy scalar.\n",
            "  if not hasattr(numpy, tp_name):\n",
            "/usr/local/lib/python3.10/dist-packages/pl_bolts/__init__.py:11: FutureWarning: In the future `np.bool` will be defined as the corresponding NumPy scalar.\n",
            "  if not hasattr(numpy, tp_name):\n",
            "/usr/local/lib/python3.10/dist-packages/pl_bolts/models/self_supervised/amdim/amdim_module.py:34: UnderReviewWarning: The feature generate_power_seq is currently marked under review. The compatibility with other Lightning projects is not guaranteed and API may change at any time. The API and functionality may change without warning in future releases. More details: https://lightning-bolts.readthedocs.io/en/latest/stability.html\n",
            "  \"lr_options\": generate_power_seq(LEARNING_RATE_CIFAR, 11),\n",
            "/usr/local/lib/python3.10/dist-packages/pl_bolts/models/self_supervised/amdim/amdim_module.py:92: UnderReviewWarning: The feature FeatureMapContrastiveTask is currently marked under review. The compatibility with other Lightning projects is not guaranteed and API may change at any time. The API and functionality may change without warning in future releases. More details: https://lightning-bolts.readthedocs.io/en/latest/stability.html\n",
            "  contrastive_task: Union[FeatureMapContrastiveTask] = FeatureMapContrastiveTask(\"01, 02, 11\"),\n",
            "/usr/local/lib/python3.10/dist-packages/pl_bolts/losses/self_supervised_learning.py:228: UnderReviewWarning: The feature AmdimNCELoss is currently marked under review. The compatibility with other Lightning projects is not guaranteed and API may change at any time. The API and functionality may change without warning in future releases. More details: https://lightning-bolts.readthedocs.io/en/latest/stability.html\n",
            "  self.nce_loss = AmdimNCELoss(tclip)\n"
          ]
        }
      ],
      "source": [
        "import gzip\n",
        "import json\n",
        "import random\n",
        "import shutil\n",
        "import tarfile\n",
        "import warnings\n",
        "from os.path import join as pathjoin\n",
        "\n",
        "import bitsandbytes as bnb\n",
        "import gdown\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import pl_bolts\n",
        "import pytorch_lightning as pl\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from peft import (LoraConfig, TaskType, get_peft_model,\n",
        "                  prepare_model_for_kbit_training)\n",
        "from sklearn.cluster import KMeans\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader, Dataset, Subset\n",
        "from tqdm import tqdm\n",
        "from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer,\n",
        "                          BitsAndBytesConfig, EncoderDecoderModel,\n",
        "                          SwitchTransformersForConditionalGeneration)\n",
        "from transformers.trainer_pt_utils import get_parameter_names\n",
        "\n",
        "import wandb\n",
        "\n",
        "warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')\n",
        "\n",
        "# Few fixes for Linux\n",
        "if(usingColab):\n",
        "    import locale\n",
        "    import os\n",
        "\n",
        "    def getpreferredencoding(do_setlocale = True):\n",
        "        return \"UTF-8\"\n",
        "    locale.getpreferredencoding = getpreferredencoding\n",
        "\n",
        "    # Java and CUDA\n",
        "    os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-11-openjdk-amd64\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85QLEQCUJfBl"
      },
      "source": [
        "## Check CUDA support"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "br6pL5UiJef3",
        "outputId": "fa55baaa-7fa1-4ea5-f1cf-390001b5819e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Cuda available: True\n",
            "GPU: Tesla T4\n",
            "Total memory: 14.7 GB\n",
            "===================================================\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
            "  and should_run_async(code)\n"
          ]
        }
      ],
      "source": [
        "if(torch.cuda.is_available()):\n",
        "    device = torch.device(\"cuda\")\n",
        "    print('Cuda available: {}'.format(torch.cuda.is_available()))\n",
        "    gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())\n",
        "    print(\"GPU: \" + gpu_name)\n",
        "    if('RTX 3060' in gpu_name):\n",
        "        print(\"\\t- Using RTX 3060. Setting MatMul precision to High.\")\n",
        "        # Faster, but less precise\n",
        "        torch.set_float32_matmul_precision(\"high\")\n",
        "    print(\"Total memory: {:.1f} GB\".format((float(torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)))))\n",
        "    print(\"===================================================\")\n",
        "else:\n",
        "    device = torch.device(\"cpu\")\n",
        "    print('Cuda not available, so using CPU. Please consider switching to a GPU runtime before running the notebook!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UWJVpGjqpJ5X"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "First of all configure the next cell (select the right student): if you are not listed in the list of the keys, add your name and congifure the path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u6RYi0NMpSk5",
        "outputId": "ae2a4ce3-429c-40d6-dddf-926bbb42c4aa"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using Professor's Notebook Setup.\n"
          ]
        }
      ],
      "source": [
        "# Student\n",
        "student = \"Professor\" # For Professor, please insert \"Professor\"\n",
        "\n",
        "# Special case for professor's notebook setup\n",
        "if(student == 'Professor'):\n",
        "  print(\"Using Professor's Notebook Setup.\")\n",
        "  msmarco_v2_doc_dir = \"/content\"\n",
        "  msmarco_v2_passage_dir = \"/content\"\n",
        "else:\n",
        "  # Project directory\n",
        "  proj_dict = {\"Emanuele\": [\"ColabNotebooks/DeepLearning/Progetto\", ''],\n",
        "              \"NN\": [\"DeepLearningProject/Progetto\", ''],\n",
        "              \"Gianmarco_Uni\": [\"Università/2. Magistrale/II ANNO/I° Semestre/Deep Learning/Progetto\", 'a61cd0e8be88d9f5059ac69c00e7fbb3c5cb33b5'], # Shortcut to Emanuele Project directory + WandB key\n",
        "              \"Gianmarco_Personal\": [\"DeepLearningProj/Progetto\", 'a61cd0e8be88d9f5059ac69c00e7fbb3c5cb33b5'],\n",
        "              }\n",
        "\n",
        "  msmarco_v2_doc_dir = f\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_doc\"\n",
        "  msmarco_v2_passage_dir = f\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_passage\"\n",
        "\n",
        "  # LOCAL VARIABLES!! Use only for local execution\n",
        "  if (not usingColab):\n",
        "    msmarco_v2_doc_dir = \"msmarco_v2_doc\"\n",
        "    msmarco_v2_passage_dir = \"msmarco_v2_passage\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tf2VWnoq8vOj"
      },
      "source": [
        "### Drive connection\n",
        "We put the Drive connection down here, since we have a special case for the Professor, where we don't want to use Drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-QgjxtgMb249"
      },
      "outputs": [],
      "source": [
        "# Import Drive if using Colab\n",
        "if(usingColab) and (student != 'Professor'):\n",
        "  drive.mount(\"/content/drive/\", force_remount=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gqrta7HRpPgF"
      },
      "source": [
        "### Download\n",
        "Official download of MSMarco Dataset from https://microsoft.github.io/msmarco/TREC-Deep-Learning.html\n",
        "- msmarco_v2_doc.tar = 32.27GB (Unpacked: 112GB)\n",
        "- msmarco_v2_passage.tar = 20.27GB"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Q0eLPHOb24-"
      },
      "outputs": [],
      "source": [
        "# Special case for the Professor. Only download the Chunk 00 and the Train bits.\n",
        "#\n",
        "# We download only the Chunk 00, since MSMarco is composed by 59 Chunks, and the total size is 112 GB.\n",
        "# Other than that, we will later use a portion of Chunk 00 to train the model, due to the limited computational resources.\n",
        "if(student == 'Professor'):\n",
        "    Chunk_URL = 'https://drive.google.com/uc?id=' + '1B8t6dxkDZ2DLomMN5CikmgTeerLiJgGU' + '&export=download&confirm=t'\n",
        "    Doc_Train_Queries_URL = 'https://drive.google.com/uc?id=' + '1Pn3oOQG1-6Y1zD01xgiPi0wgPKWT4WiT' + '&export=download&confirm=t'\n",
        "    Doc_Train_Qrels_URL = 'https://drive.google.com/uc?id=' + '1PvLVMMkSmB91BO29GOS1Xa0mhN0RxzVV' + '&export=download&confirm=t'\n",
        "    Doc_Train_Top100_URL = 'https://drive.google.com/uc?id=' + '1VgwtS-6s5lT4eKHE_uKEie_8zqmQJkca' + '&export=download&confirm=t'\n",
        "\n",
        "    gdown.download(Chunk_URL, quiet=False)\n",
        "    gdown.download(Doc_Train_Queries_URL, quiet=False)\n",
        "    gdown.download(Doc_Train_Qrels_URL, quiet=False)\n",
        "    gdown.download(Doc_Train_Top100_URL, quiet=False)\n",
        "\n",
        "    print(f\"Files downloaded to /content/\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nu-Es2CEpVzD"
      },
      "outputs": [],
      "source": [
        "# download_dataset = True\n",
        "#   - Starts the whole download process and store it in Google Drive. This takes approximately 2 hours between download/extracting.\n",
        "# download_dataset = False\n",
        "#   - It assumes that the Dataset has already been downloaded and extracted to Drive, so it simply loads from it\n",
        "download_dataset = False\n",
        "\n",
        "if(download_dataset and student != 'Professor'):\n",
        "      !wget --header \"X-Ms-Version: 2019-12-12\" https://msmarco.z22.web.core.windows.net/msmarcoranking/msmarco_v2_doc.tar\n",
        "      !wget --header \"X-Ms-Version: 2019-12-12\" https://msmarco.z22.web.core.windows.net/msmarcoranking/msmarco_v2_passage.tar\n",
        "      !wget --header \"X-Ms-Version: 2019-12-12\" https://msmarco.z22.web.core.windows.net/msmarcoranking/docv2_train_queries.tsv\n",
        "      !wget --header \"X-Ms-Version: 2019-12-12\" https://msmarco.z22.web.core.windows.net/msmarcoranking/docv2_train_qrels.tsv\n",
        "\n",
        "      # Copy the Dataset to Drive for easier access\n",
        "      shutil.copy(\"/content/msmarco_v2_doc.tar\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_doc.tar\")\n",
        "      shutil.copy(\"/content/msmarco_v2_passage.tar\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_passage.tar\")\n",
        "      shutil.copy(\"/content/docv2_train_queries.tsv\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/docv2_train_queries.tsv\")\n",
        "      shutil.copy(\"/content/docv2_train_qrels.tsv\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/docv2_train_qrels.tsv\")\n",
        "\n",
        "      # Extract both Documents and Passages into .gz files\n",
        "      tar = tarfile.open(F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_doc.tar\")\n",
        "      tar.extractall()\n",
        "      tar.close()\n",
        "      shutil.copytree(\"/content/msmarco_v2_doc\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_doc\")\n",
        "      !rm -rf \"/content/msmarco_v2_doc\" # Remove it from Colab otherwise disk will get full. Further loading will be done by Drive\n",
        "      # ------------------------- #\n",
        "      tar = tarfile.open(F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_passage.tar\")\n",
        "      tar.extractall()\n",
        "      tar.close()\n",
        "      shutil.copytree(\"/content/msmarco_v2_passage\",F\"/content/drive/MyDrive/{proj_dict[student][0]}/msmarco_v2_passage\")\n",
        "      !rm -rf \"/content/msmarco_v2_passage\" # Remove it from Colab otherwise disk will get full. Further loading will be done by Drive\n",
        "\n",
        "      # Extract Documents starting from .gz files\n",
        "      for doc in tqdm(os.listdir(msmarco_v2_doc_dir), desc=\"Extracting docs .gz files\"):\n",
        "        with gzip.open(pathjoin(msmarco_v2_doc_dir, doc), 'rb') as f_in:\n",
        "          with open(pathjoin(msmarco_v2_doc_dir, doc.split(\".\")[0]), 'wb') as f_out:\n",
        "              shutil.copyfileobj(f_in, f_out)\n",
        "              os.remove(pathjoin(msmarco_v2_doc_dir, doc))\n",
        "\n",
        "      # Extract Passages starting from .gz files\n",
        "      for doc in tqdm(os.listdir(msmarco_v2_passage_dir), desc=\"Extracting passages .gz files\"):\n",
        "        with gzip.open(pathjoin(msmarco_v2_passage_dir, doc), 'rb') as f_in:\n",
        "          with open(pathjoin(msmarco_v2_passage_dir, doc.split(\".\")[0]), 'wb') as f_out:\n",
        "              shutil.copyfileobj(f_in, f_out)\n",
        "              os.remove(pathjoin(msmarco_v2_passage_dir, doc))\n",
        "else:\n",
        "  if(usingColab) and student != 'Professor':\n",
        "    print(F\"Not downloading Dataset.\\nAssuming it is already present in: /content/drive/MyDrive/{proj_dict[student][0]}/\")\n",
        "  else:\n",
        "    print(F\"Not downloading Dataset.\\nAssuming it is already present in: {msmarco_v2_doc_dir}/\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mJ54NoQCprLs"
      },
      "source": [
        "### Link queries with documents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gEhuFqFmpr7R"
      },
      "outputs": [],
      "source": [
        "# Obtain queries\n",
        "if(student == 'Professor'):\n",
        "  docv2_train_queries = \"/content/docv2_train_queries.tsv\"\n",
        "  docv2_train_qrels = \"/content/docv2_train_qrels.tsv\"\n",
        "else:\n",
        "  docv2_train_queries = pathjoin(F\"/content/drive/MyDrive/{proj_dict[student][0]}\", \"docv2_train_queries.tsv\") # query id - query\n",
        "  docv2_train_qrels = pathjoin(F\"/content/drive/MyDrive/{proj_dict[student][0]}\", \"docv2_train_qrels.tsv\") # query id - ? - doc id - ?\n",
        "\n",
        "# LOCAL VARIABLES!! Use only for local execution\n",
        "if (not usingColab):\n",
        "  docv2_train_queries = \"docv2_train_queries.tsv\"\n",
        "  docv2_train_qrels = \"docv2_train_qrels.tsv\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rL6nQRFcy3t3"
      },
      "outputs": [],
      "source": [
        "df_docv2_train_queries = pd.read_csv(docv2_train_queries, sep='\\t', header=None, names=[\"Query_ID\", \"Query\"])\n",
        "df_docv2_train_qrels = pd.read_csv(docv2_train_qrels, sep='\\t', header=None, names=[\"Query_ID\", \"Iteration\", \"Doc_ID\", \"Relevance\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t897WdP8uTny"
      },
      "outputs": [],
      "source": [
        "# Merge the two dataframes on 'Query_ID'\n",
        "query_doc_df = pd.merge(df_docv2_train_queries, df_docv2_train_qrels, on='Query_ID')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ksoYRUeQHr-d"
      },
      "source": [
        "### Show a simple example of document retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EYv9cSW-KUgK"
      },
      "outputs": [],
      "source": [
        "df_docv2_train_queries.head(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jDkVyWMHKVJP"
      },
      "outputs": [],
      "source": [
        "df_docv2_train_qrels.head(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IldMqJN-KVqt"
      },
      "outputs": [],
      "source": [
        "query_doc_df.head(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vx4siltSHr-d"
      },
      "outputs": [],
      "source": [
        "def get_document(document_id):\n",
        "    (string1, string2, bundlenum, position) = document_id.split('_')\n",
        "    assert string1 == 'msmarco' and string2 == 'doc'\n",
        "\n",
        "    with open(f'{msmarco_v2_doc_dir}/msmarco_doc_{bundlenum}', 'rt', encoding='utf8') as in_fh:\n",
        "        in_fh.seek(int(position))\n",
        "        json_string = in_fh.readline()\n",
        "        document = json.loads(json_string)\n",
        "        assert document['docid'] == document_id\n",
        "        return document\n",
        "\n",
        "document = get_document('msmarco_doc_00_2823093')\n",
        "print(document.keys())\n",
        "print(\"-------------------------\")\n",
        "print(F\"Document ID: {document['docid']}\")\n",
        "print(\"-------------------------\")\n",
        "print(F\"URL: {document['url']}\")\n",
        "print(\"-------------------------\")\n",
        "print(F\"Title: {document['title']}\")\n",
        "print(\"-------------------------\")\n",
        "print(F\"Headings: {document['headings']}\")\n",
        "print(\"-------------------------\")\n",
        "print(F\"Body: {document['body']}\")\n",
        "print(\"## ---------------------------------- ##\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2V736BmieLdG"
      },
      "outputs": [],
      "source": [
        "# Specify the Query_ID\n",
        "target_query_id = 100015\n",
        "\n",
        "# Access rows where 'Query_ID' is equal to the target value\n",
        "matching_query = query_doc_df[query_doc_df['Query_ID'] == target_query_id]\n",
        "\n",
        "# Retrieve the Doc_ID from the matching rows\n",
        "matched_doc_ids = matching_query['Doc_ID'].tolist()\n",
        "\n",
        "# Print the results\n",
        "print(F\"Query_ID: {target_query_id}\")\n",
        "print(F\"Matched Doc_IDs: {matched_doc_ids}\")\n",
        "\n",
        "# Assuming it has found multiple documents related to this query\n",
        "print(\"\\n## ---------------------------------- ##\")\n",
        "for doc in matched_doc_ids:\n",
        "  document = get_document(F'{doc}')\n",
        "  print(F\"Document ID: {document['docid']}\")\n",
        "  print(F\"URL: {document['url']}\")\n",
        "  print(\"-------------------------\")\n",
        "  print(F\"Title: {document['title']}\")\n",
        "  print(\"-------------------------\")\n",
        "  print(F\"Headings: {document['headings']}\")\n",
        "  print(\"-------------------------\")\n",
        "  print(F\"Body: {document['body']}\")\n",
        "  print(\"## ---------------------------------- ##\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wyceQY3s53QY"
      },
      "source": [
        "### Documents/Doc_ID Representation Engines\n",
        "- For Documents, the only representation needed is on the Document itself. For this reason, we work on the \"Body\" of the Document.\n",
        "- For Doc_IDs, the only representation needed is the DocID, hence solutions are developed only for it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ajSj2ySR58mI"
      },
      "outputs": [],
      "source": [
        "class DocRepresentationEngine():\n",
        "    def __init__(self, strategy):\n",
        "        self.strategy = strategy\n",
        "        self.columns = ['Doc_ID', self.strategy]\n",
        "        self.doc_count = 0\n",
        "        assert strategy in ['direct_indexing', 'set_indexing', 'summarization', 'default'], 'Strategy not recognized! Select the correct one (direct_indexing, set_indexing, summarization, default).'\n",
        "\n",
        "        # If the strategy is set_indexing, download the stopwords from NLTK\n",
        "        if(strategy == 'set_indexing'):\n",
        "            import nltk\n",
        "            nltk.download('stopwords')\n",
        "            from nltk.corpus import stopwords\n",
        "            self.stop_words = set(stopwords.words('english'))\n",
        "\n",
        "    # This function will be called by the DataLoader, depending on the strategy chosen\n",
        "    def forward(self,chunk):\n",
        "      if(self.strategy == 'direct_indexing'):\n",
        "          return self.direct_indexing(chunk)\n",
        "      elif(self.strategy == 'set_indexing'):\n",
        "          return self.set_indexing(chunk)\n",
        "      elif(self.strategy == 'summarization'):\n",
        "          return self.default(chunk) # The way this is implemented is tricky. More on DatasetDSI_DocQuery\n",
        "        #   return self.summarization(chunk)\n",
        "      else:\n",
        "          return self.default(chunk)\n",
        "\n",
        "    # Default strategy.\n",
        "    # It will take the whole body and store them in the dataframe\n",
        "    def default(self,chunk):\n",
        "        full_dict = {}\n",
        "\n",
        "        for index, row in chunk.iterrows():\n",
        "            full_dict[index] = {}\n",
        "            doc_id = row['Doc_ID']\n",
        "            full_dict[index] = {'Doc_ID': doc_id, self.strategy: row['Body']}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index', columns=self.columns)\n",
        "        self.doc_count = len(df)\n",
        "        return df\n",
        "\n",
        "    # Direct Indexing strategy\n",
        "    # It will take the first L tokens of the body and store them in the dataframe\n",
        "    def direct_indexing(self,chunk):\n",
        "        L = 32\n",
        "        full_dict = {}\n",
        "\n",
        "        for index, row in chunk.iterrows():\n",
        "            full_dict[index] = {}\n",
        "            body_tokens = row['Body'].split(\" \")\n",
        "            body = body_tokens[:L]\n",
        "            body = \" \".join([word for word in body])\n",
        "            doc_id = row['Doc_ID']\n",
        "            full_dict[index] = {'Doc_ID': doc_id, self.strategy: body}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index', columns=self.columns)\n",
        "        self.doc_count = len(df)\n",
        "        return df\n",
        "\n",
        "    # Summarization strategy\n",
        "    # It will summarize the body in maximum 'max_length' tokens and store them in the dataframe\n",
        "    def summarization(self, chunk):\n",
        "        full_dict = {}\n",
        "\n",
        "        # Falcon with 4-Bit for faster inference\n",
        "        from transformers import pipeline, T5ForConditionalGeneration\n",
        "        model_id = \"Falconsai/text_summarization\"\n",
        "\n",
        "        bnb_config = BitsAndBytesConfig(\n",
        "            load_in_4bit=True,\n",
        "            bnb_4bit_use_double_quant=True,\n",
        "            bnb_4bit_quant_type=\"nf4\",\n",
        "            bnb_4bit_compute_dtype=torch.bfloat16\n",
        "        )\n",
        "\n",
        "        model_4bit = T5ForConditionalGeneration.from_pretrained(\n",
        "            model_id,\n",
        "            device_map=\"auto\",\n",
        "            quantization_config=bnb_config,\n",
        "        )\n",
        "\n",
        "        model_4bit.eval() # Set the model to inference mode\n",
        "\n",
        "        tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "\n",
        "        max_length = 32\n",
        "        limit_length = 20000\n",
        "\n",
        "        hugging_pipe = pipeline(\n",
        "            \"summarization\",\n",
        "            model=model_4bit,\n",
        "            tokenizer=tokenizer,\n",
        "            use_cache=True,\n",
        "            device_map=\"auto\",\n",
        "            max_length=max_length,\n",
        "            min_length=16,\n",
        "            do_sample=True,\n",
        "            top_k=10,\n",
        "            num_return_sequences=1,\n",
        "            eos_token_id=tokenizer.eos_token_id,\n",
        "            pad_token_id=tokenizer.eos_token_id,\n",
        "        )\n",
        "\n",
        "        for index, row in tqdm(chunk.iterrows(), total=len(chunk), desc=\"Summarizing documents\"):\n",
        "            full_dict[index] = {}\n",
        "            doc_id = row['Doc_ID']\n",
        "\n",
        "            # Why summarize if the document has less tokens than the summary length?\n",
        "            if len(row['Body'].split()) < max_length:\n",
        "                summary = row['Body']\n",
        "            else:\n",
        "                if len(row['Body']) > limit_length:\n",
        "                    row['Body'] = row['Body'][:limit_length]  # Limit the length to 20k characters\n",
        "                with torch.no_grad():\n",
        "                    summary = hugging_pipe(row['Body'])[0]['summary_text'] # Crashes because doc with index 60 has length = 745k. All others are < 20k\n",
        "\n",
        "            full_dict[index] = {'Doc_ID': doc_id, self.strategy: summary, chunk.columns[4]: row[chunk.columns[4]]}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index', columns=['Doc_ID', self.strategy, chunk.columns[4]])\n",
        "        self.doc_count = len(df)\n",
        "        del bnb_config, model_4bit, tokenizer, hugging_pipe\n",
        "        return df\n",
        "\n",
        "    # Set Indexing strategy\n",
        "    # It will remove stopwords and duplicates from the body and store them in the dataframe\n",
        "    def set_indexing(self,chunk):\n",
        "        full_dict = {}\n",
        "\n",
        "        for index, row in chunk.iterrows():\n",
        "            full_dict[index] = {}\n",
        "            body = set(row['Body'].split()) - self.stop_words\n",
        "            doc_id = row['Doc_ID']\n",
        "            full_dict[index] = {'Doc_ID': doc_id, self.strategy: body}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index', columns=self.columns)\n",
        "        self.doc_count = len(df)\n",
        "        return df"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "haIUDb7fb25C"
      },
      "outputs": [],
      "source": [
        "class DocIDRepresentationEngine():\n",
        "    def __init__(self, strategy, doc_count):\n",
        "        self.strategy = strategy\n",
        "        self.used_doc_ids = set()\n",
        "        self.columns = ['Doc_ID',self.strategy] # Do not confuse! Doc_ID is the MsMarco Original identifier, while self.strategy returns the doc_id generated by the algorithm\n",
        "        self.doc_count = doc_count # This got passed from the DocRepresentationEngine\n",
        "        assert strategy in ['unstructured_atomic', 'naively_structured', 'semantically_structured'], 'Strategy not recognized! Select the correct one (unstructured_atomic, naively_structured, semantically_structured).'\n",
        "\n",
        "    # This function will be called by the DataLoader, depending on the strategy chosen\n",
        "    def forward(self, chunk):\n",
        "      if(self.strategy == 'unstructured_atomic'):\n",
        "          return self.unstructred_atomic(chunk)\n",
        "      elif(self.strategy == 'naively_structured'):\n",
        "          return self.naively_structured(chunk)\n",
        "      else:\n",
        "          raise Exception(\"To implement\")\n",
        "          return self.semantically_structured(chunk)\n",
        "\n",
        "    # Needed for the Unstructured/Naively Structured Identifiers\n",
        "    def generate_unique_doc_id(self):\n",
        "        doc_id = random.randint(0, self.doc_count) # From 0 to the max number of possibile document\n",
        "\n",
        "        # If the doc_id is already used, generate a new one\n",
        "        while doc_id in self.used_doc_ids:\n",
        "            doc_id = random.randint(0, self.doc_count)\n",
        "\n",
        "        self.used_doc_ids.add(doc_id)\n",
        "        return doc_id\n",
        "\n",
        "    # The most naive way to represent documents is assign each an\n",
        "    # arbitrary (and possibly random) unique integer identifier.\n",
        "    def unstructred_atomic(self,chunk):\n",
        "        full_dict = {}\n",
        "\n",
        "        for index, row in chunk.iterrows():\n",
        "            full_dict[index] = {}\n",
        "\n",
        "            doc_id = str(index).zfill(len(str(self.doc_count)))  # Random unique indexing (as expressed in the paper) | Example: 0001728\n",
        "\n",
        "            full_dict[index] = {'Doc_ID': row['Doc_ID'], self.strategy: doc_id}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index', columns=self.columns)\n",
        "        return df\n",
        "\n",
        "    # Naively Structured Identifiers\n",
        "    # It will take the original Doc_ID and split it into tokens\n",
        "    def naively_structured(self,chunk):\n",
        "        full_dict = {}\n",
        "\n",
        "        for index, row in chunk.iterrows():\n",
        "            doc_id = self.generate_unique_doc_id() #                            | Example: 1728\n",
        "            doc_id_str = str(doc_id)  # Convert it into a string                | Example: '1728'\n",
        "            tokens = list(doc_id_str) # Get a tokenizable version of the Doc_ID | Example: ['1', '7', '2', '8']\n",
        "            full_dict[index] = {'Doc_ID': row['Doc_ID'], self.strategy: tokens}\n",
        "\n",
        "        df = pd.DataFrame.from_dict(full_dict, orient='index')\n",
        "        return df\n",
        "\n",
        "    # Needed for Semantically Structured Identifiers\n",
        "    def cluster_documents(self, document_embeddings, number_of_clusters=10):\n",
        "        kmeans = KMeans(n_clusters=number_of_clusters, random_state=2047315) # For reproducibility\n",
        "        clusters = kmeans.fit_predict(document_embeddings)\n",
        "        clustered_documents = [document_embeddings[clusters == i] for i in range(number_of_clusters)]\n",
        "        return clustered_documents\n",
        "\n",
        "    def generate_semantic_ids(self, cluster, max_docs_per_cluster):\n",
        "        # This function should be implemented in the future\n",
        "        raise Exception(\"To implement\")\n",
        "\n",
        "    # Needs to be properly tested!!\n",
        "    def semantically_structured(self):\n",
        "      number_of_clusters = 10\n",
        "      max_docs_per_cluster = 100\n",
        "\n",
        "      # Il problema qua è che bisogna avere, per prima cosa, gli embeddings di tutti i documenti.\n",
        "      # Significa che bisogna far passare un \"small 8-layer BERT model\" per tutti i body nel DataFrame\n",
        "      document_embeddings = None\n",
        "\n",
        "      clusters = self.cluster_documents(document_embeddings, number_of_clusters=number_of_clusters) # C_1:10 ← Cluster(X_1:N , k = 10)\n",
        "      doc_ids = []                                                                                  # J ← empty list\n",
        "\n",
        "      # As per pseudo-code\n",
        "      for i in range(10):                                                                           # for i = 0 to 9 do\n",
        "          current_cluster = [str(i)] * len(clusters[i])                                             # J_current ← [i] ∗ |C_i+1|\n",
        "\n",
        "          if len(clusters[i]) > max_docs_per_cluster:                                               # if |C_i+1| > c then\n",
        "              rest_cluster = self.generate_semantic_ids(clusters[i], max_docs_per_cluster)              # J_rest ← GENERATESEMANTICIDS(C_i+1)\n",
        "          else:                                                                                     # else\n",
        "              rest_cluster = [str(j) for j in range(len(clusters[i]))]                                  # J_rest ← [0, . . . , |C_i+1| − 1]\n",
        "\n",
        "          cluster_ids = [current + rest for current, rest in zip(current_cluster, rest_cluster)]    # J_cluster ←elementwiseStrConcat(J_current, J_rest)\n",
        "          doc_ids.extend(cluster_ids)                                                               # J ← J.appendElements(Jcluster)\n",
        "\n",
        "          # Manca il riordinamento: J ← reorderToOriginal(J, X1:N , C1:10)\n",
        "          # Significa che molto probabilmente bisogna riordinare i doc_ids basandosi\n",
        "          # su come sono arrivati gli embeddings iniziali\n",
        "\n",
        "      return doc_ids"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zFxmkZg3Zqyv"
      },
      "source": [
        "### Datasets\n",
        "\n",
        "This section is responsable to create the DataLoader in a simple manner, ready to be used for training.<br>\n",
        "The `DatasetDSI` Class has to cover all the combination from Document Representation and DocID Representation.<br>\n",
        "\n",
        "In general, we identify a single sample as:<br>\n",
        "`[ms_marco_original_doc_id, document_representation, doc_id_representation]`<br>\n",
        "\n",
        "More on `DatasetDSI_DocQuery` in the `DataLoaders` Chapter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YeGM6PBgay-2"
      },
      "outputs": [],
      "source": [
        "class DatasetDSI(Dataset):\n",
        "  '''\n",
        "    train_dataset_chunk: the single MSMarco file we want to process\n",
        "    doc_full_iterator: the iterator that returns a full rapresentation of a document\n",
        "    doc_rep_strategy: use to bootstrap the DocRepresentationEngine\n",
        "    doc_ids_rep_strategy: use to bootstrap the DocIDRepresentationEngine\n",
        "  '''\n",
        "  def __init__(self, train_dataset_chunk='00', doc_rep_strategy='direct_indexing', doc_ids_rep_strategy=\"unstructured_atomic\", MAX_DOCS = 1000):\n",
        "    self.train_dataset_chunk = train_dataset_chunk\n",
        "    self.MAX_DOCS = MAX_DOCS      # Max: 200000 | Default: 1000\n",
        "    self.df = self.generate_df()  # Get MAX_DOCS documents from the chunk\n",
        "\n",
        "    # Run the DocRepresentationEngine on the Chunk\n",
        "    self.doc_rep_engine = DocRepresentationEngine(doc_rep_strategy)\n",
        "    self.docs_rep = self.doc_rep_engine.forward(self.df)\n",
        "\n",
        "    # Run the DocIDRepresentationEngine on the Chunk\n",
        "    self.docID_rep_engine = DocIDRepresentationEngine(doc_ids_rep_strategy, self.doc_rep_engine.doc_count)\n",
        "    self.docs_id_rep = self.docID_rep_engine.forward(self.df)\n",
        "    print(f\"-- Considering as dataset only the first {self.MAX_DOCS} rows\")\n",
        "\n",
        "  # Generate the DataFrame from the MSMarco chunk chosen.\n",
        "  # Optionally we can limit the number of documents to be processed for faster testing\n",
        "  def generate_df(self):\n",
        "    dict_current_idx = {}\n",
        "    columns = [\"Doc_ID\", \"Title\", \"Body\"]\n",
        "    with open(f'{msmarco_v2_doc_dir}/msmarco_doc_{self.train_dataset_chunk}', 'rt', encoding='utf8') as in_fh:\n",
        "      i = 0\n",
        "      for line in tqdm(in_fh, desc=f\"Generating chunks for file msmarco_doc_{self.train_dataset_chunk}\"):\n",
        "        # TODO // Remove this if if we want to load the whole file\n",
        "        if i < self.MAX_DOCS:\n",
        "          line_dict = json.loads(line)\n",
        "\n",
        "          # Update the dictionary\n",
        "          dict_current_idx[str(i)] = {\n",
        "              \"Doc_ID\":  line_dict['docid'],\n",
        "              \"Title\": line_dict['title'],\n",
        "              \"Body\":  line_dict['body'],\n",
        "          }\n",
        "        else:\n",
        "          break\n",
        "\n",
        "        i += 1\n",
        "\n",
        "    df = pd.DataFrame.from_dict(\n",
        "        dict_current_idx, orient='index', columns=columns)\n",
        "    print(\"Chunk correctly generated. Processing DocIDs and Documents now...\")\n",
        "    return df\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.df)\n",
        "\n",
        "  # This function will be called by the DataLoader when we call a batch\n",
        "  def __getitem__(self, index):\n",
        "    sample = self.df.iloc[index]        # Get the row from the DataFrame (Mainly for retrieving original MSMarco Doc_ID)\n",
        "    doc_rep = self.docs_rep.iloc[index] # Get the Document representation based on the index\n",
        "    doc_id_rep = self.docs_id_rep.iloc[index] # Get the Document ID representation based on the index\n",
        "\n",
        "    # Return the sample as: Doc_ID, Doc_Representation, DocID_Representation\n",
        "    return {\n",
        "        \"Doc_ID\":  sample['Doc_ID'],\n",
        "        # \"Title\": sample['Title'],\n",
        "        # \"Body\":\t sample['Body'],\n",
        "        self.doc_rep_engine.strategy: doc_rep[self.doc_rep_engine.strategy],\n",
        "        self.docID_rep_engine.strategy: doc_id_rep[self.docID_rep_engine.strategy]\n",
        "    }\n",
        "\n",
        "  # Collate function for the DataLoader\n",
        "  # This gets called after the __getitem__ and is used to prepare the batch\n",
        "  # by applying some transformations to the doc_rep\n",
        "  def collate_fn(self, batch):\n",
        "    # body = [item[\"Body\"] for item in batch]\n",
        "    # title = [item[\"Title\"] for item in batch]\n",
        "    doc_id = [item[\"Doc_ID\"] for item in batch]\n",
        "\n",
        "    doc_rep = [item[self.doc_rep_engine.strategy] for item in batch]\n",
        "    doc_id_rep = [item[self.docID_rep_engine.strategy] for item in batch]\n",
        "\n",
        "    return {\n",
        "        \"Doc_ID\":  doc_id,\n",
        "        # \"Title\": title,\n",
        "        # \"Body\":\t body,\n",
        "        self.doc_rep_engine.strategy: doc_rep,\n",
        "        self.docID_rep_engine.strategy: doc_id_rep\n",
        "    }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VflefjdIb25C"
      },
      "outputs": [],
      "source": [
        "class DatasetDSI_DocQuery(Dataset):\n",
        "  '''This is the DSI Class that allows the creation of a dataframe that contains the queries (for retrieval task)'''\n",
        "  def __init__(self, train_dataset: DatasetDSI, query_doc_df: pd.DataFrame, strategy: str):\n",
        "    self.train_dataset_docs_rep = train_dataset.docs_rep      # columns ['Doc_ID', 'Doc Repres. Strategy']\n",
        "    self.train_dataset_docs_id = train_dataset.docs_id_rep    # columns ['Doc_ID', 'Doc ID Repres. Strategy']\n",
        "    self.query_doc_df = query_doc_df # Dataframe containing ALL the queries and ALL the documents (330k) # columns ['Query_ID', 'Query', 'Iteration', 'Doc_ID', 'Relevance']\n",
        "\n",
        "    # Merge dataframes\n",
        "    # We first merge the train_dataset_docs_rep with the train_dataset_docs_id on 'Doc_ID', hence having a Dataframe with columns ['Doc_ID', 'Doc Repres. Strategy', 'Doc ID Repres. Strategy']\n",
        "    self.filtered_docs_with_queries = pd.merge(self.train_dataset_docs_rep, self.train_dataset_docs_id, on='Doc_ID')\n",
        "\n",
        "    # Now we merge the filtered_docs_with_queries with the train_dataset.df on 'Doc_ID', hence having a Dataframe with columns ['Doc_ID', 'Title', 'Body', 'Doc Repres. Strategy', 'Doc ID Repres. Strategy']\n",
        "    self.filtered_docs_with_queries = pd.merge(train_dataset.df, self.filtered_docs_with_queries, on='Doc_ID')\n",
        "\n",
        "    # Finally, we want to filter the documents that have at least one query. We can do this by taking only the documents that are in the query_doc_df\n",
        "    self.filtered_docs_with_queries = self.filtered_docs_with_queries[self.filtered_docs_with_queries['Doc_ID'].isin(self.query_doc_df['Doc_ID'])] # Take only the docs with at least one query\n",
        "\n",
        "    # We use the summarization strategy in this brutal way due to the fact that if we did summarization\n",
        "    # on the whole Dataset, we had to summarize up to MAX_DOCS documents, which is not what we want.\n",
        "    #\n",
        "    # We want to only summarize the documents that have queries, which can be found in self.filtered_docs_with_queries.\n",
        "    if(strategy == 'summarization'):\n",
        "      summarizations = 'https://drive.google.com/uc?id=' + '1S_M69b2SAN2WyvudEKgeHBkMcU2vlerv' + '&export=download&confirm=t'\n",
        "\n",
        "      gdown.download(summarizations, \"summarization.csv\", quiet=False)\n",
        "      # self.filtered_docs_with_queries = train_dataset.doc_rep_engine.summarization(self.filtered_docs_with_queries)\n",
        "      # Read the summarization.csv file and specify the index column as the first one, then specify the type of the columns (all str)\n",
        "      # Read the CSV file using the custom function\n",
        "\n",
        "      self.filtered_docs_with_queries = pd.read_csv(\"summarization.csv\", sep=\",\", dtype={\n",
        "            'Doc_ID': 'string',\n",
        "            'summarization': 'string',\n",
        "            'unstructured_atomic': 'string',\n",
        "      })\n",
        "      self.filtered_docs_with_queries = self.filtered_docs_with_queries.apply(self.fix_row, column_name='idx', axis=1)\n",
        "      self.filtered_docs_with_queries.set_index('idx', inplace=True)\n",
        "    else:\n",
        "      # Drop Title and Body Column\n",
        "      self.filtered_docs_with_queries = self.filtered_docs_with_queries.drop(columns=['Title', 'Body'])\n",
        "\n",
        "    self.columns = self.filtered_docs_with_queries.columns.values.tolist()\n",
        "\n",
        "  # Define a function to read the CSV file\n",
        "  def fix_row(self, row: pd.Series, column_name: str) -> pd.Series:\n",
        "      value = row[column_name]\n",
        "      formated_value = str(value).split(',')\n",
        "      if len(formated_value) > 1:\n",
        "          i = 0\n",
        "          j = 0\n",
        "          # We now need to check if in the list there is a string that needs to be concatenated. We can check this if it starts with a \" or ends with a \".\n",
        "          # If it starts with a \" and ends with a \", then it means that the string has been split and we need to concatenate it.\n",
        "          # We can do this by checking if the first element of the list starts with a \" and the last element ends with a \".\n",
        "          for i in range(len(formated_value)):\n",
        "            if(formated_value[i].startswith('\"')): # Starts with a comma. Let's find the end in positions i+1, i+2, ...\n",
        "              for j in range(i+1, len(formated_value)):\n",
        "                if(formated_value[j].endswith('\"')): # Ends with a comma. We found the end of the string\n",
        "                  formated_value[i] = ','.join(formated_value[i:j+1]) # Join the string\n",
        "                  # We need to remove the elements from i+1 to j\n",
        "                  formated_value = formated_value[:i+1] + formated_value[j+1:]\n",
        "                  break\n",
        "              break\n",
        "\n",
        "          return pd.Series(dict(zip(row.keys(), formated_value)))\n",
        "      return row\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.filtered_docs_with_queries)\n",
        "\n",
        "  def __getitem__(self, index):\n",
        "    sample = self.filtered_docs_with_queries.iloc[index]\n",
        "\n",
        "    return {\n",
        "        self.columns[0]: sample['Doc_ID'],\n",
        "        self.columns[1]: sample[self.columns[1]],\n",
        "        self.columns[2]: sample[self.columns[2]],\n",
        "    }\n",
        "\n",
        "  def collate_fn(self, batch):\n",
        "    '''\n",
        "      Here we perform some preprocessing\n",
        "    '''\n",
        "    doc_id = [item[\"Doc_ID\"] for item in batch]\n",
        "    doc_rep = [item[self.columns[1]] for item in batch]\n",
        "    doc_id_rep = [item[self.columns[2]] for item in batch]\n",
        "\n",
        "    return {\n",
        "        \"Doc_ID\":  doc_id,\n",
        "        self.columns[1]: doc_rep,\n",
        "        self.columns[2]: doc_id_rep\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9_cXx_8Fm9Fq"
      },
      "source": [
        "### DataLoaders\n",
        "We have two main Datasets:\n",
        "- `DatasetDSI`,  which contains the full set of documents (from the chunk chosen)\n",
        "- `DatasetDSI_DocQuery`, which includes documents associated with queries for retrieval tasks.\n",
        "\n",
        "In MSMarco, not all documents have a query linked to them, so we can't properly perform the retrieval task, hence the creation of the second Dataset.\n",
        "\n",
        "When training, we want a dataset that combines the complete set of documents along with retrieval examples.<br>\n",
        "To achieve this, the model includes the full `DatasetDSI` and add a portion (80%) of retrieval examples from `DatasetDSI_DocQuery`.<br>\n",
        "\n",
        "Train Dataset:\n",
        "  - Full `DatasetDSI`\n",
        "  - 0.8 from `DatasetDSI_DocQuery`\n",
        "  \n",
        "Validation Dataset:\n",
        "  - 0.1 from `DatasetDSI_DocQuery`\n",
        "  \n",
        "Test Dataset:\n",
        "  - 0.1 from `DatasetDSI_DocQuery`\n",
        "\n",
        "In this way, we ensure that the model never performs retrieval on Validation and Test Documents, as it is limited to the 0.8 portion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WBrFhp3SHr-f"
      },
      "outputs": [],
      "source": [
        "dsi_config = {\n",
        "    # ---- GENERAL CONFIGURATIONS ----\n",
        "    \"seed\": 42,                                                     # Seed for reproducibility\n",
        "    \"MAX_LENGTH\": 256,                                              # Max length for the Tokenizer when tokenizing the input\n",
        "    \"MAX_DOCS\": 10000,                                              # Max number of documents to consider\n",
        "    \"batch_size\": 4,                                                # Batch size | Default 4 otherwise it will crash\n",
        "    # ---- TRAINING CONFIGURATIONS ----\n",
        "    \"train_dataset_chunk\": \"00\",                                    # The single MSMarco file we want to process\n",
        "    \"document_representation_strategy\": \"direct_indexing\",            # Strategy for the document representation | \"direct_indexing\", \"set_indexing\", \"inverted_indexing\", \"summarization\", \"default\"\n",
        "                                                                    #   \"summarization\" only works when inserting MAX_DOCS = 10000\n",
        "    \"document_id_representation_strategy\": \"unstructured_atomic\",   # Strategy for the document ID representation | \"unstructured_atomic\", \"naively_structured\", \"semantically_structured\"\n",
        "    \"training_type\": \"generative\",                                  # Training type | \"discriminative\" / \"generative\"\n",
        "    \"enable_multitask_prompting\" : True,                            # If True, it will enable the multitask prompting | \"Generate a doc_id etc..\"\n",
        "    \"train_indexing_retrieval_ratio\": 0.3,                          # Ratio between indexing and retrieval tasks | Default: 0.3 (70% for indexing, 30% for retrieval)\n",
        "    # ---- W&B CONFIGURATIONS ----\n",
        "    \"wandb_configs\": {\"enable\": False, \"group_id\": \"v0.9.3\"},        # WandB configurations\n",
        "    \"student\": student,                                             # Student name\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5rQxiLojVC26",
        "outputId": "8dcd2862-4e32-4b03-dc17-48bb6833ecaf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Document Representation Strategy: direct_indexing\n",
            "Document ID Representation Strategy: unstructured_atomic\n",
            "Train Dataset Chunk: 00\n",
            "-----------------------------------\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Generating chunks for file msmarco_doc_00: 0it [00:00, ?it/s]"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Generating chunks for file msmarco_doc_00: 200000it [00:22, 8707.35it/s] \n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Chunk correctly generated. Processing DocIDs and Documents now...\n",
            "-- Considering as dataset only the first 200000 rows\n"
          ]
        }
      ],
      "source": [
        "doc_rep_strategy = dsi_config[\"document_representation_strategy\"]\n",
        "doc_ids_rep_strategy = dsi_config[\"document_id_representation_strategy\"]\n",
        "train_dataset_chunk = dsi_config[\"train_dataset_chunk\"]\n",
        "\n",
        "print(\"Document Representation Strategy: {}\".format(doc_rep_strategy))\n",
        "print(\"Document ID Representation Strategy: {}\".format(doc_ids_rep_strategy))\n",
        "print(\"Train Dataset Chunk: {}\".format(train_dataset_chunk))\n",
        "print(\"-----------------------------------\")\n",
        "\n",
        "# This contains ALL the documents.\n",
        "train_dataset = DatasetDSI(train_dataset_chunk=train_dataset_chunk, doc_rep_strategy=doc_rep_strategy, doc_ids_rep_strategy=doc_ids_rep_strategy, MAX_DOCS=dsi_config['MAX_DOCS'])\n",
        "\n",
        "# This filters out the Documents which have a query(ies), suitable for Retrieval Task.\n",
        "train_docs_with_queries_dataset = DatasetDSI_DocQuery(train_dataset, query_doc_df, dsi_config['document_representation_strategy'])\n",
        "\n",
        "# This is the number of labels we need in order to generate a Linear Layer for the model\n",
        "# when the training type is \"discriminative\".\n",
        "dsi_config['labels_count'] = len(train_docs_with_queries_dataset.filtered_docs_with_queries[doc_ids_rep_strategy])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8FKJ_xm4m9Fq"
      },
      "outputs": [],
      "source": [
        "# -------------- RETRIEVAL SPLIT -------------- #\n",
        "np.random.seed(dsi_config[\"seed\"])\n",
        "\n",
        "# Define the ratios for validation and test datasets\n",
        "train_ratio = 0.8       # 80% of retrieval Docs for training\n",
        "validation_ratio = 0.1  # 10% of retrieval Docs for validation\n",
        "test_ratio = 0.1        # 10% of retrieval Docs for testing\n",
        "\n",
        "# Calculate the number of samples for validation and test datasets\n",
        "num_samples = len(train_docs_with_queries_dataset)\n",
        "num_train_samples = int(train_ratio * num_samples)\n",
        "num_val_samples = int(validation_ratio * num_samples)\n",
        "num_test_samples = int(test_ratio * num_samples)\n",
        "\n",
        "# Create an array of indices for the dataset\n",
        "indices = np.arange(num_samples)\n",
        "np.random.shuffle(indices) # Shuffle the indices\n",
        "\n",
        "# Split the shuffled indices into validation and test indices\n",
        "train_indices = indices[:num_train_samples]\n",
        "val_indices = indices[num_train_samples:num_train_samples + num_val_samples]\n",
        "test_indices = indices[num_val_samples:num_val_samples + num_test_samples]\n",
        "\n",
        "# Create subsets of the dataset for validation and test\n",
        "train_subset_with_queries_filtered = Subset(train_docs_with_queries_dataset, train_indices)\n",
        "val_subset = Subset(train_docs_with_queries_dataset, val_indices)\n",
        "test_subset = Subset(train_docs_with_queries_dataset, test_indices)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OeN8P5arLmar",
        "outputId": "eaece50c-f033-4f83-cec5-bbf2aab8276c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train Dataloader length: 117\n",
            "Validation Dataloader length: 15\n",
            "Test Dataloader length: 15\n",
            "------------------------------------\n",
            "Train Docs with Queries Dataloader length: 35\n"
          ]
        }
      ],
      "source": [
        "# Define number of workers for dataloaders\n",
        "#num_workers = int(os.cpu_count() / 2) if os.cpu_count() > 2 else 1 # Not setting maximum number of workers to avoid overloading the system\n",
        "num_workers = 0\n",
        "\n",
        "# Create dataloaders for training, validation, and test sets\n",
        "train_dataloader = DataLoader(train_subset_with_queries_filtered, batch_size=dsi_config['batch_size'], shuffle=True, collate_fn=train_docs_with_queries_dataset.collate_fn, num_workers=num_workers)\n",
        "\n",
        "# Filter the train dataloader to have a smaller subset for Retrieval Task\n",
        "# [---------------------------------|---]\n",
        "# [                I                | R ]\n",
        "batches_to_keep = int(len(train_dataloader) * dsi_config[\"train_indexing_retrieval_ratio\"])\n",
        "np.random.seed(dsi_config[\"seed\"])\n",
        "np.random.shuffle(train_indices)\n",
        "train_subset_with_queries_filtered = Subset(train_docs_with_queries_dataset, train_indices[:batches_to_keep*dsi_config['batch_size']]) # This will be passed to the model\n",
        "train_dataloader_filtered = DataLoader(train_subset_with_queries_filtered, batch_size=dsi_config['batch_size'], shuffle=True, collate_fn=train_docs_with_queries_dataset.collate_fn, num_workers=num_workers)\n",
        "\n",
        "val_dataloader = DataLoader(val_subset, batch_size=dsi_config['batch_size'], shuffle=False, collate_fn=train_docs_with_queries_dataset.collate_fn, num_workers=num_workers)\n",
        "test_dataloader = DataLoader(test_subset, batch_size=dsi_config['batch_size'], shuffle=False, collate_fn=train_docs_with_queries_dataset.collate_fn, num_workers=num_workers)\n",
        "\n",
        "print(f\"Train Dataloader length: {len(train_dataloader)}\")\n",
        "print(f\"Validation Dataloader length: {len(val_dataloader)}\")\n",
        "print(f\"Test Dataloader length: {len(test_dataloader)}\")\n",
        "print(\"------------------------------------\")\n",
        "print(f\"Train Docs with Queries Dataloader length: {len(train_dataloader_filtered)}\")\n",
        "\n",
        "# Update some configurations\n",
        "dsi_config[\"train_begin_retrieval\"] = len(train_dataloader) - len(train_dataloader_filtered)    # The index where the retrieval task begins | Dynamic\n",
        "dsi_config[\"token_tokenization_max_length\"] = train_dataset.docID_rep_engine.doc_count          # The max length for the tokenization. Used by tokenizing the labels and the generate() function\n",
        "\n",
        "if(dsi_config[\"enable_multitask_prompting\"]):\n",
        "    dsi_config[\"multitask_prompting_indexing\"] = f\"Generate a document identifier with {len(str(train_dataset.MAX_DOCS))} digits between 0 and 9 for the following document:\"\n",
        "    dsi_config[\"multitask_prompting_retrieval\"] = f\"Generate a document identifier with {len(str(train_dataset.MAX_DOCS))} digits between 0 and 9 as response to the following query:\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEGUoJIdLQRI"
      },
      "source": [
        "## Model\n",
        "- Quantization: https://huggingface.co/docs/transformers/quantization#8-bit <br>\n",
        "- Notebook on FineTuning using Lora & BitsAndBytes: https://github.com/huggingface/peft/blob/main/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb\n",
        "\n",
        "BitsAndBytes is needed to load up large models into small RAM systems through model quantization to 8 or 4-bits precision.<br>\n",
        "However, after a model is quantized it isn’t typically further trained for downstream tasks because training can be unstable due to the lower precision of the weights and activations.<br>\n",
        "But since PEFT methods only add extra trainable parameters, this allows you to train a quantized model with a PEFT adapter on top!<br>\n",
        "Combining quantization with PEFT can be a good strategy for training even the largest models on a single GPU. For example, QLoRA is a method that quantizes a model to 4/8-bits and then trains it with LoRA. <br>\n",
        "\n",
        "- LoRa = One would keep the base model in 32 or 16 bits in memory, and then train the parameter weights.\n",
        "- QLoRa = Apply LoRa to a quantized model (like a 4-bit model)\n",
        "\n",
        "The various models tested can be found in:\n",
        "- SwitchTransformers = https://huggingface.co/collections/google/switch-transformers-release-6548c35c6507968374b56d1f\n",
        "- Flan-T5 = https://huggingface.co/collections/google/flan-t5-release-65005c39e3201fff885e22fb\n",
        "\n",
        "| Model Name               | All Parameters | 8-Bit Parameters | QLoRa Trainable Params (Percentage) | Hidden Size Encoder  | Target modules                        |\n",
        "|--------------------------|----------------|------------------|-------------------------------------|----------------------|---------------------------------------|\n",
        "| 'google/switch-base-8'   | 619M           | 24.722.688\t   | 1,327,104 (0.0214%)                 | 768                  | target_modules=[\"q\", \"k\", \"v\"]        |\n",
        "| 'flan-t5-base'           | 248M           | /                | 1,327,104 (0.0533%)         \t     | 768                  | target_modules=[\"q\", \"k\", \"v\"]        |\n",
        "| 'bert-base-uncased'      | 139M           | /                | 1,327,104 (0.0929%)         \t     | 768                  | target_modules=[\"q\", \"k\", \"v\"]        |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GoKHVwm3LW69"
      },
      "source": [
        "### Foundation model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AjM5DKWcfLf_"
      },
      "outputs": [],
      "source": [
        "class FoundationModel(pl.LightningModule):\n",
        "    def __init__(self, model_id=\"google/flan-t5-base\", encoder_id='bert-base-uncased', decoder_id='bert-base-uncased', MAX_LENGTH=512, fine_tuning=False):\n",
        "        super().__init__()\n",
        "        # Misc variables\n",
        "        self.MAX_LENGTH = MAX_LENGTH\n",
        "\n",
        "        if(not fine_tuning):\n",
        "            # We first define the quantization configuration we want to use\n",
        "            # Here we use the Bits and Bytes configuration with 8-bit quantization\n",
        "            quantization_config = BitsAndBytesConfig(\n",
        "                load_in_8bit=True,                      # Quantize the model to 8-bits when you load it\n",
        "                bnb_8bit_use_double_quant=True,         # To use a nested quantization scheme to quantize the already quantized weights\n",
        "                bnb_8bit_quant_type=\"nf4\",              # To use a special 8-bit data type for weights initialized from a normal distribution\n",
        "                bnb_8bit_compute_dtype=torch.bfloat16,  # To use bfloat16 for faster computation\n",
        "            )\n",
        "\n",
        "            # Now that the quantized config is ready, let’s set up a configuration for further shrink of the parameters.\n",
        "            # We will use the LoraConfig class from the PEFT library.\n",
        "            lora_config = LoraConfig(\n",
        "                task_type=TaskType.SEQ_2_SEQ_LM,        # We are using the model for a seq2seq task\n",
        "                inference_mode=False,\n",
        "                r=8,                                    # Lora attention dimension | Default: 8\n",
        "                lora_alpha=32,                          # The alpha parameter for Lora scaling | Default: 8\n",
        "                lora_dropout=0.1,                       # The dropout probability for Lora layers | Default: 0.0\n",
        "                target_modules=[\"q\", \"k\", \"v\"],         # The names of the modules to apply Lora to | Can also be [\"q\", \"k\", \"v\", \"o\", \"wi\", \"wo\"] if we want to include weights_in / weights_out of FFN\n",
        "            )\n",
        "        else:\n",
        "            quantization_config = None\n",
        "            lora_config = None\n",
        "\n",
        "        # Load the model and tokenizer\n",
        "        if(model_id != 'bert-base-uncased') and (encoder_id == None) and (decoder_id == None):\n",
        "            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id, device_map=\"auto\", use_cache=False, quantization_config=quantization_config)\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(model_id, model_max_length=self.MAX_LENGTH) # Set reasonable default for models without max length\n",
        "            self.is_encoder_decoder = False\n",
        "            print(\"Using AutoModelForSeq2SeqLM\")\n",
        "        else:\n",
        "            self.model = EncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id, quantization_config=quantization_config, tie_encoder_decoder=True)\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(encoder_id, model_max_length=self.MAX_LENGTH) # Set reasonable default for models without max length\n",
        "            self.is_encoder_decoder = True\n",
        "            # Configurations of the EncoderDecoderModel\n",
        "            self.model.config.decoder_start_token_id = self.tokenizer.cls_token_id\n",
        "            self.model.config.eos_token_id = self.tokenizer.sep_token_id\n",
        "            self.model.config.pad_token_id = self.tokenizer.pad_token_id\n",
        "            self.model.config.vocab_size = self.model.config.encoder.vocab_size\n",
        "            lora_config.target_modules = {'query', 'key', 'value'}\n",
        "            print(\"Using EncoderDecoderModel, hence bert-base-uncased.\")\n",
        "\n",
        "        if(fine_tuning):\n",
        "            # Set requires_grad=True for the last 4 layers of the encoder and decoder\n",
        "            for name, param in self.model.named_parameters():\n",
        "                if name.startswith(\"encoder.block\") and int(name.split(\".\")[2]) >= 8:\n",
        "                    param.requires_grad = True\n",
        "                elif name.startswith(\"decoder.block\") and int(name.split(\".\")[2]) >= 8:\n",
        "                    param.requires_grad = True\n",
        "                else:\n",
        "                    param.requires_grad = False\n",
        "\n",
        "        # Function to preprocess the quantized model for training.\n",
        "        if(quantization_config is not None):\n",
        "            self.model = prepare_model_for_kbit_training(self.model, use_gradient_checkpointing=False)\n",
        "\n",
        "        # Should fix: UserWarning: None of the inputs have requires_grad=True. Gradients will be None.\n",
        "        # Should fix: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
        "        if hasattr(self.model, \"enable_input_require_grads\"):\n",
        "            self.model.enable_input_require_grads()\n",
        "        else:\n",
        "            def make_inputs_require_grad(module, input, output):\n",
        "                output.requires_grad_(True)\n",
        "\n",
        "            self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n",
        "\n",
        "        if(lora_config is not None):\n",
        "            self.model = get_peft_model(self.model, lora_config)\n",
        "            print(\"Encoder info:\")\n",
        "            self.model.print_trainable_parameters()\n",
        "        else:\n",
        "            print(f\"Encoder info: {self.get_n_trainable_parameters()} trainable parameters\")\n",
        "\n",
        "    # This is no more useful, the forward if this model is called directly in the DSI model\n",
        "    def forward(self, input):\n",
        "        pass\n",
        "\n",
        "    # Print number of trainable parameters\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bGeTZvJ6Hr-g"
      },
      "source": [
        "### Custom Encoder-Decoder\n",
        "Here we use SwitchTransformer as an Encoder and FlanT5 as Decoder, both instantiated with QLoRa.<br>\n",
        "\n",
        "The mechanism follows this scheme:\n",
        "- Forward pass in Encoder for document representation\n",
        "- Forward pass in Decoder with loss computation\n",
        "- Backpropagation of the loss\n",
        "\n",
        "This is done for both indexing and retrival tasks, but in the latter, the Query is used for the Forward Pass of the Encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bnqJLCXbHr-g"
      },
      "source": [
        "#### Encoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xcBi04h9Hr-g"
      },
      "outputs": [],
      "source": [
        "class SwitchTransformer(pl.LightningModule):\n",
        "    def __init__(self, model_id=\"google/switch-base-8\", MAX_LENGTH=512):\n",
        "        super().__init__()\n",
        "        # Misc variables\n",
        "        self.MAX_LENGTH = MAX_LENGTH\n",
        "\n",
        "        # We first define the quantization configuration we want to use\n",
        "        # Here we use the Bits and Bytes configuration with 8-bit quantization\n",
        "        quantization_config = BitsAndBytesConfig(\n",
        "            load_in_8bit=True,                      # Quantize the model to 8-bits when you load it\n",
        "            bnb_8bit_use_double_quant=True,         # To use a nested quantization scheme to quantize the already quantized weights\n",
        "            bnb_8bit_quant_type=\"nf4\",              # To use a special 8-bit data type for weights initialized from a normal distribution\n",
        "            bnb_8bit_compute_dtype=torch.bfloat16,  # To use bfloat16 for faster computation\n",
        "        )\n",
        "\n",
        "        # Now that the quantized config is ready, let’s set up a configuration for further shrink of the parameters.\n",
        "        # We will use the LoraConfig class from the PEFT library.\n",
        "        lora_config = LoraConfig(\n",
        "            task_type=TaskType.SEQ_2_SEQ_LM,        # We are using the model for a seq2seq task\n",
        "            inference_mode=False,\n",
        "            r=8,                                    # Lora attention dimension | Default: 8\n",
        "            lora_alpha=32,                          # The alpha parameter for Lora scaling | Default: 8\n",
        "            lora_dropout=0.1,                       # The dropout probability for Lora layers | Default: 0.0\n",
        "            target_modules=[\"q\", \"k\", \"v\"],         # The names of the modules to apply Lora to | Can also be [\"q\", \"k\", \"v\", \"o\", \"wi\", \"wo\"] if we want to include weights_in / weights_out of FFN\n",
        "        )\n",
        "\n",
        "        # Load up Switch Transformer\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(model_id, model_max_length=self.MAX_LENGTH) # Set reasonable default for models without max length\n",
        "        self.model = SwitchTransformersForConditionalGeneration.from_pretrained(model_id, device_map=\"auto\", use_cache=False, quantization_config=quantization_config)\n",
        "\n",
        "        # Function to preprocess the quantized model for training.\n",
        "        self.model = prepare_model_for_kbit_training(self.model, use_gradient_checkpointing=False)\n",
        "\n",
        "        # Should fix: UserWarning: None of the inputs have requires_grad=True. Gradients will be None.\n",
        "        # Should fix: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
        "        if hasattr(self.model, \"enable_input_require_grads\"):\n",
        "            self.model.enable_input_require_grads()\n",
        "        else:\n",
        "            def make_inputs_require_grad(module, input, output):\n",
        "                output.requires_grad_(True)\n",
        "\n",
        "            self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n",
        "\n",
        "        self.model = get_peft_model(self.model, lora_config)\n",
        "        print(\"Encoder info:\")\n",
        "        self.model.print_trainable_parameters()\n",
        "\n",
        "    # Forward pass for Encoder.\n",
        "    # input = Documents (or Queries) that needs to be tokenized and forwarded through the model\n",
        "    def forward(self, input):\n",
        "        docs_input = self.tokenizer(input, add_special_tokens=False, return_tensors='pt', padding='max_length', truncation=True) # [batch_size, max_length_tokenizer] -> [2, 256]\n",
        "        # inputs = self.tokenizer(input, add_special_tokens=False, return_tensors='pt', max_length=64, padding=\"max_length\", truncation=True)\n",
        "        docs_input[\"input_ids\"].to(device)\n",
        "        docs_input[\"attention_mask\"].to(device)\n",
        "\n",
        "        # Calling the whole model, but effectively retain the Encoder last hidden states\n",
        "        outputs = self.model(**docs_input, decoder_input_ids=torch.zeros_like(docs_input[\"input_ids\"], device=device), output_hidden_states=True)\n",
        "\n",
        "        last_hidden_states = outputs.encoder_last_hidden_state\n",
        "        encoder_hidden_states = outputs.encoder_hidden_states\n",
        "\n",
        "        return (last_hidden_states, encoder_hidden_states), docs_input, outputs\n",
        "\n",
        "    # Print number of trainable parameters\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67xlJwOQHr-k"
      },
      "source": [
        "#### Decoder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PpQaQ32DHr-k"
      },
      "outputs": [],
      "source": [
        "class FlanT5(pl.LightningModule):\n",
        "    def __init__(self, model_id=\"google/flan-t5-base\", MAX_LENGTH=512):\n",
        "        super().__init__()\n",
        "        # Misc variables\n",
        "        self.MAX_LENGTH = MAX_LENGTH\n",
        "\n",
        "        # We first define the quantization configuration we want to use\n",
        "        # Here we use the Bits and Bytes configuration with 8-bit quantization\n",
        "        quantization_config = BitsAndBytesConfig(\n",
        "            load_in_8bit=True,                      # Quantize the model to 8-bits when you load it\n",
        "            bnb_8bit_use_double_quant=True,         # To use a nested quantization scheme to quantize the already quantized weights\n",
        "            bnb_8bit_quant_type=\"nf4\",              # To use a special 8-bit data type for weights initialized from a normal distribution\n",
        "            bnb_8bit_compute_dtype=torch.bfloat16,  # To use bfloat16 for faster computation\n",
        "        )\n",
        "\n",
        "        # Now that the quantized config is ready, let’s set up a configuration for further shrink of the parameters.\n",
        "        # We will use the LoraConfig class from the PEFT library.\n",
        "        lora_config = LoraConfig(\n",
        "            task_type=TaskType.SEQ_2_SEQ_LM,        # We are using the model for a seq2seq task\n",
        "            inference_mode=False,\n",
        "            r=8,                                    # Lora attention dimension | Default: 8\n",
        "            lora_alpha=32,                          # The alpha parameter for Lora scaling | Default: 8\n",
        "            lora_dropout=0.1,                       # The dropout probability for Lora layers | Default: 0.0\n",
        "            target_modules=[\"q\", \"k\", \"v\"],         # The names of the modules to apply Lora to | Can also be [\"q\", \"k\", \"v\", \"o\", \"wi\", \"wo\"] if we want to include weights_in / weights_out of FFN\n",
        "        )\n",
        "\n",
        "        # Standard call for initializing FlanT5\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(model_id, model_max_length=self.MAX_LENGTH)\n",
        "        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id, device_map=\"auto\", use_cache=False, quantization_config=quantization_config)\n",
        "\n",
        "        # Function to preprocess the quantized model for training.\n",
        "        self.model = prepare_model_for_kbit_training(self.model, use_gradient_checkpointing=False)\n",
        "\n",
        "        # Should fix: UserWarning: None of the inputs have requires_grad=True. Gradients will be None.\n",
        "        # Should fix: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
        "        if hasattr(self.model, \"enable_input_require_grads\"):\n",
        "            self.model.enable_input_require_grads()\n",
        "        else:\n",
        "            def make_inputs_require_grad(module, input, output):\n",
        "                output.requires_grad_(True)\n",
        "\n",
        "            self.model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)\n",
        "\n",
        "        self.model = get_peft_model(self.model, lora_config)\n",
        "        print(\"Decoder info:\")\n",
        "        self.model.print_trainable_parameters()\n",
        "\n",
        "    # Forward pass for Decoder.\n",
        "    #   - embeddings        : embeddings arriving from the Encoder Module\n",
        "    #   - decoder_input_ids : input_ids from the Encoder's Tokenizer (for Doc & Query)\n",
        "    #   - labels            : Tensor of doc_ids tokenized to compute the loss\n",
        "    def forward(self, embeddings, decoder_input_ids, decoder_attention_mask, labels):\n",
        "\n",
        "      # https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py#L1709 | We can skip directly to Decoder by passing encoder_outputs.\n",
        "      # encoder_outputs is a tuple composed of: (last_hidden_state, optional: hidden_states, optional: attentions)\n",
        "      #     -> last_hidden_state : Shape (batch_size, sequence_length, hidden_size)\n",
        "      #        is a sequence of hidden states at the output of the last layer of the encoder.\n",
        "      #        Used in the cross-attention of the decoder.\n",
        "      # In this case, we would pass the last_hidden_state of the encoder to the decoder.\n",
        "      return self.model(\n",
        "          decoder_input_ids=decoder_input_ids,              # Use decoder_input_ids for the decoder input | Indices of decoder input sequence tokens in the vocabulary.\n",
        "          decoder_attention_mask=decoder_attention_mask,\n",
        "          encoder_outputs=embeddings,                       # Pass embeddings coming from Encoder as inputs_embeds for the decoder\n",
        "          labels=labels                                     # Labels for computing the sequence classification/regression loss | All labels set to -100 are ignored (masked), the loss is only computed for labels in [0, ..., model.vocab_size] # Default value = 32128\n",
        "      )\n",
        "\n",
        "    # Print number of trainable parameters\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nMsbPSrH2iYv"
      },
      "source": [
        "### Configs & WandB support"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NslC8o0u2iYv",
        "outputId": "e7c9c281-7a47-443c-e7b3-b7ac7fdb3f4f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mslimshadys\u001b[0m (\u001b[33msapienza_ml_2022_23\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: C:\\Users\\Gianmarco\\.netrc\n"
          ]
        }
      ],
      "source": [
        "# Recall that:\n",
        "# proj_dict[student][0] -> Project directory\n",
        "# proj_dict[student][1] -> WandB key\n",
        "if(dsi_config[\"wandb_configs\"][\"enable\"]) and student != 'Professor':\n",
        "    wandb.login(key = proj_dict[student][1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9aV-yUcub25Q"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "This method edits a configuration parameter given a run_id.\n",
        "\n",
        ":param run_id: It's the ID of the run (13qvc9uo, 83cvs1ff, etc.)\n",
        ":param project_name: The name of the WANDB project\n",
        ":param parameter: The name of the parameter to change\n",
        ":param new_value: The value to give to the parameter passed above\n",
        "\n",
        ":return: Edit the needed parameter and returns\n",
        "\n",
        "@ Gianmarco Scarano\n",
        "\"\"\"\n",
        "def configChanger(run_id='None', project_name='None', parameter='None', new_value = 'None'):\n",
        "\n",
        "    if run_id == 'None' or project_name == 'None' or parameter == 'None' or new_value == 'None':\n",
        "        raise NotImplementedError(\"Check parameters! You should provide:\\n- run_ID\\n- Project name\\n- Parameter\\n- New Value\")\n",
        "\n",
        "    run = wandb.Api().run(F\"{project_name}/{run_id}\")\n",
        "\n",
        "    # EXAMPLE: We edit the model_name parameter\n",
        "    run.config[parameter] = new_value\n",
        "\n",
        "    run.update()\n",
        "    wandb.finish()\n",
        "    return\n",
        "\n",
        "# configChanger(run_id='8gxx5lpy', project_name='DeepLearning-DSI', parameter='model_name', new_value='google/switch-base-8')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eD8Q89j-jej1"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZ1kKCW2b25R"
      },
      "source": [
        "### Generative Approach"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJGP_bDob25R"
      },
      "source": [
        "#### DSI Model with Foundation Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UZr44RFYb25R"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "  -- DSI Model --\n",
        "\n",
        "  Parameters:\n",
        "    - config: The configuration dictionary populated throughout the whole Code\n",
        "    - model: The model to use (populated when instantiating the Foundation Model)\n",
        "    - query_doc_df: The dataframe containing the queries\n",
        "        This is needed for retrieving DocIDs during the retrieval task\n",
        "    - training_dataloader_with_queries_filtered: The dataloader containing the filtered documents with queries\n",
        "        Due to Lightning's limitations, we need to pass the dataloader to the model in order to be able to iterate over it\n",
        "'''\n",
        "class DSI_Model(pl.LightningModule):\n",
        "    def __init__(self,\n",
        "                 config: dict,\n",
        "                 model: pl.LightningModule,\n",
        "                 query_doc_df: pd.DataFrame,\n",
        "                 training_dataloader_with_queries_filtered: DataLoader,\n",
        "                 ):\n",
        "        super().__init__()\n",
        "\n",
        "        # Variables\n",
        "        self.config = config\n",
        "        self.MAX_LENGTH = config[\"MAX_LENGTH\"]\n",
        "\n",
        "        self.token_tokenization_max_length = config['token_tokenization_max_length'] # Define the label max_length of token for the tokenization\n",
        "        self.train_indexing_retrieval_ratio = config[\"train_indexing_retrieval_ratio\"] # Docs with Query for Training\n",
        "        self.train_begin_retrieval = config[\"train_begin_retrieval\"] # At which index the retrieval task begins\n",
        "\n",
        "        self.enable_multitask_prompting = config[\"enable_multitask_prompting\"]\n",
        "        if(self.enable_multitask_prompting):\n",
        "          self.multitask_prompting_indexing = config[\"multitask_prompting_indexing\"]\n",
        "          self.multitask_prompting_retrieval = config[\"multitask_prompting_retrieval\"]\n",
        "\n",
        "        # Foundation Model\n",
        "        self.model = model.to(device)\n",
        "\n",
        "        # Dataloader\n",
        "        # We need to keep a copy of the training dataloader, due to the fact that we\n",
        "        # must reset the training dataloader every time we finish an epoch\n",
        "        # If we don't do this, we will have a StopIteration error\n",
        "        self.copy_training_dataloader = training_dataloader_with_queries_filtered\n",
        "        self.training_dataloader_with_queries_filtered = iter(training_dataloader_with_queries_filtered)\n",
        "\n",
        "        # Log losses\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "        self.validation_retrieval_losses = []\n",
        "\n",
        "        self.training_loss = []\n",
        "        self.validation_loss = []\n",
        "        self.test_accuracy = []\n",
        "\n",
        "        self.best_training_loss = float('inf')\n",
        "        self.best_validation_loss = float('inf')\n",
        "        self.best_test_accuracy = 0.0\n",
        "\n",
        "        # WandB\n",
        "        self.enable_wandb = config[\"wandb_configs\"][\"enable\"]\n",
        "        if(self.enable_wandb):\n",
        "          self.wandb_run = wandb.init(\n",
        "              project='DeepLearning-DSI',\n",
        "              group=config[\"wandb_configs\"][\"group_id\"],\n",
        "              config=config,)\n",
        "        else:\n",
        "          self.wandb_run = None\n",
        "\n",
        "        # Dataframes\n",
        "        self.query_doc_df = query_doc_df  # Query - Doc ID Dataframe\n",
        "        self.top_100_df = None            # Populated only if the test_type is \"top_k\"\n",
        "\n",
        "        # The current task\n",
        "        self.current_task = \"indexing\" # \"indexing\" / \"retrieval\"\n",
        "        self.test_type = \"accuracy\" # \"accuracy\" / \"top_k\"\n",
        "\n",
        "        # Create a Dict DocID2Index and Index2DocID\n",
        "        # Useful when we need to retrieve the DocID from the index and vice versa (Test phase)\n",
        "        self.docID2label = {}\n",
        "        self.label2DocID = {}\n",
        "        filtered_docs = train_docs_with_queries_dataset.filtered_docs_with_queries\n",
        "        for _, row in filtered_docs.iterrows():\n",
        "          self.docID2label[row['Doc_ID']] = str(row[config['document_id_representation_strategy']])\n",
        "          self.label2DocID[str(row[config['document_id_representation_strategy']])] = row['Doc_ID']\n",
        "\n",
        "        # Generation variables\n",
        "        self.num_return_sequences = 5\n",
        "        self.num_beams = 5\n",
        "\n",
        "    '''\n",
        "      Generate method:\n",
        "        - docs_input: The input for the model\n",
        "        - labels_att_mask: The attention mask for the labels\n",
        "    '''\n",
        "    def generate(self, docs_input, labels_att_mask):\n",
        "        docs_input_ids = docs_input['input_ids'].to(device)\n",
        "        docs_input_att_mask = docs_input[\"attention_mask\"].to(device)\n",
        "\n",
        "        # Here we might want to force the model to generate a specific token at the beginning of the sequence\n",
        "        force_words = [\"0\"]\n",
        "        force_words_ids = self.model.tokenizer(force_words, add_special_tokens=False).input_ids\n",
        "\n",
        "        # If we are using the BertUncased Model, we need to pass the decoder_start_token_id\n",
        "        # If num_beams > 0, then it means that we are using the Constrainted Beam Search\n",
        "        with torch.no_grad():\n",
        "          generated_ids = self.model.model.generate(\n",
        "              inputs=docs_input_ids, # [bs, 256]\n",
        "              attention_mask=docs_input_att_mask, # [bs, 256]\n",
        "              # decoder_attention_mask=labels_att_mask, # [bs, 4]\n",
        "              max_length=len(str(self.token_tokenization_max_length)),\n",
        "              decoder_start_token_id = self.model.tokenizer.cls_token_id if self.model.is_encoder_decoder else None,\n",
        "              #force_words_ids=force_words_ids,\n",
        "              num_beams=self.num_beams,\n",
        "              num_return_sequences=self.num_return_sequences,\n",
        "              do_sample=True,\n",
        "              top_k=50,\n",
        "              temperature=0.6,\n",
        "          )\n",
        "\n",
        "        return self.model.tokenizer.batch_decode(\n",
        "            generated_ids,\n",
        "            skip_special_tokens=True,\n",
        "            clean_up_tokenization_spaces=True\n",
        "        )\n",
        "    '''\n",
        "      The forward method of the DSI model that differentiate slightly the procedure to\n",
        "      follow according to the task that has to be performed (indexing/retrieval)\n",
        "    '''\n",
        "    def forward(self, batch, training=True):\n",
        "        torch.cuda.empty_cache() # Clear CUDA cache before forward pass\n",
        "\n",
        "        if self.current_task == \"indexing\": # Indexing Task\n",
        "          # T5 Prompt Setup as suggested in the paper\n",
        "          if self.enable_multitask_prompting:\n",
        "            docs_body = [f\"{self.multitask_prompting_indexing} {item}\" for item in batch[train_dataset.doc_rep_engine.strategy]]\n",
        "          else:\n",
        "            docs_body = batch[train_dataset.doc_rep_engine.strategy]\n",
        "\n",
        "          # Labels are the Doc_ID(s) processed with the strategy chosen\n",
        "          labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          docs_input = self.model.tokenizer(docs_body, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "          labels_inputs = self.model.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=len(str(self.token_tokenization_max_length)), padding=\"max_length\", truncation=True)\n",
        "\n",
        "          labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "          labels_att_mask = labels_inputs['attention_mask'].to(device)\n",
        "          input_ids = docs_input['input_ids'].to(device)\n",
        "          attention_mask = docs_input['attention_mask'].to(device)\n",
        "\n",
        "          labels_input_ids[labels_input_ids == self.model.tokenizer.pad_token_id] = -100\n",
        "\n",
        "          # For training -> Includes loss computation as well\n",
        "          if(training):\n",
        "            # The model is called with the input_ids and the labels\n",
        "            # The input_ids will flow through the Encoder-Decoder model\n",
        "            return self.model.model(\n",
        "                input_ids,\n",
        "                attention_mask=attention_mask,\n",
        "                decoder_input_ids=None,\n",
        "                decoder_attention_mask=labels_att_mask,\n",
        "                labels=labels_input_ids\n",
        "            )\n",
        "          else: # For validation and test\n",
        "            return self.generate(docs_input, labels_att_mask)\n",
        "        else:\n",
        "          # Retrieval Task. The process is the exact same as the indexing task,\n",
        "          # but we need to consider the queries associated to the documents, not the documents themselves.\n",
        "          queries = batch[\"Query\"]\n",
        "\n",
        "          queries = [str(item[0]) for item in queries]\n",
        "          if self.enable_multitask_prompting:\n",
        "            queries = [f\"{self.multitask_prompting_retrieval} {query}\" for query in queries]\n",
        "\n",
        "          labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          queries_input = self.model.tokenizer(queries, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "          labels_inputs = self.model.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=len(str(self.token_tokenization_max_length)), padding=\"max_length\", truncation=True)\n",
        "\n",
        "          labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "          labels_att_mask = labels_inputs['attention_mask'].to(device)\n",
        "          labels_input_ids[labels_input_ids == self.model.tokenizer.pad_token_id] = -100\n",
        "          input_ids = queries_input['input_ids'].to(device)\n",
        "          attention_mask = queries_input['attention_mask'].to(device)\n",
        "\n",
        "          if(training):\n",
        "            return self.model.model(\n",
        "                input_ids,\n",
        "                attention_mask=attention_mask,\n",
        "                decoder_input_ids=None,\n",
        "                decoder_attention_mask=labels_att_mask,\n",
        "                labels=labels_input_ids\n",
        "            )\n",
        "          else:\n",
        "            return self.generate(queries_input, labels_att_mask)\n",
        "\n",
        "    '''\n",
        "      A retrival batch is created starting from a batch considering the questions associated to doc\n",
        "      in the current batch.\n",
        "    '''\n",
        "    def prepare_sample_for_retrieval(self, batch):\n",
        "        batch['Queries_ID'] = []\n",
        "        batch['Query'] = []\n",
        "        doc_ids = batch['Doc_ID']\n",
        "\n",
        "        for doc_id in doc_ids:\n",
        "            query_ids = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query_ID'].values\n",
        "\n",
        "            if(len(query_ids) == 0):\n",
        "              raise Exception(f\"Document {doc_id} without queries!\")\n",
        "            else:\n",
        "              queries = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query'].values\n",
        "\n",
        "              batch['Queries_ID'].append(query_ids)\n",
        "              batch['Query'].append(queries)\n",
        "\n",
        "        return batch\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        # ========================================================================\n",
        "        # Retrieval step | Every self.train_indexing_retrieval_ratio batches\n",
        "        #\n",
        "        # Training step => Indexing all the documents and perform retrival\n",
        "        #                  (at the end of each epoch considering the given ratio)\n",
        "        #                  on just the document chosen as training docs.\n",
        "        # ========================================================================\n",
        "\n",
        "        # Indexing step\n",
        "        self.current_task = 'indexing'\n",
        "\n",
        "        output = self.forward(batch, training=True)\n",
        "        loss_idx = output.loss\n",
        "\n",
        "        self.training_indexing_losses.append(loss_idx.cpu().detach().item())\n",
        "\n",
        "        total_step_loss = loss_idx\n",
        "\n",
        "        # Switch to retrieval task if the batch index is greater than the train_begin_retrieval value\n",
        "        # This is done to alternate between indexing and retrieval tasks as suggested in the paper\n",
        "        if batch_idx >= self.train_begin_retrieval:\n",
        "          self.current_task = 'retrieval'\n",
        "\n",
        "          # If we are in the retrieval task, we need to prepare the sample for retrieval\n",
        "          # Hence, we need to consider the queries associated to the documents in the current batch\n",
        "          batch = next(self.training_dataloader_with_queries_filtered)\n",
        "          batch = self.prepare_sample_for_retrieval(batch)\n",
        "\n",
        "          output = self.forward(batch, training=True)\n",
        "          loss_retrieval = output.loss\n",
        "\n",
        "          self.training_retrieval_losses.append(loss_retrieval.cpu().detach().item())\n",
        "\n",
        "          total_step_loss = (total_step_loss + loss_retrieval) / 2\n",
        "\n",
        "        self.log(\"training_loss\", total_step_loss.cpu().detach().item())\n",
        "        self.training_loss.append(total_step_loss.cpu().detach().item())\n",
        "\n",
        "        return total_step_loss\n",
        "\n",
        "    # Validation => Retrieval step for all the documents chosen to be in validation split\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        self.current_task = 'retrieval'\n",
        "\n",
        "        # Retrieval step (same as retrieval step in training phase)\n",
        "        batch = self.prepare_sample_for_retrieval(batch)\n",
        "        bs = len(batch['Doc_ID'])\n",
        "\n",
        "        output = self.forward(batch, training=True)\n",
        "        loss = output.loss\n",
        "\n",
        "        self.log(\"validation_loss\", loss.cpu().detach().item(), batch_size=bs)\n",
        "        self.validation_loss.append(loss.cpu().detach().item())\n",
        "        self.validation_retrieval_losses.append(loss.cpu().detach().item())\n",
        "\n",
        "        return loss\n",
        "\n",
        "    # Test => Retrieval step for all the documents chosen to be in test split\n",
        "    def test_step(self, batch, batch_idx):\n",
        "      self.current_task = 'retrieval'\n",
        "\n",
        "      # Retrieval step (same as retrieval step in training/validation phase)\n",
        "      batch = self.prepare_sample_for_retrieval(batch)\n",
        "      bs = len(batch['Doc_ID'])\n",
        "\n",
        "      output = self.forward(batch, training=False)\n",
        "\n",
        "      if(self.test_type == \"accuracy\"):\n",
        "        accuracy = self.run_accuracy(output, batch)\n",
        "      else:\n",
        "        accuracy = self.run_top_k(output, batch)\n",
        "\n",
        "      self.log(\"test_accuracy\", accuracy, batch_size=bs)\n",
        "      self.test_accuracy.append(accuracy)\n",
        "\n",
        "      return accuracy\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model are in the top-k relevant document of the queries from the batch\n",
        "    '''\n",
        "    def run_top_k(self, output, batch):\n",
        "      # Before trying to find if the outputs are relevant w.r.t query, we first check if the query is present in the Top100 DF.\n",
        "      list_df = []\n",
        "      batch_queries = batch['Queries_ID']\n",
        "      for query in batch_queries:\n",
        "          q = self.top_100_df[self.top_100_df['Query_ID'] == query[0]]\n",
        "          if len(q) > 0:\n",
        "              list_df.append(q)\n",
        "          else:\n",
        "              pass\n",
        "              #print(f\"There are no queries associated with batch query {query[0]}!\")\n",
        "      batch_top_100 = pd.concat(list_df)\n",
        "\n",
        "      # Check if num_return_sequences > 1, if so, we need to iterate over the output differently\n",
        "      if(self.num_return_sequences > 1):\n",
        "          batch_predictions = []\n",
        "          final_returns = []\n",
        "\n",
        "          # Iterate over each prediction sequence in the output\n",
        "          for i in range(0, len(output), self.num_return_sequences):\n",
        "              # Get the predictions for the current batch\n",
        "              batch_predictions.append(output[i:i+self.num_return_sequences])\n",
        "\n",
        "          # Extract DocIDs from the output\n",
        "          for idx, b_pred in enumerate(batch_predictions):\n",
        "              model_returns = []\n",
        "              for p in b_pred:\n",
        "                  try:\n",
        "                      model_returns.append(self.label2DocID[p])\n",
        "                      self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "                  except:\n",
        "                      model_returns.append('')\n",
        "                      #print(f\"There is no DocID for model prediction: {p}!\")\n",
        "              final_returns.append(model_returns)\n",
        "\n",
        "          # Now we need to:\n",
        "          # - Check if the batch predictions are in the top 100 for their respective queries\n",
        "          correct = 0\n",
        "          for idx, predictions in enumerate(final_returns):\n",
        "              for prediction in predictions:\n",
        "                  if prediction in batch_top_100[batch_top_100['Query_ID'] == batch_queries[idx][0]]['Doc_ID'].values:\n",
        "                      correct += 1\n",
        "          return float(correct / len(batch_predictions))\n",
        "      else: # num_return_sequences == 1 / 0\n",
        "          # Extract DocIDs from the output\n",
        "          model_returns = []\n",
        "          for idx, b_pred in enumerate(output):\n",
        "              try:\n",
        "                  model_returns.append(self.label2DocID[b_pred])\n",
        "                  self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "              except:\n",
        "                  model_returns.append('')\n",
        "                  #print(f\"There is no DocID for model prediction: {b_pred}!\")\n",
        "\n",
        "          # Now we need to:\n",
        "          # - Take the first query from the batch\n",
        "          # - Check if the first model_returns is in the top 100 for the first query\n",
        "          # - Loop this process for all queries/outputs\n",
        "          correct = 0\n",
        "          for idx, query in enumerate(batch_queries):\n",
        "              if(model_returns[idx] in batch_top_100[batch_top_100['Query_ID'] == query[0]]['Doc_ID'].values):\n",
        "                  correct += 1\n",
        "\n",
        "          return float(correct / len(output))\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model is equal to the label present in the batch.\n",
        "    '''\n",
        "    def run_accuracy(self, output, batch):\n",
        "      batch_predictions = []\n",
        "\n",
        "      if(len(output) > self.config['batch_size']): # It means that we are in the generation phase with num_return_sequences > 1\n",
        "        # Iterate over each prediction sequence in the output\n",
        "        for i in range(0, len(output), self.num_return_sequences):\n",
        "            # Get the predictions for the current batch\n",
        "            batch_predictions.append(output[i:i+self.num_return_sequences])\n",
        "\n",
        "        correct = 0\n",
        "        for idx, pred in enumerate(batch_predictions): # Enters a single batch of predictions.\n",
        "          # Compute the accuracy\n",
        "          for b_pred in pred: # Enters a single prediction\n",
        "            try:\n",
        "              model_return = self.label2DocID[b_pred]\n",
        "              self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "            except:\n",
        "              model_return = ''\n",
        "              #print(f\"The model did not output a correct DocID for prediction {b_pred}!\")\n",
        "\n",
        "            if model_return == str(batch['Doc_ID'][idx]):\n",
        "              correct += 1\n",
        "        return correct / len(batch_predictions)\n",
        "\n",
        "      else:\n",
        "        for idx, b_pred in enumerate(output):\n",
        "          try:\n",
        "            model_return = self.label2DocID[b_pred]\n",
        "            self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "          except:\n",
        "            model_return = ''\n",
        "            #print(f\"The model did not output a correct DocID for prediction {b_pred}!\")\n",
        "\n",
        "          if model_return == str(batch['Doc_ID'][idx]):\n",
        "            correct += 1\n",
        "        return correct / len(batch['Doc_ID'])\n",
        "\n",
        "    def on_train_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.training_loss) / len(self.training_loss) if len(self.training_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_training_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'TRAINING'}\")\n",
        "\n",
        "        best_loss = getattr(self, f\"best_training_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_training_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.training_indexing_losses)!=0):\n",
        "          loss_indexing = sum(loss for loss in self.training_indexing_losses) / len(self.training_indexing_losses) if len(self.training_indexing_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_indexing_loss\", loss_indexing, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Indexing Loss => {loss_indexing:.4f}\")\n",
        "\n",
        "        if(len(self.training_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.training_retrieval_losses) / len(self.training_retrieval_losses) if len(self.training_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "            self.wandb_run.log({f\"avg_training_loss\": epoch_loss})\n",
        "            if(len(self.training_indexing_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_indexing_loss\": loss_indexing})\n",
        "            if(len(self.training_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_retrieval_loss\": loss_retrieval})\n",
        "            self.wandb_run.log({f\"epoch\": self.current_epoch})\n",
        "\n",
        "        # At the end of each training epoch we select a differet subset of training document\n",
        "        # on which perform the training retrieval. This ensure to have stocasticity selection\n",
        "        # while performing the retrieval on training docs\n",
        "        self.training_dataloader_with_queries_filtered = iter(self.copy_training_dataloader)\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.training_loss = []\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "\n",
        "    def on_validation_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.validation_loss) / len(self.validation_loss) if len(self.validation_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_validation_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'VALIDATION'}\")\n",
        "\n",
        "        best_loss = getattr(self, f\"best_validation_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_validation_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.validation_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.validation_retrieval_losses) / len(self.validation_retrieval_losses) if len(self.validation_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_validation_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "            self.wandb_run.log({f\"avg_validation_loss\": epoch_loss})\n",
        "            if(len(self.validation_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_validation_retrieval_loss\": loss_retrieval})\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.validation_loss = []\n",
        "        self.validation_retrieval_losses = []\n",
        "\n",
        "    def on_test_epoch_end(self) -> None:\n",
        "      epoch_accuracy = sum(acc for acc in self.test_accuracy) / len(self.test_accuracy) if len(self.test_accuracy) > 0 else float('nan')\n",
        "      self.log(f\"avg_test_accuracy\", epoch_accuracy, batch_size=self.config['batch_size'])\n",
        "\n",
        "      print(f\"| Epoch {self.current_epoch} | {'TEST'}\")\n",
        "\n",
        "      best_accuracy = getattr(self, f\"best_test_accuracy\")\n",
        "\n",
        "      # Update best loss if current loss is better\n",
        "      if epoch_accuracy < best_accuracy:\n",
        "          setattr(self, f\"best_test_accuracy\", epoch_accuracy)\n",
        "\n",
        "      print(f\"\\t- Total Accuracy => {epoch_accuracy:.4f}\")\n",
        "\n",
        "      if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "          self.wandb_run.log({f\"avg_test_accuracy\": epoch_accuracy})\n",
        "\n",
        "      # Empty the list of the losses for the next epoch\n",
        "      self.test_accuracy = []\n",
        "\n",
        "    def save_model_pytorch_api(self, name='best', epoch=''):\n",
        "        torch.save(self.state_dict(), f'{name}_ep-{epoch}.pt')\n",
        "\n",
        "    # In this case, we configure the model in 3 ways:\n",
        "    #   - Adam 8-Bit\n",
        "    #   - AdamW 8-Bit\n",
        "    #   - Normal Adam\n",
        "    #   - Normal AdamW\n",
        "    # There is an option to also set the weight decay for the parameters\n",
        "    def configure_optimizers(self):\n",
        "\n",
        "        # # ---- Grouped Parameters instead of self.parameters() directly to Adam ----\n",
        "        # decay_parameters = get_parameter_names(self.model, [nn.LayerNorm])\n",
        "        # decay_parameters = [name for name in decay_parameters if \"bias\" not in name]\n",
        "\n",
        "        # optimizer_grouped_parameters = [\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n not in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        # ]\n",
        "        # # ---------------------------------------------------------------------------\n",
        "\n",
        "        # Configs\n",
        "        epsilon = 1e-8  # Default\n",
        "        lr = 3e-4       # Default: 1e-3 | Can be changed also to 3e-4\n",
        "\n",
        "        # # Adam 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.Adam8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # AdamW 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.AdamW8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # Normal Adam\n",
        "        # adam_bnb_optim = bnb.optim.Adam(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # Standard PyTorch AdamW\n",
        "        adam_bnb_optim = torch.optim.AdamW(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # # Standard PyTorch Adam\n",
        "        # adam_bnb_optim = torch.optim.Adam(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # TODO: Add a Linear Warmup LR (maybe https://lightning-flash.readthedocs.io/en/stable/api/generated/flash.core.optimizers.LinearWarmupCosineAnnealingLR.html)\n",
        "        scheduler = pl_bolts.optimizers.lr_scheduler.LinearWarmupCosineAnnealingLR(adam_bnb_optim, warmup_epochs=2, warmup_start_lr=0.0, eta_min=0.0, max_epochs=self.trainer.max_epochs)\n",
        "        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(adam_bnb_optim, T_max=self.trainer.max_epochs)\n",
        "\n",
        "        return {\"optimizer\": adam_bnb_optim, \"lr_scheduler\": scheduler}\n",
        "\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return self.model.get_n_trainable_parameters()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XsLjz5nKb25S"
      },
      "outputs": [],
      "source": [
        "# Define the model\n",
        "#\n",
        "# 'model_name' = 'bert-base-uncased'  : The system will load the Bert-Uncased model\n",
        "# 'model_name' = 'google/flan-t5-base': The system will load the Flan-T5 model\n",
        "#\n",
        "# All the configs will be dynamically set according to the model chosen\n",
        "dsi_config['model_name'] = 'bert-base-uncased' # 'google/flan-t5-base' | 'bert-base-uncased'\n",
        "dsi_config['fine_tuning'] = False\n",
        "\n",
        "# Load the model\n",
        "if(dsi_config['model_name'] != 'bert-base-uncased'):\n",
        "    model = FoundationModel(model_id=dsi_config['model_name'], encoder_id=None, decoder_id=None, MAX_LENGTH=dsi_config[\"MAX_LENGTH\"], fine_tuning=dsi_config['fine_tuning']) # Flan-T5 Model\n",
        "    dsi_config[\"hidden_size\"] = model.model.config.hidden_size\n",
        "else:\n",
        "    dsi_config['encoder_model'] = 'bert-base-uncased'\n",
        "    dsi_config['decoder_model'] = 'bert-base-uncased'\n",
        "    model = FoundationModel(model_id=None, encoder_id=dsi_config['encoder_model'], decoder_id=dsi_config['decoder_model'], MAX_LENGTH=dsi_config[\"MAX_LENGTH\"], fine_tuning=dsi_config['fine_tuning']) # Bert-Uncased Model\n",
        "    dsi_config[\"hidden_size\"] = model.model.encoder.config.hidden_size\n",
        "    dsi_config.pop(\"model_name\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3NPidr31b25S"
      },
      "outputs": [],
      "source": [
        "# Create actual DSI Model\n",
        "dsi_model = DSI_Model(dsi_config,\n",
        "                        model=model,\n",
        "                        query_doc_df=query_doc_df,\n",
        "                        training_dataloader_with_queries_filtered=train_dataloader_filtered,)\n",
        "\n",
        "total_params = dsi_model.get_n_trainable_parameters()\n",
        "total_params = \"{:,}\".format(total_params).replace(\",\", \".\")\n",
        "print(\"--------------------------------------------------------------\")\n",
        "print(f\"Total number of trainable parameters for DSI Model: {total_params}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vtav475hb25T"
      },
      "source": [
        "#### DSI Model with Encoder-Decoder\n",
        "SwitchTransformer + FlanT5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sfj1rCcFb25V"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "  -- DSI Model for the custom Encoder Decoder model--\n",
        "\n",
        "  Parameters:\n",
        "    - config: The configuration dictionary populated throughout the whole Code\n",
        "    - encoder: The encoder to use (SwitchTransformer)\n",
        "    - decoder: The decoder to use (Flan-T5)\n",
        "    - query_doc_df: The dataframe containing the queries\n",
        "        This is needed for retrieving DocIDs during the retrieval task\n",
        "    - training_dataloader_with_queries_filtered: The dataloader containing the filtered documents with queries\n",
        "        Due to Lightning's limitations, we need to pass the dataloader to the model in order to be able to iterate over it\n",
        "'''\n",
        "class DSI_EncoderDecoder(pl.LightningModule):\n",
        "    def __init__(self,\n",
        "                 config: dict,\n",
        "                 encoder: pl.LightningModule,\n",
        "                 decoder: pl.LightningModule,\n",
        "                 query_doc_df: pd.DataFrame,\n",
        "                 training_dataloader_with_queries_filtered: DataLoader,\n",
        "                 ):\n",
        "        super().__init__()\n",
        "\n",
        "        # Variables\n",
        "        self.config = config\n",
        "        self.MAX_LENGTH = config[\"MAX_LENGTH\"]\n",
        "\n",
        "        self.token_tokenization_max_length = config['token_tokenization_max_length'] # Define the label max_length of token for the tokenization\n",
        "        self.train_indexing_retrieval_ratio = config[\"train_indexing_retrieval_ratio\"] # Docs with Query for Training\n",
        "        self.train_begin_retrieval = config[\"train_begin_retrieval\"] # At which index the retrieval task begins\n",
        "\n",
        "        self.enable_multitask_prompting = config[\"enable_multitask_prompting\"]\n",
        "        if(self.enable_multitask_prompting):\n",
        "          self.multitask_prompting_indexing = config[\"multitask_prompting_indexing\"]\n",
        "          self.multitask_prompting_retrieval = config[\"multitask_prompting_retrieval\"]\n",
        "\n",
        "        # Encoder and Decoder\n",
        "        self.encoder = encoder.to(device)\n",
        "        self.decoder = decoder.to(device)\n",
        "\n",
        "        # Dataloader\n",
        "        # We need to keep a copy of the training dataloader, due to the fact that we\n",
        "        # must reset the training dataloader every time we finish an epoch\n",
        "        # If we don't do this, we will have a StopIteration error\n",
        "        self.copy_training_dataloader = training_dataloader_with_queries_filtered\n",
        "        self.training_dataloader_with_queries_filtered = iter(training_dataloader_with_queries_filtered)\n",
        "\n",
        "        # Log losses\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "        self.validation_retrieval_losses = []\n",
        "        self.test_retrieval_losses = []\n",
        "\n",
        "        self.training_loss = []\n",
        "        self.validation_loss = []\n",
        "        self.test_accuracy = []\n",
        "\n",
        "        self.best_training_loss = float('inf')\n",
        "        self.best_validation_loss = float('inf')\n",
        "        self.best_test_accuracy = 0.0\n",
        "\n",
        "        # WandB\n",
        "        self.enable_wandb = config[\"wandb_configs\"][\"enable\"]\n",
        "        if(self.enable_wandb):\n",
        "          self.wandb_run = wandb.init(\n",
        "              project='DeepLearning-DSI',\n",
        "              group=config[\"wandb_configs\"][\"group_id\"],\n",
        "              config=config,)\n",
        "        else:\n",
        "          self.wandb_run = None\n",
        "\n",
        "        # Dataframes\n",
        "        self.query_doc_df = query_doc_df  # Query - Doc ID Dataframe\n",
        "        self.top_100_df = None            # Populated only if the test_type is \"top_k\"\n",
        "\n",
        "        # The current task\n",
        "        self.current_task = \"indexing\"  # \"indexing\" / \"retrieval\"\n",
        "        self.test_type = \"accuracy\"     # \"accuracy\" / \"top_k\"\n",
        "\n",
        "        # Create a Dict DocID2Index and Index2DocID\n",
        "        # Useful when we need to retrieve the DocID from the index and vice versa (Test phase)\n",
        "        self.docID2label = {}\n",
        "        self.label2DocID = {}\n",
        "        filtered_docs = train_docs_with_queries_dataset.filtered_docs_with_queries\n",
        "        for index, row in filtered_docs.iterrows():\n",
        "          self.docID2label[row['Doc_ID']] = str(row[config['document_id_representation_strategy']])\n",
        "          self.label2DocID[str(row[config['document_id_representation_strategy']])] = row['Doc_ID']\n",
        "\n",
        "        # Generation variables\n",
        "        self.num_return_sequences = 5\n",
        "        self.num_beams = 5\n",
        "\n",
        "    '''\n",
        "      Generate method:\n",
        "        - outputs_encoder: The outputs coming from the Encoder model\n",
        "        - enc_input_ids: The input_ids coming from the Encoder Tokenizer\n",
        "    '''\n",
        "    def generate(self, outputs_encoder, enc_input_ids):\n",
        "        # Switch to CUDA first\n",
        "        enc_input_ids = enc_input_ids.to(device)\n",
        "        self.decoder = self.decoder.to(device)\n",
        "\n",
        "        # Here we might want to force the model to generate a specific token at the beginning of the sequence\n",
        "        force_words = [\"0\"]\n",
        "        force_words_ids = self.encoder.tokenizer(force_words, add_special_tokens=False).input_ids\n",
        "\n",
        "        # In this case, we have already computed the Encoder part from SwitchTransformer (Encoder Model)\n",
        "        # Hence, we just want to generate the output from the Decoder part (Flan-T5)\n",
        "        #\n",
        "        # If num_beams > 0, then it means that we are using the Constrainted Beam Search\n",
        "        with torch.no_grad():\n",
        "          output = self.decoder.model.generate(\n",
        "              input_ids=None,                                               # No input_ids are needed for the Encoder part the Flan-T5 model\n",
        "              decoder_input_ids=enc_input_ids,                              # Use generated SwitchTransformer input_ids for the Decoder (Flan-T5) input ids\n",
        "              encoder_hidden_states=outputs_encoder.encoder_hidden_states,  # Use generated SwitchTransformer hidden states as the Encoder (Flan-T5) hidden states (skipping Encoding of Flan-T5)\n",
        "              bos_token_id=self.decoder.tokenizer.pad_token_id,             # The beginning of the sequence token. Needed for generation\n",
        "              max_length=len(str(self.token_tokenization_max_length)),      # Because the number of possbile doc_ids is a certain values, this parameters ensure that the generation\n",
        "                                                                            # create at maximum the number of token useful to format a doc_id | Example: Max Length = 5 -> '00012'\n",
        "              #force_words_ids=force_words_ids,                              # Use the force_words_ids to force the model to generate a specific token at the beginning of the sequence\n",
        "              do_sample=True,\n",
        "              use_cache=True,\n",
        "              top_k=50,\n",
        "              temperature=0.6,\n",
        "\n",
        "          )\n",
        "\n",
        "        return self.decoder.tokenizer.batch_decode(\n",
        "          output,\n",
        "          skip_special_tokens=True,\n",
        "          clean_up_tokenization_spaces=True\n",
        "        )\n",
        "\n",
        "    '''\n",
        "      The forward method of the DSI model that differentiate slightly the procedure to\n",
        "      follow according to the task that has to be performed (indexing/retrieval)\n",
        "    '''\n",
        "    def forward(self, batch, training=True):\n",
        "        torch.cuda.empty_cache() # Clear CUDA cache before forward pass\n",
        "\n",
        "        if self.current_task == \"indexing\": # Indexing Task\n",
        "          # T5 Prompt Setup as suggested in the paper\n",
        "          if self.enable_multitask_prompting:\n",
        "            docs_body = [f\"{self.multitask_prompting_indexing} {item}\" for item in batch[train_dataset.doc_rep_engine.strategy]]\n",
        "          else:\n",
        "            docs_body = batch[train_dataset.doc_rep_engine.strategy]\n",
        "\n",
        "          # Labels are the Doc_ID(s) processed with the strategy chosen\n",
        "          labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          # Calling the encoder\n",
        "          # Returns:\n",
        "          #   - embeddings: (last_hidden_states, encoder_hidden_states)\n",
        "          #   - docs_input_tokenized: The tokenized input of the documents\n",
        "          #   - outputs_encoder: The outputs of the Encoder model\n",
        "          embeddings, docs_input_tokenized, outputs_encoder = self.encoder(docs_body)\n",
        "\n",
        "          # Switch to CUDA\n",
        "          embeddings[0].to(device)\n",
        "          input_ids = docs_input_tokenized[\"input_ids\"].to(device)\n",
        "          attention_mask = docs_input_tokenized[\"attention_mask\"].to(device)\n",
        "          outputs_encoder['encoder_hidden_states'] = [item.to(device) for item in outputs_encoder['encoder_hidden_states']]\n",
        "\n",
        "          # For training -> Includes loss computation as well\n",
        "          if(training):\n",
        "            # Treat the labels as tokenizable strings\n",
        "            labels_inputs = self.encoder.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "            labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "            #labels_attention_mask = labels_inputs[\"attention_mask\"]\n",
        "\n",
        "            # Mask the labels corresponding to the padding tokens of the Encoder Tokenizer with -100\n",
        "            # Such that they are not taken into account in the loss computation\n",
        "            labels_input_ids[labels_input_ids == self.encoder.tokenizer.pad_token_id] = -100\n",
        "\n",
        "            return self.decoder(\n",
        "              embeddings,                                   # Embeddings coming from Encoder Model's last 4 Hidden Layers\n",
        "              decoder_input_ids=input_ids,                  # Input_ids coming from Encoder Tokenizer\n",
        "              decoder_attention_mask=attention_mask,        # Attention mask coming from Encoder Tokenizer\n",
        "              labels=labels_input_ids,                      # Input_ids of labels tokenized with Encoder Tokenizer,\n",
        "            )\n",
        "          else: # For validation and test\n",
        "            return self.generate(outputs_encoder, input_ids)\n",
        "        else:\n",
        "          # Retrieval Task. The process is the exact same as the indexing task,\n",
        "          # but we need to consider the queries associated to the documents, not the documents themselves.\n",
        "          queries = batch[\"Query\"]\n",
        "\n",
        "          queries = [str(item[0]) for item in queries]\n",
        "          if self.enable_multitask_prompting:\n",
        "            queries = [f\"{self.multitask_prompting_retrieval} {query}\" for query in queries]\n",
        "\n",
        "          labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          # Calling the encoder\n",
        "          # embeddings = (last_hidden_states, encoder_hidden_states)\n",
        "          embeddings, docs_input_tokenized, outputs_encoder = self.encoder(queries)\n",
        "\n",
        "          # Switch to CUDA\n",
        "          embeddings[0].to(device)\n",
        "          input_ids = docs_input_tokenized[\"input_ids\"].to(device)\n",
        "          attention_mask = docs_input_tokenized[\"attention_mask\"].to(device)\n",
        "          outputs_encoder['encoder_hidden_states'] = [item.to(device) for item in outputs_encoder['encoder_hidden_states']]\n",
        "\n",
        "          if(training):\n",
        "            # Treat the labels as tokenizable strings\n",
        "            labels_inputs = self.encoder.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "            labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "            #labels_attention_mask = labels_inputs[\"attention_mask\"]\n",
        "\n",
        "            # Mask the labels corresponding to the padding tokens of the Encoder Tokenizer with -100\n",
        "            # Such that they are not taken into account in the loss computation\n",
        "            labels_input_ids[labels_input_ids == self.encoder.tokenizer.pad_token_id] = -100\n",
        "\n",
        "            return self.decoder(\n",
        "                embeddings,                             # Embeddings coming from Encoder Model's last 4 Hidden Layers\n",
        "                decoder_input_ids=input_ids,            # Input_ids coming from Encoder Tokenizer\n",
        "                decoder_attention_mask=attention_mask,  # Attention mask coming from Encoder Tokenizer\n",
        "                labels=labels_input_ids,                # Input_ids of labels tokenized with Encoder Tokenizer\n",
        "            )\n",
        "          else:\n",
        "            return self.generate(outputs_encoder, input_ids)\n",
        "\n",
        "    '''\n",
        "      A retrival batch is created starting from a batch considering the questions associated to doc\n",
        "      in the current batch.\n",
        "    '''\n",
        "    def prepare_sample_for_retrieval(self, batch):\n",
        "        batch['Queries_ID'] = []\n",
        "        batch['Query'] = []\n",
        "        doc_ids = batch['Doc_ID']\n",
        "\n",
        "        for doc_id in doc_ids:\n",
        "            query_ids = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query_ID'].values\n",
        "\n",
        "            if(len(query_ids) == 0):\n",
        "              raise Exception(f\"Document {doc_id} without queries!\")\n",
        "            else:\n",
        "              queries = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query'].values\n",
        "\n",
        "              batch['Queries_ID'].append(query_ids)\n",
        "              batch['Query'].append(queries)\n",
        "\n",
        "        return batch\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        # ========================================================================\n",
        "        # Retrieval step | Every self.train_indexing_retrieval_ratio batches\n",
        "        #\n",
        "        # Training step => Indexing all the documents and perform retrival\n",
        "        #                  (at the end of each epoch considering the given ratio)\n",
        "        #                  on just the document chosen as training docs.\n",
        "        # ========================================================================\n",
        "\n",
        "        # Indexing step\n",
        "        self.current_task = 'indexing'\n",
        "\n",
        "        output = self.forward(batch, training=True)\n",
        "        loss_idx = output.loss\n",
        "\n",
        "        self.training_indexing_losses.append(loss_idx.cpu().detach().item())\n",
        "\n",
        "        total_step_loss = loss_idx\n",
        "\n",
        "        # Switch to retrieval task if the batch index is greater than the train_begin_retrieval value\n",
        "        # This is done to alternate between indexing and retrieval tasks as suggested in the paper\n",
        "        if batch_idx >= self.train_begin_retrieval:\n",
        "          self.current_task = 'retrieval'\n",
        "\n",
        "          # If we are in the retrieval task, we need to prepare the sample for retrieval\n",
        "          # Hence, we need to consider the queries associated to the documents in the current batch\n",
        "          batch = next(self.training_dataloader_with_queries_filtered)\n",
        "          batch = self.prepare_sample_for_retrieval(batch)\n",
        "\n",
        "          output = self.forward(batch, training=True)\n",
        "          loss_retrieval = output.loss\n",
        "\n",
        "          self.training_retrieval_losses.append(loss_retrieval.cpu().detach().item())\n",
        "\n",
        "          total_step_loss = (total_step_loss + loss_retrieval) / 2\n",
        "\n",
        "        self.log(\"training_loss\", total_step_loss.cpu().detach().item())\n",
        "        self.training_loss.append(total_step_loss.cpu().detach().item())\n",
        "\n",
        "        return total_step_loss\n",
        "\n",
        "    # Validation => Retrieval step for all the documents chosen to be in validation split\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "      self.current_task = 'retrieval'\n",
        "\n",
        "      # Retrieval step (same as retrieval step in training phase)\n",
        "      batch = self.prepare_sample_for_retrieval(batch)\n",
        "      bs = len(batch['Doc_ID'])\n",
        "\n",
        "      output = self.forward(batch, training=True)\n",
        "      loss = output.loss\n",
        "\n",
        "      self.log(\"validation_loss\", loss.cpu().detach().item(), batch_size=bs)\n",
        "      self.validation_loss.append(loss.cpu().detach().item())\n",
        "      self.validation_retrieval_losses.append(loss.cpu().detach().item())\n",
        "\n",
        "      return loss\n",
        "\n",
        "    # Test => Retrieval step for all the documents chosen to be in test split\n",
        "    def test_step(self, batch, batch_idx):\n",
        "      self.current_task = 'retrieval'\n",
        "\n",
        "      # Retrieval step (same as retrieval step in training/validation phase)\n",
        "      batch = self.prepare_sample_for_retrieval(batch)\n",
        "      bs = len(batch['Doc_ID'])\n",
        "\n",
        "      output = self.forward(batch, training=False)\n",
        "\n",
        "      if(self.test_type == \"accuracy\"):\n",
        "        accuracy = self.run_accuracy(output, batch)\n",
        "      else:\n",
        "        accuracy = self.run_top_k(output, batch)\n",
        "\n",
        "      self.log(\"accuracy\", accuracy, batch_size=bs)\n",
        "      # self.log(\"test_loss\", loss.cpu().detach().item(), batch_size=bs)\n",
        "      # self.test_loss.append(loss.cpu().detach().item())\n",
        "      # self.test_retrieval_losses.append(loss.cpu().detach().item())\n",
        "\n",
        "      return accuracy\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model are in the top-k relevant document of the queries from the batch\n",
        "    '''\n",
        "    def run_top_k(self, output, batch):\n",
        "      # Before trying to find if the outputs are relevant w.r.t query, we first check if the query is present in the Top100 DF.\n",
        "      list_df = []\n",
        "      batch_queries = batch['Queries_ID']\n",
        "      for query in batch_queries:\n",
        "          q = self.top_100_df[self.top_100_df['Query_ID'] == query[0]]\n",
        "          if len(q) > 0:\n",
        "              list_df.append(q)\n",
        "          else:\n",
        "              pass\n",
        "              #print(f\"There are no queries associated with batch query {query[0]}!\")\n",
        "      batch_top_100 = pd.concat(list_df)\n",
        "\n",
        "      # Check if num_return_sequences > 1, if so, we need to iterate over the output differently\n",
        "      if(self.num_return_sequences > 1):\n",
        "          batch_predictions = []\n",
        "          final_returns = []\n",
        "\n",
        "          # Iterate over each prediction sequence in the output\n",
        "          for i in range(0, len(output), self.num_return_sequences):\n",
        "              # Get the predictions for the current batch\n",
        "              batch_predictions.append(output[i:i+self.num_return_sequences])\n",
        "\n",
        "          # Extract DocIDs from the output\n",
        "          for idx, b_pred in enumerate(batch_predictions):\n",
        "              model_returns = []\n",
        "              for p in b_pred:\n",
        "                  try:\n",
        "                      model_returns.append(self.label2DocID[p])\n",
        "                      self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "                  except:\n",
        "                      model_returns.append('')\n",
        "                      #print(f\"There is no DocID for model prediction: {p}!\")\n",
        "              final_returns.append(model_returns)\n",
        "\n",
        "          # Now we need to:\n",
        "          # - Check if the batch predictions are in the top 100 for their respective queries\n",
        "          correct = 0\n",
        "          for idx, predictions in enumerate(final_returns):\n",
        "              for prediction in predictions:\n",
        "                  if prediction in batch_top_100[batch_top_100['Query_ID'] == batch_queries[idx][0]]['Doc_ID'].values:\n",
        "                      correct += 1\n",
        "          return float(correct / len(batch_predictions))\n",
        "      else: # num_return_sequences == 1 / 0\n",
        "          # Extract DocIDs from the output\n",
        "          model_returns = []\n",
        "          for idx, b_pred in enumerate(output):\n",
        "              try:\n",
        "                  model_returns.append(self.label2DocID[b_pred])\n",
        "                  self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "              except:\n",
        "                  model_returns.append('')\n",
        "                  #print(f\"There is no DocID for model prediction: {b_pred}!\")\n",
        "\n",
        "          # Now we need to:\n",
        "          # - Take the first query from the batch\n",
        "          # - Check if the first model_returns is in the top 100 for the first query\n",
        "          # - Loop this process for all queries/outputs\n",
        "          correct = 0\n",
        "          for idx, query in enumerate(batch_queries):\n",
        "              if(model_returns[idx] in batch_top_100[batch_top_100['Query_ID'] == query[0]]['Doc_ID'].values):\n",
        "                  correct += 1\n",
        "\n",
        "          return float(correct / len(output))\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model is equal to the label present in the batch.\n",
        "    '''\n",
        "    def run_accuracy(self, output, batch):\n",
        "      batch_predictions = []\n",
        "\n",
        "      if(len(output) > self.config['batch_size']): # It means that we are in the generation phase with num_return_sequences > 1\n",
        "        # Iterate over each prediction sequence in the output\n",
        "        for i in range(0, len(output), self.num_return_sequences):\n",
        "            # Get the predictions for the current batch\n",
        "            batch_predictions.append(output[i:i+self.num_return_sequences])\n",
        "\n",
        "        correct = 0\n",
        "        for idx, pred in enumerate(batch_predictions): # Enters a single batch of predictions.\n",
        "          # Compute the accuracy\n",
        "          for b_pred in pred: # Enters a single prediction\n",
        "            try:\n",
        "              model_return = self.label2DocID[b_pred]\n",
        "              self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "            except:\n",
        "              model_return = ''\n",
        "              #print(f\"The model did not output a correct DocID for prediction {b_pred}!\")\n",
        "\n",
        "            if model_return == str(batch['Doc_ID'][idx]):\n",
        "              correct += 1\n",
        "        return correct / len(batch_predictions)\n",
        "\n",
        "      else:\n",
        "        correct = 0\n",
        "        for idx, b_pred in enumerate(output):\n",
        "          try:\n",
        "            model_return = self.label2DocID[b_pred]\n",
        "            self.model_found_label += 1 # Increment the number of labels found | Instantiated before starting the test phase\n",
        "          except:\n",
        "            model_return = ''\n",
        "            #print(f\"The model did not output a correct DocID for prediction {b_pred}!\")\n",
        "\n",
        "          if model_return == str(batch['Doc_ID'][idx]):\n",
        "            correct += 1\n",
        "        return correct / len(batch)\n",
        "\n",
        "    def on_train_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.training_loss) / len(self.training_loss) if len(self.training_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_training_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'TRAINING'}\")\n",
        "\n",
        "        best_loss = getattr(self, f\"best_training_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_training_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.training_indexing_losses)!=0):\n",
        "          loss_indexing = sum(loss for loss in self.training_indexing_losses) / len(self.training_indexing_losses) if len(self.training_indexing_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_indexing_loss\", loss_indexing, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Indexing Loss => {loss_indexing:.4f}\")\n",
        "\n",
        "        if(len(self.training_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.training_retrieval_losses) / len(self.training_retrieval_losses) if len(self.training_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "            self.wandb_run.log({f\"avg_training_loss\": epoch_loss})\n",
        "            if(len(self.training_indexing_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_indexing_loss\": loss_indexing})\n",
        "            if(len(self.training_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_retrieval_loss\": loss_retrieval})\n",
        "            self.wandb_run.log({f\"epoch\": self.current_epoch})\n",
        "\n",
        "        # At the end of each training epoch we select a differet subset of training document\n",
        "        # on which perform the training retrieval. This ensure to have stocasticity selection\n",
        "        # while performing the retrieval on training docs\n",
        "        self.training_dataloader_with_queries_filtered = iter(self.copy_training_dataloader)\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.training_loss = []\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "\n",
        "    def on_validation_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.validation_loss) / len(self.validation_loss) if len(self.validation_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_validation_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'VALIDATION'}\")\n",
        "\n",
        "        best_loss = getattr(self, f\"best_validation_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_validation_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.validation_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.validation_retrieval_losses) / len(self.validation_retrieval_losses) if len(self.validation_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_validation_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "            self.wandb_run.log({f\"avg_validation_loss\": epoch_loss})\n",
        "            if(len(self.validation_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_validation_retrieval_loss\": loss_retrieval})\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.validation_loss = []\n",
        "        self.validation_retrieval_losses = []\n",
        "\n",
        "    def on_test_epoch_end(self) -> None:\n",
        "      epoch_accuracy = sum(acc for acc in self.test_accuracy) / len(self.test_accuracy) if len(self.test_accuracy) > 0 else float('nan')\n",
        "      self.log(f\"avg_test_accuracy\", epoch_accuracy, batch_size=self.config['batch_size'])\n",
        "\n",
        "      print(f\"| Epoch {self.current_epoch} | {'TEST'}\")\n",
        "\n",
        "      best_accuracy = getattr(self, f\"best_test_accuracy\")\n",
        "\n",
        "      # Update best loss if current loss is better\n",
        "      if epoch_accuracy < best_accuracy:\n",
        "          setattr(self, f\"best_test_accuracy\", epoch_accuracy)\n",
        "\n",
        "      print(f\"\\t- Total Accuracy => {epoch_accuracy:.4f}\")\n",
        "\n",
        "      if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "          self.wandb_run.log({f\"avg_test_accuracy\": epoch_accuracy})\n",
        "\n",
        "      # Empty the list of the losses for the next epoch\n",
        "      self.test_accuracy = []\n",
        "\n",
        "    def save_model_pytorch_api(self, name='best', epoch=''):\n",
        "        torch.save(self.state_dict(), f'{name}_ep-{epoch}.pt')\n",
        "\n",
        "    # In this case, we configure the model in 3 ways:\n",
        "    #   - Adam 8-Bit\n",
        "    #   - AdamW 8-Bit\n",
        "    #   - Normal Adam\n",
        "    #   - Normal AdamW\n",
        "    # There is an option to also set the weight decay for the parameters\n",
        "    def configure_optimizers(self):\n",
        "\n",
        "        # # ---- Grouped Parameters instead of self.parameters() directly to Adam ----\n",
        "        # decay_parameters = get_parameter_names(self.model, [nn.LayerNorm])\n",
        "        # decay_parameters = [name for name in decay_parameters if \"bias\" not in name]\n",
        "\n",
        "        # optimizer_grouped_parameters = [\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n not in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        # ]\n",
        "        # # ---------------------------------------------------------------------------\n",
        "\n",
        "        # Configs\n",
        "        epsilon = 1e-8  # Default\n",
        "        lr = 3e-4       # Default: 1e-3 | Can be changed also to 3e-4\n",
        "\n",
        "        # # Adam 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.Adam8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # AdamW 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.AdamW8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # Normal Adam\n",
        "        # adam_bnb_optim = bnb.optim.Adam(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # Standard PyTorch AdamW\n",
        "        adam_bnb_optim = torch.optim.AdamW(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # # Standard PyTorch Adam\n",
        "        # adam_bnb_optim = torch.optim.Adam(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # TODO: Add a Linear Warmup LR (maybe https://lightning-flash.readthedocs.io/en/stable/api/generated/flash.core.optimizers.LinearWarmupCosineAnnealingLR.html)\n",
        "        scheduler = pl_bolts.optimizers.lr_scheduler.LinearWarmupCosineAnnealingLR(adam_bnb_optim, warmup_epochs=2, warmup_start_lr=0.0, eta_min=0.0, max_epochs=self.trainer.max_epochs)\n",
        "        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(adam_bnb_optim, T_max=self.trainer.max_epochs)\n",
        "\n",
        "        return {\"optimizer\": adam_bnb_optim, \"lr_scheduler\": scheduler}\n",
        "\n",
        "    # Print number of trainable parameters\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return self.encoder.get_n_trainable_parameters() + self.decoder.get_n_trainable_parameters()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UShZhADWb25V",
        "outputId": "d426d8c3-d0dc-455c-a0f4-7c9b6b35c850"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Encoder info:\n",
            "trainable params: 1,327,104 || all params: 620,666,112 || trainable%: 0.21381931030898624\n",
            "Decoder info:\n",
            "trainable params: 1,327,104 || all params: 248,904,960 || trainable%: 0.5331770005708203\n"
          ]
        }
      ],
      "source": [
        "# Our Model\n",
        "dsi_config['encoder_model'] = \"google/switch-base-8\"\n",
        "dsi_config['decoder_model'] = \"google/flan-t5-base\"\n",
        "\n",
        "encoder = SwitchTransformer(model_id=dsi_config['encoder_model'], MAX_LENGTH=dsi_config[\"MAX_LENGTH\"])\n",
        "decoder = FlanT5(model_id=dsi_config['decoder_model'], MAX_LENGTH=dsi_config[\"MAX_LENGTH\"])\n",
        "dsi_config[\"hidden_size\"] = decoder.model.config.hidden_size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WZRHTOHtlR1N"
      },
      "outputs": [],
      "source": [
        "# Create actual DSI Model\n",
        "dsi_model = DSI_EncoderDecoder(dsi_config,\n",
        "                        encoder=encoder,\n",
        "                        decoder=decoder,\n",
        "                        query_doc_df=query_doc_df,\n",
        "                        training_dataloader_with_queries_filtered=train_dataloader_filtered,)\n",
        "\n",
        "total_params = dsi_model.get_n_trainable_parameters()\n",
        "total_params = \"{:,}\".format(total_params).replace(\",\", \".\")\n",
        "print(\"--------------------------------------------------------------\")\n",
        "print(f\"Total number of trainable parameters for DSI Model: {total_params}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T1h1yGNFb25W"
      },
      "source": [
        "### Discriminative Approach\n",
        "The Discriminative approach aims at adding a Linear Projection of the Decoder embeddings (through a Linear Layer) in the Doc_ID space.<br>\n",
        "By doing so, we simply use a SoftMax function to get the likelihood of each Doc_ID, given the query"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PTj5_lPTb25W"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "  -- DSI Model for Discriminative Training --\n",
        "\n",
        "  Parameters:\n",
        "    - config: The configuration dictionary populated throughout the whole Code\n",
        "    - model: The model to use (populated when instantiating the Foundation Model)\n",
        "    - query_doc_df: The dataframe containing the queries\n",
        "        This is needed for retrieving DocIDs during the retrieval task\n",
        "    - training_dataloader_with_queries_filtered: The dataloader containing the filtered documents with queries\n",
        "        Due to Lightning's limitations, we need to pass the dataloader to the model in order to be able to iterate over it\n",
        "'''\n",
        "class DSI_Discriminative(pl.LightningModule):\n",
        "    def __init__(self,\n",
        "                 config: dict,\n",
        "                 model: pl.LightningModule = None,\n",
        "                 query_doc_df: pd.DataFrame = None,\n",
        "                 training_dataloader_with_queries_filtered: DataLoader = None,\n",
        "                 ):\n",
        "        super().__init__()\n",
        "\n",
        "        # Variables\n",
        "        self.config = config\n",
        "        self.MAX_LENGTH = config[\"MAX_LENGTH\"]\n",
        "\n",
        "        self.token_tokenization_max_length = config['token_tokenization_max_length'] # Define the label max_length of token for the tokenization\n",
        "        self.train_indexing_retrieval_ratio = config[\"train_indexing_retrieval_ratio\"] # Docs with Query for Training\n",
        "        self.train_begin_retrieval = config[\"train_begin_retrieval\"] # At which index the retrieval task begins\n",
        "\n",
        "        self.enable_multitask_prompting = config[\"enable_multitask_prompting\"]\n",
        "        if(self.enable_multitask_prompting):\n",
        "          self.multitask_prompting_indexing = config[\"multitask_prompting_indexing\"]\n",
        "          self.multitask_prompting_retrieval = config[\"multitask_prompting_retrieval\"]\n",
        "\n",
        "        # Model\n",
        "        self.model = model.to(device)\n",
        "\n",
        "        # Generative or Discriminative Training\n",
        "        self.cross_entropy_loss = nn.CrossEntropyLoss()           # Cross Entropy Loss for Discriminative Training\n",
        "        self.batch_norm = nn.BatchNorm1d(config[\"hidden_size\"])   # Batch Normalization for Discriminative Training\n",
        "        self.relu = nn.ReLU()                                     # ReLU for Discriminative Training\n",
        "        self.linear_layer = nn.Linear(config[\"hidden_size\"], config['labels_count'], device=device) # Linear Layer for Discriminative Training\n",
        "        self.softmax = nn.Softmax(dim=1) # Softmax for Discriminative Training\n",
        "        self.dropout = nn.Dropout(p=0.2) # Dropout for Discriminative Training\n",
        "\n",
        "        # Dataloader\n",
        "        # We need to keep a copy of the training dataloader, due to the fact that we\n",
        "        # must reset the training dataloader every time we finish an epoch\n",
        "        # If we don't do this, we will have a StopIteration error\n",
        "        self.copy_training_dataloader = training_dataloader_with_queries_filtered\n",
        "        self.training_dataloader_with_queries_filtered = iter(training_dataloader_with_queries_filtered)\n",
        "\n",
        "        # Log losses\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "        self.validation_retrieval_losses = []\n",
        "        self.test_retrieval_losses = []\n",
        "\n",
        "        self.training_loss = []\n",
        "        self.validation_loss = []\n",
        "        self.test_accuracy = []\n",
        "\n",
        "        self.best_training_loss = float('inf')\n",
        "        self.best_validation_loss = float('inf')\n",
        "        self.best_test_accuracy = 0.0\n",
        "\n",
        "        # WandB\n",
        "        self.enable_wandb = config[\"wandb_configs\"][\"enable\"]\n",
        "        if(self.enable_wandb):\n",
        "          self.wandb_run = wandb.init(\n",
        "              project='DeepLearning-DSI',\n",
        "              group=config[\"wandb_configs\"][\"group_id\"],\n",
        "              config=config,)\n",
        "        else:\n",
        "          self.wandb_run = None\n",
        "\n",
        "        # Dataframes\n",
        "        self.query_doc_df = query_doc_df  # Query - Doc ID Dataframe\n",
        "        self.top_100_df = None            # Populated only if the test_type is \"top_k\"\n",
        "\n",
        "        # The current task\n",
        "        self.current_task = \"indexing\"  # \"indexing\" / \"retrieval\"\n",
        "        self.test_type = \"accuracy\"     # \"accuracy\" / \"top_k\"\n",
        "\n",
        "        # Create a Dict DocID2Index and Index2DocID\n",
        "        # Useful when we need to retrieve the DocID from the index and vice versa (Test phase)\n",
        "        self.docID2label = {}\n",
        "        self.label2DocID = {}\n",
        "        filtered_docs = train_docs_with_queries_dataset.filtered_docs_with_queries\n",
        "        for index, row in filtered_docs.iterrows():\n",
        "          self.docID2label[row['Doc_ID']] = str(row[config['document_id_representation_strategy']])\n",
        "          self.label2DocID[str(row[config['document_id_representation_strategy']])] = row['Doc_ID']\n",
        "\n",
        "        # Generation variables\n",
        "        self.num_return_sequences = 5\n",
        "        self.num_beams = 5\n",
        "\n",
        "    '''\n",
        "      The forward method of the DSI model that differentiate slightly the procedure to\n",
        "      follow according to the task that has to be performed (indexing/retrieval)\n",
        "    '''\n",
        "    def forward(self, batch, training=True):\n",
        "\n",
        "        torch.cuda.empty_cache() # Clear CUDA cache before forward pass\n",
        "\n",
        "        if self.current_task == \"indexing\": # Indexing Task\n",
        "          # T5 Prompting Setup as suggested in the paper\n",
        "          if self.enable_multitask_prompting:\n",
        "            docs_body = [f\"{self.multitask_prompting_indexing} {item}\" for item in batch[train_dataset.doc_rep_engine.strategy]]\n",
        "          else:\n",
        "            docs_body = batch[train_dataset.doc_rep_engine.strategy]\n",
        "\n",
        "          # Labels are the Doc_ID(s) processed with the strategy chosen\n",
        "          # labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          docs_input = self.model.tokenizer(docs_body, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "          # labels_inputs = self.model.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=len(str(self.token_tokenization_max_length)), padding=\"max_length\", truncation=True)\n",
        "\n",
        "          # labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "          # labels_att_mask = labels_inputs['attention_mask'].to(device)\n",
        "          input_ids = docs_input['input_ids'].to(device)\n",
        "          attention_mask = docs_input['attention_mask'].to(device)\n",
        "\n",
        "          # labels_input_ids[labels_input_ids == self.model.tokenizer.pad_token_id] = -100\n",
        "\n",
        "          # In this case, we calculate the output of the model, regardless of the training phase\n",
        "          # Then, we pass the output to the discriminative_step method to compute the loss\n",
        "          output = self.model.model(\n",
        "              input_ids,                     # Input_ids coming from Encoder Tokenizer\n",
        "              attention_mask=attention_mask, # Attention mask coming from Encoder Tokenizer\n",
        "              decoder_input_ids=torch.zeros_like(input_ids, device=device),\n",
        "              decoder_attention_mask=None,\n",
        "              labels=None,                   # Labels are not since we do not want to compute the loss\n",
        "              output_hidden_states=True\n",
        "          )\n",
        "          # If training, run a discriminative step to compute the loss\n",
        "          if(training):\n",
        "              loss, _ = self.discriminative_step(batch['Doc_ID'], output)\n",
        "              return loss\n",
        "          else: # If not training, run a discriminative step to compute the prediction\n",
        "              _, prediction = self.discriminative_step(batch['Doc_ID'], output)\n",
        "              return prediction\n",
        "        else:\n",
        "          # Retrival Task. The process is the exact same as the indexing task,\n",
        "          # but we need to consider the queries associated to the documents, not the documents themselves.\n",
        "          queries = batch[\"Query\"]\n",
        "\n",
        "          queries = [str(item[0]) for item in queries]\n",
        "          if self.enable_multitask_prompting:\n",
        "            queries = [f\"{self.multitask_prompting_retrieval} {query}\" for query in queries]\n",
        "\n",
        "          # labels = [str(item) for item in batch[train_dataset.docID_rep_engine.strategy]]\n",
        "\n",
        "          queries_input = self.model.tokenizer(queries, add_special_tokens=False, return_tensors='pt', max_length=self.MAX_LENGTH, padding=\"max_length\", truncation=True)\n",
        "          # labels_inputs = self.model.tokenizer(labels, add_special_tokens=False, return_tensors='pt', max_length=len(str(self.token_tokenization_max_length)), padding=\"max_length\", truncation=True)\n",
        "\n",
        "          # labels_input_ids = labels_inputs[\"input_ids\"].to(device)\n",
        "          # labels_att_mask = labels_inputs['attention_mask'].to(device)\n",
        "          # labels_input_ids[labels_input_ids == self.model.tokenizer.pad_token_id] = -100\n",
        "          input_ids = queries_input['input_ids'].to(device)\n",
        "          attention_mask = queries_input['attention_mask'].to(device)\n",
        "\n",
        "          # For training. Also includes loss computation\n",
        "          output = self.model.model(\n",
        "              input_ids,\n",
        "              attention_mask=attention_mask,\n",
        "              decoder_input_ids=torch.zeros_like(input_ids, device=device),\n",
        "              decoder_attention_mask=None,\n",
        "              labels=None,\n",
        "              output_hidden_states=True\n",
        "          )\n",
        "          if(training):\n",
        "              loss, _ = self.discriminative_step(batch['Doc_ID'], output)\n",
        "              return loss\n",
        "          else:\n",
        "              _, prediction = self.discriminative_step(batch['Doc_ID'], output)\n",
        "              return prediction\n",
        "\n",
        "    '''\n",
        "      A retrieval batch is created starting from a batch considering the questions associated to doc\n",
        "      in the current batch.\n",
        "    '''\n",
        "    def prepare_sample_for_retrieval(self, batch):\n",
        "        batch['Queries_ID'] = []\n",
        "        batch['Query'] = []\n",
        "        doc_ids = batch['Doc_ID']\n",
        "\n",
        "        for doc_id in doc_ids:\n",
        "            query_ids = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query_ID'].values\n",
        "\n",
        "            if(len(query_ids) == 0):\n",
        "              raise Exception(f\"Document {doc_id} without queries!\")\n",
        "            else:\n",
        "              queries = self.query_doc_df[self.query_doc_df['Doc_ID'] == doc_id]['Query'].values\n",
        "\n",
        "              batch['Queries_ID'].append(query_ids)\n",
        "              batch['Query'].append(queries)\n",
        "\n",
        "        return batch\n",
        "\n",
        "    # Generate the target tensor for the discriminative step\n",
        "    # This is needed such that we can compute the loss\n",
        "    # This will be a one hot encoded tensor with 1s at the index of the Doc_ID\n",
        "    def generate_target_tensor(self, document_ids: list):\n",
        "        x = torch.zeros((len(document_ids), self.config['labels_count']), device=device)\n",
        "\n",
        "        for index, id in enumerate(document_ids):\n",
        "            x[index, list(self.docID2label.keys()).index(id)] = 1\n",
        "\n",
        "        return x\n",
        "\n",
        "    def discriminative_step(self, list_doc_ids, output):\n",
        "        # Stack the last 4 Decoder Hidden Layers and take the max value\n",
        "        decoder_hidden_layer = torch.mean(torch.stack(output.decoder_hidden_states[-4:]), dim=0)\n",
        "        decoder_hidden_layer = torch.max(decoder_hidden_layer, dim=1)[0] # Dim = [batch_size, hidden_size]\n",
        "        decoder_hidden_layer = decoder_hidden_layer.to(device)\n",
        "\n",
        "        # Pass the max_decoder to the Linear Layer + Softmax + BatchNorm\n",
        "        x = self.batch_norm(decoder_hidden_layer)   # [batch_size, hidden_size]\n",
        "        x = self.relu(x)                            # [batch_size, hidden_size]\n",
        "        x = self.linear_layer(x)                    # [batch_size, labels_count]\n",
        "        x = self.dropout(x)                         # [batch_size, labels_count]\n",
        "        prediction = self.softmax(x)                # Apply the SoftMax on top of the Linear Layer\n",
        "\n",
        "        # Generate a target tensor starting from the DocIDs in the batch\n",
        "        # It will be a one hot encoded tensor with 1s at the index of the Doc_ID\n",
        "        target_tensor = self.generate_target_tensor(list_doc_ids)\n",
        "\n",
        "        # Get the indices of the max value in the target tensor and the prediction tensor\n",
        "        target_indices = torch.argmax(target_tensor, dim=1)   # Indices in target tensor\n",
        "        # prediction_indices = torch.argmax(prediction, dim=1)  # Indices in prediction tensor\n",
        "\n",
        "        # CrossEntropyLoss between the indices and the target tensor\n",
        "        loss = self.cross_entropy_loss(x, target_indices)\n",
        "\n",
        "        return loss, prediction  # [batch_size, labels_count]\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        # ========================================================================\n",
        "        # Retrieval step | Every self.train_indexing_retrieval_ratio batches\n",
        "        #\n",
        "        # Training step => Indexing all the documents and perform retrieval\n",
        "        #                  (at the end of each epoch considering the given ratio)\n",
        "        #                  on just the document chosen as training docs.\n",
        "        # ========================================================================\n",
        "\n",
        "        # Indexing step\n",
        "        self.current_task = 'indexing'\n",
        "\n",
        "        loss_idx = self.forward(batch, training=True)\n",
        "\n",
        "        self.training_indexing_losses.append(loss_idx.cpu().detach().item())\n",
        "        total_step_loss = loss_idx\n",
        "\n",
        "        # Switch to retrieval task if the batch index is greater than the train_begin_retrieval value\n",
        "        # This is done to alternate between indexing and retrieval tasks as suggested in the paper\n",
        "        if batch_idx >= self.train_begin_retrieval:\n",
        "          self.current_task = 'retrieval'\n",
        "\n",
        "          # If we are in the retrieval task, we need to prepare the sample for retrieval\n",
        "          # Hence, we need to consider the queries associated to the documents in the current batch\n",
        "          batch = next(self.training_dataloader_with_queries_filtered)\n",
        "          batch = self.prepare_sample_for_retrieval(batch)\n",
        "\n",
        "          loss_retrieval = self.forward(batch, training=True)\n",
        "\n",
        "          self.training_retrieval_losses.append(loss_retrieval.cpu().detach().item())\n",
        "          total_step_loss = (total_step_loss + loss_retrieval) / 2\n",
        "\n",
        "        self.log(\"training_loss\", total_step_loss.cpu().detach().item())\n",
        "        self.training_loss.append(total_step_loss.cpu().detach().item())\n",
        "\n",
        "        return total_step_loss\n",
        "\n",
        "    # Validation => Retrieval step for all the documents chosen to be in validation split\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "      self.current_task = 'retrieval'\n",
        "\n",
        "      # Retrieval step (same as retrieval step in training phase)\n",
        "      batch = self.prepare_sample_for_retrieval(batch)\n",
        "      bs = len(batch['Doc_ID'])\n",
        "\n",
        "      loss = self.forward(batch, training=True)\n",
        "\n",
        "      self.log(\"validation_loss\", loss.cpu().detach().item(), batch_size=bs)\n",
        "      self.validation_loss.append(loss.cpu().detach().item())\n",
        "      self.validation_retrieval_losses.append(loss.cpu().detach().item())\n",
        "\n",
        "      return loss\n",
        "\n",
        "    # Test => Retrieval step for all the documents chosen to be in test split\n",
        "    def test_step(self, batch, batch_idx):\n",
        "      self.current_task = 'retrieval'\n",
        "\n",
        "      # Retrieval step (same as retrieval step in training/validation phase)\n",
        "      batch = self.prepare_sample_for_retrieval(batch)\n",
        "      bs = len(batch['Doc_ID'])\n",
        "\n",
        "      predictions = self.forward(batch, training=False)\n",
        "\n",
        "      # Get the indices of the max value in the prediction tensor\n",
        "      prediction_indices = torch.argmax(predictions, dim=1)  # Indices in prediction tensor\n",
        "\n",
        "      if(self.test_type == \"accuracy\"):\n",
        "        accuracy = self.run_accuracy(prediction_indices, batch)\n",
        "      else:\n",
        "        accuracy = self.run_top_k(prediction_indices, batch)\n",
        "\n",
        "      # # Retrieve index from index2doc dictionary\n",
        "      # docs = [self.label2DocID[index.item()] for index in predictions]\n",
        "      # print(docs)\n",
        "\n",
        "      # self.log(\"test_loss\", loss.cpu().detach().item(), batch_size=bs)\n",
        "      # self.test_loss.append(loss.cpu().detach().item())\n",
        "      # self.test_retrieval_losses.append(loss.cpu().detach().item())\n",
        "      self.log(\"accuracy\", accuracy, batch_size=bs)\n",
        "      self.test_accuracy.append(accuracy)\n",
        "\n",
        "      return accuracy\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model are in the top-k relevant document of the queries from the batch\n",
        "    '''\n",
        "    def run_top_k(self, prediction_indices, batch):\n",
        "      # Before trying to find if the outputs are relevant w.r.t query, we first check if the query is present in the Top100 DF.\n",
        "      list_df = []\n",
        "      batch_queries = batch['Queries_ID']\n",
        "      for query in batch_queries:\n",
        "          q = self.top_100_df[self.top_100_df['Query_ID'] == query[0]]\n",
        "          if len(q) > 0:\n",
        "              list_df.append(q)\n",
        "          else:\n",
        "              print(f\"There are no queries associated with batch query {query[0]}!\")\n",
        "      batch_top_100 = pd.concat(list_df)\n",
        "\n",
        "      # Extract DocIDs from the output\n",
        "      docs = [self.label2DocID[list(self.label2DocID.keys())[index.item()]] for index in prediction_indices]\n",
        "\n",
        "      # Now we need to:\n",
        "      # - Take the first query from the batch\n",
        "      # - Check if the first model_returns is in the top 100 for the first query\n",
        "      # - Loop this process for all queries/outputs\n",
        "      correct = 0\n",
        "      for idx, query in enumerate(batch_queries):\n",
        "          if(docs[idx] in batch_top_100[batch_top_100['Query_ID'] == query[0]]['Doc_ID'].values):\n",
        "              correct += 1\n",
        "\n",
        "      final = float(correct / len(prediction_indices))\n",
        "      return final\n",
        "\n",
        "    '''\n",
        "      Checks if the output from the model is equal to the label present in the batch.\n",
        "    '''\n",
        "    def run_accuracy(self, prediction_indices, batch):\n",
        "      correct = 0\n",
        "\n",
        "      # Extract DocIDs from the output\n",
        "      docs = [self.label2DocID[list(self.label2DocID.keys())[index.item()]] for index in prediction_indices]\n",
        "\n",
        "      for idx, doc_id_pred in enumerate(docs):\n",
        "        if doc_id_pred == str(batch['Doc_ID'][idx]):\n",
        "          correct += 1\n",
        "\n",
        "      return correct / len(batch['Doc_ID'])\n",
        "\n",
        "    def on_train_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.training_loss) / len(self.training_loss) if len(self.training_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_training_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'TRAINING'}\")\n",
        "        best_loss = getattr(self, f\"best_training_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_training_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.training_indexing_losses)!=0):\n",
        "          loss_indexing = sum(loss for loss in self.training_indexing_losses) / len(self.training_indexing_losses) if len(self.training_indexing_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_indexing_loss\", loss_indexing, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Indexing Loss => {loss_indexing:.4f}\")\n",
        "\n",
        "        if(len(self.training_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.training_retrieval_losses) / len(self.training_retrieval_losses) if len(self.training_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_training_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and self.wandb_run is not None:\n",
        "            self.wandb_run.log({f\"avg_training_loss\": epoch_loss})\n",
        "            if(len(self.training_indexing_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_indexing_loss\": loss_indexing})\n",
        "            if(len(self.training_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_training_retrieval_loss\": loss_retrieval})\n",
        "            self.wandb_run.log({f\"epoch\": self.current_epoch})\n",
        "\n",
        "        # At the end of each training epoch we select a differet subset of training document\n",
        "        # on which perform the training retrieval. This ensure to have stocasticity selection\n",
        "        # while performing the retrieval on training docs\n",
        "        self.training_dataloader_with_queries_filtered = iter(self.copy_training_dataloader)\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.training_loss = []\n",
        "        self.training_indexing_losses = []\n",
        "        self.training_retrieval_losses = []\n",
        "\n",
        "    def on_validation_epoch_end(self) -> None:\n",
        "        epoch_loss = sum(loss for loss in self.validation_loss) / len(self.validation_loss) if len(self.validation_loss) > 0 else float('nan')\n",
        "        self.log(f\"avg_validation_loss\", epoch_loss, batch_size=self.config['batch_size'])\n",
        "\n",
        "        print(f\"| Epoch {self.current_epoch} | {'VALIDATION'}\")\n",
        "\n",
        "        best_loss = getattr(self, f\"best_validation_loss\")\n",
        "\n",
        "        # Update best loss if current loss is better\n",
        "        if epoch_loss < best_loss:\n",
        "            setattr(self, f\"best_validation_loss\", epoch_loss)\n",
        "\n",
        "        if(len(self.validation_retrieval_losses)!=0):\n",
        "          loss_retrieval = sum(loss for loss in self.validation_retrieval_losses) / len(self.validation_retrieval_losses) if len(self.validation_retrieval_losses) > 0 else float('nan')\n",
        "          self.log(f\"avg_validation_retrieval_loss\", loss_retrieval, batch_size=self.config['batch_size'])\n",
        "          print(f\"\\t- Retrieval Loss => {loss_retrieval:.4f}\")\n",
        "\n",
        "        print(f\"\\t- Total Loss => {epoch_loss:.4f}\")\n",
        "\n",
        "        if self.enable_wandb and self.wandb_run is not None:\n",
        "            self.wandb_run.log({f\"avg_validation_loss\": epoch_loss})\n",
        "            if(len(self.validation_retrieval_losses)!=0):\n",
        "              self.wandb_run.log({f\"avg_validation_retrieval_loss\": loss_retrieval})\n",
        "\n",
        "        # Empty the list of the losses for the next epoch\n",
        "        self.validation_loss = []\n",
        "        self.validation_retrieval_losses = []\n",
        "\n",
        "    def on_test_epoch_end(self) -> None:\n",
        "      epoch_accuracy = sum(acc for acc in self.test_accuracy) / len(self.test_accuracy) if len(self.test_accuracy) > 0 else float('nan')\n",
        "      self.log(f\"avg_test_accuracy\", epoch_accuracy, batch_size=self.config['batch_size'])\n",
        "\n",
        "      print(f\"| Epoch {self.current_epoch} | {'TEST'}\")\n",
        "\n",
        "      best_accuracy = getattr(self, f\"best_test_accuracy\")\n",
        "\n",
        "      # Update best loss if current loss is better\n",
        "      if epoch_accuracy < best_accuracy:\n",
        "          setattr(self, f\"best_test_accuracy\", epoch_accuracy)\n",
        "\n",
        "      print(f\"\\t- Total Accuracy => {epoch_accuracy:.4f}\")\n",
        "\n",
        "      if self.enable_wandb and not self.wandb_run._is_finished:\n",
        "          self.wandb_run.log({f\"avg_test_accuracy\": epoch_accuracy})\n",
        "\n",
        "      # Empty the list of the losses for the next epoch\n",
        "      self.test_accuracy = []\n",
        "\n",
        "    def save_model_pytorch_api(self, name='best', epoch=''):\n",
        "        torch.save(self.state_dict(), f'{name}_ep-{epoch}.pt')\n",
        "\n",
        "    # In this case, we configure the model in 3 ways:\n",
        "    #   - Adam 8-Bit\n",
        "    #   - AdamW 8-Bit\n",
        "    #   - Normal Adam\n",
        "    # There is an option to also set the weight decay for the parameters\n",
        "    def configure_optimizers(self):\n",
        "\n",
        "        # # ---- Grouped Parameters instead of self.parameters() directly to Adam ----\n",
        "        # decay_parameters = get_parameter_names(self.model, [nn.LayerNorm])\n",
        "        # decay_parameters = [name for name in decay_parameters if \"bias\" not in name]\n",
        "\n",
        "        # optimizer_grouped_parameters = [\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        #     {\n",
        "        #         \"params\": [p for n, p in self.model.named_parameters() if n not in decay_parameters],\n",
        "        #         \"weight_decay\": 0.0,\n",
        "        #     },\n",
        "        # ]\n",
        "        # # ---------------------------------------------------------------------------\n",
        "\n",
        "        # Configs\n",
        "        epsilon = 1e-8  # Default\n",
        "        lr = 3e-4       # Default: 1e-3 | Can be changed also to 3e-4\n",
        "\n",
        "        # # Adam 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.Adam8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # AdamW 8-Bit\n",
        "        # adam_bnb_optim = bnb.optim.AdamW8bit(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # # Normal Adam\n",
        "        # adam_bnb_optim = bnb.optim.Adam(\n",
        "        #     self.parameters(), # Can also be \"optimizer_grouped_parameters\"\n",
        "        #     eps=epsilon,\n",
        "        #     lr=lr,\n",
        "        # )\n",
        "\n",
        "        # Standard PyTorch AdamW\n",
        "        adam_bnb_optim = torch.optim.AdamW(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # # Standard PyTorch Adam\n",
        "        # adam_bnb_optim = torch.optim.Adam(self.parameters(), eps=epsilon, lr=lr)\n",
        "\n",
        "        # TODO: Add a Linear Warmup LR (maybe https://lightning-flash.readthedocs.io/en/stable/api/generated/flash.core.optimizers.LinearWarmupCosineAnnealingLR.html)\n",
        "        scheduler = pl_bolts.optimizers.lr_scheduler.LinearWarmupCosineAnnealingLR(adam_bnb_optim, warmup_epochs=2, warmup_start_lr=0.0, eta_min=0.0, max_epochs=self.trainer.max_epochs)\n",
        "        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(adam_bnb_optim, T_max=self.trainer.max_epochs)\n",
        "\n",
        "        return {\"optimizer\": adam_bnb_optim, \"lr_scheduler\": scheduler}\n",
        "\n",
        "    # Print number of trainable parameters\n",
        "    def get_n_trainable_parameters(self):\n",
        "        return self.model.get_n_trainable_parameters()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qlx6OySmb25W",
        "outputId": "ffb19080-d061-4a43-baa6-6cc70e6c4ee6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using AutoModelForSeq2SeqLM\n",
            "Encoder info:\n",
            "trainable params: 1,327,104 || all params: 248,904,960 || trainable%: 0.5331770005708203\n"
          ]
        }
      ],
      "source": [
        "# ---- DISCRIMINATIVE TRAINING ---- #\n",
        "dsi_config['model_name'] = 'google/flan-t5-base' # Only Flan-T5 Model if we want to perform discriminative training\n",
        "\n",
        "model = FoundationModel(model_id=dsi_config['model_name'], encoder_id=None, decoder_id=None, MAX_LENGTH=dsi_config[\"MAX_LENGTH\"])\n",
        "dsi_config[\"hidden_size\"] = model.model.config.hidden_size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        },
        "id": "ESH8VdGcb25W",
        "outputId": "cada7dfd-18cb-494a-bab0-085b40ba58a0"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "Tracking run with wandb version 0.16.3"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "Run data is saved locally in <code>h:\\.shortcut-targets-by-id\\1xApeZaLC56HWK_2RiQ-7AHot2Rra0RUj\\Progetto\\wandb\\run-20240216_085917-5jdjww1r</code>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "Syncing run <strong><a href='https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI/runs/5jdjww1r' target=\"_blank\">glistening-kumquat-79</a></strong> to <a href='https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              " View project at <a href='https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI' target=\"_blank\">https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI</a>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              " View run at <a href='https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI/runs/5jdjww1r' target=\"_blank\">https://wandb.ai/sapienza_ml_2022_23/DeepLearning-DSI/runs/5jdjww1r</a>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--------------------------------------------------------------\n",
            "Total number of trainable parameters for DSI Model: 1.327.104\n"
          ]
        }
      ],
      "source": [
        "dsi_model = DSI_Discriminative(dsi_config,\n",
        "                        model=model,\n",
        "                        query_doc_df=query_doc_df,\n",
        "                        training_dataloader_with_queries_filtered=train_dataloader_filtered,)\n",
        "\n",
        "total_params = dsi_model.get_n_trainable_parameters()\n",
        "total_params = \"{:,}\".format(total_params).replace(\",\", \".\")\n",
        "print(\"--------------------------------------------------------------\")\n",
        "print(f\"Total number of trainable parameters for DSI Model: {total_params}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5FaLhlptjIm"
      },
      "source": [
        "### Show a simple example of training/inference\n",
        "- If you would like to try a forward pass for training, set _training=`False`_, otherwise set it to `True` for inference purposes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bfTQQiTVMEa0"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "    Obtain the top 100 documents for each query ID in the training set.\n",
        "'''\n",
        "def obtain_top_100_df():\n",
        "    # Obtain queries\n",
        "    if(student == 'Professor'):\n",
        "        docv2_train_top = \"/content/top100.train.txt\"\n",
        "    else:\n",
        "        docv2_train_top = pathjoin(F\"/content/drive/MyDrive/{proj_dict[student][0]}\", \"top100.train.txt\") # qid, “Q0”, docid, rank, score, runstring\n",
        "\n",
        "    # LOCAL VARIABLES!! Use only for local execution\n",
        "    if (not usingColab):\n",
        "        docv2_train_top = \"top100.train.txt\"\n",
        "        docv2_train_top = \"top100.train.txt\"\n",
        "\n",
        "    # Here we merge our filtered dataset with the query_doc_df to obtain the dataset with the query IDs\n",
        "    # columns = ['Doc_ID', 'Doc Repres. Strategy', 'Doc ID Repres. Strategy', 'Query_ID']\n",
        "    dataset_with_queryIDs = pd.merge(train_docs_with_queries_dataset.filtered_docs_with_queries, query_doc_df[['Query_ID','Doc_ID']], on='Doc_ID')\n",
        "\n",
        "    # Read the top100.train.txt file. This file contains the top 100 documents for each query ID.\n",
        "    df_docv2_train_top100 = pd.read_csv(docv2_train_top, header=None, sep=\" \", names=[\"Query_ID\", \"Q0\", \"Doc_ID\", \"Rank\", \"Score\",\"Relevance\"])\n",
        "\n",
        "    # Filter out df_docv2_train_top100.\n",
        "    #\n",
        "    # After this line, we have a dataframe with the top 100 documents for each query ID, but only for the queries that we have in our dataset_with_queryIDs\n",
        "    df_docv2_train_top100_filtered = df_docv2_train_top100[df_docv2_train_top100['Query_ID'].isin(dataset_with_queryIDs['Query_ID'])]\n",
        "\n",
        "    # Here we might have two options:\n",
        "    #   - Keep all the 100 documents for each query ID (but in practice we have just documents from the chunk '00'.)\n",
        "    #   - Filter the query IDs to have just documents from the chunk '00'\n",
        "\n",
        "    # This is the second option\n",
        "    df_docv2_train_top100_filtered = df_docv2_train_top100_filtered[df_docv2_train_top100_filtered['Doc_ID'].str.startswith(f'msmarco_doc_{dsi_config[\"train_dataset_chunk\"]}_')]\n",
        "    return df_docv2_train_top100_filtered"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lUinCpqOKQaX"
      },
      "outputs": [],
      "source": [
        "# Initialize the Dataloaders such that in the cell below we can have different samples each time\n",
        "sample_doc_dataloader = iter(train_dataloader)\n",
        "sample_doc_query_dataloader = iter(val_dataloader)\n",
        "\n",
        "dsi_model = dsi_model.to(device) # Move the model to the GPU\n",
        "training = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OO68p2lDKvqn",
        "outputId": "bd775926-9f48-49d0-ef5b-d0c8caa08175"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(0.1489, device='cuda:0',\n",
            "       grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)\n"
          ]
        }
      ],
      "source": [
        "task = 'retrieval'\n",
        "\n",
        "if(task == 'retrieval'):\n",
        "  dsi_model.current_task = \"retrieval\"\n",
        "  batch = next(sample_doc_query_dataloader)\n",
        "  batch = dsi_model.prepare_sample_for_retrieval(batch)\n",
        "else:\n",
        "  dsi_model.current_task = \"indexing\"\n",
        "  batch = next(sample_doc_dataloader)\n",
        "\n",
        "# UserWarning: Input length of decoder_input_ids is 65, but `max_length` is set to 20. This can lead to unexpected behavior. You should consider increasing `max_new_tokens`.\n",
        "# UserWarning: Using the model-agnostic default `max_length`(=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
        "out = dsi_model.forward(batch, training=training)\n",
        "\n",
        "print(out)\n",
        "\n",
        "if(dsi_model.config['training_type'] == 'discriminative' and training == False):\n",
        "    dsi_model.test_type = 'top_k' # 'accuracy' / 'top_k'\n",
        "\n",
        "    # Get the indices of the max value in the prediction tensor\n",
        "    #\n",
        "    # If it gives error in this line, most probably you have set dsi_config['training_type'] = 'generative'\n",
        "    # It should be set to 'discriminative' in order to perform this test phase\n",
        "    prediction_indices = torch.argmax(out, dim=1)  # Indices in prediction tensor\n",
        "\n",
        "    if(dsi_model.test_type == \"accuracy\"):\n",
        "        accuracy = dsi_model.run_accuracy(prediction_indices, batch)\n",
        "    else:\n",
        "        if(dsi_model.top_100_df is None) and (dsi_model.test_type == \"top_k\"):\n",
        "            dsi_model.top_100_df = obtain_top_100_df()\n",
        "        accuracy = dsi_model.run_top_k(prediction_indices, batch)\n",
        "\n",
        "    print(accuracy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKdwDx76bViq"
      },
      "source": [
        "### Actual Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XYBGSWl28kRP",
        "outputId": "58480d81-79f2-4fa9-d96f-780e43788c52"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "GPU available: True (cuda), used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "HPU available: False, using: 0 HPUs\n"
          ]
        }
      ],
      "source": [
        "# Output of Model Name\n",
        "if(usingColab):\n",
        "    if(dsi_config['student'] == \"Professor\"):\n",
        "        parent_path = f\"/content/saved_models/\"\n",
        "    else:\n",
        "        parent_path = f\"/content/drive/MyDrive/{proj_dict[student][0]}/saved_models/\"\n",
        "else:\n",
        "    parent_path = f\"saved_models/\"\n",
        "\n",
        "model_name = f\"{dsi_config['document_representation_strategy']}-{dsi_config['document_id_representation_strategy']}-{str(len(train_dataset))}_rows-checkpoint\"\n",
        "\n",
        "# Add the runID to the model name if WANDB is enabled\n",
        "if(dsi_config['wandb_configs']['enable']):\n",
        "    model_name += f\"-runID_{dsi_model.wandb_run.name}\"\n",
        "\n",
        "# Checkpoint Callback\n",
        "checkpoint_callback = pl.callbacks.ModelCheckpoint(dirpath=f\"{parent_path+model_name}\", monitor=\"avg_validation_loss\", mode=\"min\", save_top_k=1)\n",
        "\n",
        "# Early Stopping Callback\n",
        "early_stop_callback = pl.callbacks.EarlyStopping(monitor=\"avg_validation_loss\", min_delta=0.0, patience=5, verbose=True, mode=\"min\")\n",
        "\n",
        "# Train parameters\n",
        "train_parameters = dict(\n",
        "    accelerator = 'auto',\n",
        "    max_epochs = 50,\n",
        "    callbacks=[checkpoint_callback], # [checkpoint_callback, early_stop_callback]\n",
        "    fast_dev_run=False,\n",
        "    gradient_clip_val=0.5,\n",
        ")\n",
        "\n",
        "trainer = pl.Trainer(**train_parameters)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlTMUcgivKnT"
      },
      "outputs": [],
      "source": [
        "# Fit the model! :)\n",
        "trainer.fit(\n",
        "    model=dsi_model,\n",
        "    train_dataloaders=train_dataloader,\n",
        "    val_dataloaders=val_dataloader\n",
        ")\n",
        "\n",
        "# Close the WandB run if it's enabled\n",
        "if(dsi_model.enable_wandb and dsi_model.wandb_run is not None):\n",
        "    dsi_model.wandb_run.finish()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oyp8Cdmpjej8"
      },
      "source": [
        "## Test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goc54mfpjej9"
      },
      "outputs": [],
      "source": [
        "'''\n",
        "    Obtain the top 100 documents for each query ID in the training set.\n",
        "'''\n",
        "def obtain_top_100_df():\n",
        "    # Obtain queries\n",
        "    if(student == 'Professor'):\n",
        "        docv2_train_top = \"/content/top100.train.txt\"\n",
        "    else:\n",
        "        docv2_train_top = pathjoin(F\"/content/drive/MyDrive/{proj_dict[student][0]}\", \"top100.train.txt\") # qid, “Q0”, docid, rank, score, runstring\n",
        "\n",
        "    # LOCAL VARIABLES!! Use only for local execution\n",
        "    if (not usingColab):\n",
        "        docv2_train_top = \"top100.train.txt\"\n",
        "        docv2_train_top = \"top100.train.txt\"\n",
        "\n",
        "    # Here we merge our filtered dataset with the query_doc_df to obtain the dataset with the query IDs\n",
        "    # columns = ['Doc_ID', 'Doc Repres. Strategy', 'Doc ID Repres. Strategy', 'Query_ID']\n",
        "    dataset_with_queryIDs = pd.merge(train_docs_with_queries_dataset.filtered_docs_with_queries, query_doc_df[['Query_ID','Doc_ID']], on='Doc_ID')\n",
        "\n",
        "    # Read the top100.train.txt file. This file contains the top 100 documents for each query ID.\n",
        "    df_docv2_train_top100 = pd.read_csv(docv2_train_top, header=None, sep=\" \", names=[\"Query_ID\", \"Q0\", \"Doc_ID\", \"Rank\", \"Score\",\"Relevance\"])\n",
        "\n",
        "    # Filter out df_docv2_train_top100.\n",
        "    #\n",
        "    # After this line, we have a dataframe with the top 100 documents for each query ID, but only for the queries that we have in our dataset_with_queryIDs\n",
        "    df_docv2_train_top100_filtered = df_docv2_train_top100[df_docv2_train_top100['Query_ID'].isin(dataset_with_queryIDs['Query_ID'])]\n",
        "\n",
        "    # Here we might have two options:\n",
        "    #   - Keep all the 100 documents for each query ID (but in practice we have just documents from the chunk '00'.)\n",
        "    #   - Filter the query IDs to have just documents from the chunk '00'\n",
        "\n",
        "    # This is the second option\n",
        "    df_docv2_train_top100_filtered = df_docv2_train_top100_filtered[df_docv2_train_top100_filtered['Doc_ID'].str.startswith(f'msmarco_doc_{dsi_config[\"train_dataset_chunk\"]}_')]\n",
        "    return df_docv2_train_top100_filtered"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D54jWUJMjej9"
      },
      "outputs": [],
      "source": [
        "dsi_model.test_type = \"accuracy\"     # \"accuracy\" or \"top_k\"\n",
        "dsi_model.model_found_label = 0      # Counter to simply keep track of at least how many Document can our Model retrieve\n",
        "if(dsi_model.top_100_df is None) and (dsi_model.test_type == \"top_k\"):\n",
        "    print(\"Obtaining Top 100 Documents for each Query ID in the Training Set...\")\n",
        "    dsi_model.top_100_df = obtain_top_100_df()\n",
        "\n",
        "trainer.test(dsi_model, dataloaders=test_dataloader)\n",
        "print(f\"In total, the Model was able to find {dsi_model.model_found_label/((len(test_dataloader)*dsi_config['batch_size'])*dsi_model.num_return_sequences)*100:.2f}% of the Documents, starting from its predictions.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8-bZd1bX4uw6"
      },
      "source": [
        "## Restore a Checkpoint and run a Test\n",
        "In this case we are restoring a Foundation Model with BERT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u4uGzmclMEa2",
        "outputId": "119f35ef-d5e5-408a-c47d-24496647b146"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Some weights of BertLMHeadModel were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['bert.encoder.layer.0.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.0.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.0.crossattention.output.dense.bias', 'bert.encoder.layer.0.crossattention.output.dense.weight', 'bert.encoder.layer.0.crossattention.self.key.bias', 'bert.encoder.layer.0.crossattention.self.key.weight', 'bert.encoder.layer.0.crossattention.self.query.bias', 'bert.encoder.layer.0.crossattention.self.query.weight', 'bert.encoder.layer.0.crossattention.self.value.bias', 'bert.encoder.layer.0.crossattention.self.value.weight', 'bert.encoder.layer.1.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.1.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.1.crossattention.output.dense.bias', 'bert.encoder.layer.1.crossattention.output.dense.weight', 'bert.encoder.layer.1.crossattention.self.key.bias', 'bert.encoder.layer.1.crossattention.self.key.weight', 'bert.encoder.layer.1.crossattention.self.query.bias', 'bert.encoder.layer.1.crossattention.self.query.weight', 'bert.encoder.layer.1.crossattention.self.value.bias', 'bert.encoder.layer.1.crossattention.self.value.weight', 'bert.encoder.layer.10.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.10.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.10.crossattention.output.dense.bias', 'bert.encoder.layer.10.crossattention.output.dense.weight', 'bert.encoder.layer.10.crossattention.self.key.bias', 'bert.encoder.layer.10.crossattention.self.key.weight', 'bert.encoder.layer.10.crossattention.self.query.bias', 'bert.encoder.layer.10.crossattention.self.query.weight', 'bert.encoder.layer.10.crossattention.self.value.bias', 'bert.encoder.layer.10.crossattention.self.value.weight', 'bert.encoder.layer.11.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.11.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.11.crossattention.output.dense.bias', 'bert.encoder.layer.11.crossattention.output.dense.weight', 'bert.encoder.layer.11.crossattention.self.key.bias', 'bert.encoder.layer.11.crossattention.self.key.weight', 'bert.encoder.layer.11.crossattention.self.query.bias', 'bert.encoder.layer.11.crossattention.self.query.weight', 'bert.encoder.layer.11.crossattention.self.value.bias', 'bert.encoder.layer.11.crossattention.self.value.weight', 'bert.encoder.layer.2.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.2.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.2.crossattention.output.dense.bias', 'bert.encoder.layer.2.crossattention.output.dense.weight', 'bert.encoder.layer.2.crossattention.self.key.bias', 'bert.encoder.layer.2.crossattention.self.key.weight', 'bert.encoder.layer.2.crossattention.self.query.bias', 'bert.encoder.layer.2.crossattention.self.query.weight', 'bert.encoder.layer.2.crossattention.self.value.bias', 'bert.encoder.layer.2.crossattention.self.value.weight', 'bert.encoder.layer.3.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.3.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.3.crossattention.output.dense.bias', 'bert.encoder.layer.3.crossattention.output.dense.weight', 'bert.encoder.layer.3.crossattention.self.key.bias', 'bert.encoder.layer.3.crossattention.self.key.weight', 'bert.encoder.layer.3.crossattention.self.query.bias', 'bert.encoder.layer.3.crossattention.self.query.weight', 'bert.encoder.layer.3.crossattention.self.value.bias', 'bert.encoder.layer.3.crossattention.self.value.weight', 'bert.encoder.layer.4.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.4.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.4.crossattention.output.dense.bias', 'bert.encoder.layer.4.crossattention.output.dense.weight', 'bert.encoder.layer.4.crossattention.self.key.bias', 'bert.encoder.layer.4.crossattention.self.key.weight', 'bert.encoder.layer.4.crossattention.self.query.bias', 'bert.encoder.layer.4.crossattention.self.query.weight', 'bert.encoder.layer.4.crossattention.self.value.bias', 'bert.encoder.layer.4.crossattention.self.value.weight', 'bert.encoder.layer.5.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.5.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.5.crossattention.output.dense.bias', 'bert.encoder.layer.5.crossattention.output.dense.weight', 'bert.encoder.layer.5.crossattention.self.key.bias', 'bert.encoder.layer.5.crossattention.self.key.weight', 'bert.encoder.layer.5.crossattention.self.query.bias', 'bert.encoder.layer.5.crossattention.self.query.weight', 'bert.encoder.layer.5.crossattention.self.value.bias', 'bert.encoder.layer.5.crossattention.self.value.weight', 'bert.encoder.layer.6.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.6.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.6.crossattention.output.dense.bias', 'bert.encoder.layer.6.crossattention.output.dense.weight', 'bert.encoder.layer.6.crossattention.self.key.bias', 'bert.encoder.layer.6.crossattention.self.key.weight', 'bert.encoder.layer.6.crossattention.self.query.bias', 'bert.encoder.layer.6.crossattention.self.query.weight', 'bert.encoder.layer.6.crossattention.self.value.bias', 'bert.encoder.layer.6.crossattention.self.value.weight', 'bert.encoder.layer.7.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.7.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.7.crossattention.output.dense.bias', 'bert.encoder.layer.7.crossattention.output.dense.weight', 'bert.encoder.layer.7.crossattention.self.key.bias', 'bert.encoder.layer.7.crossattention.self.key.weight', 'bert.encoder.layer.7.crossattention.self.query.bias', 'bert.encoder.layer.7.crossattention.self.query.weight', 'bert.encoder.layer.7.crossattention.self.value.bias', 'bert.encoder.layer.7.crossattention.self.value.weight', 'bert.encoder.layer.8.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.8.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.8.crossattention.output.dense.bias', 'bert.encoder.layer.8.crossattention.output.dense.weight', 'bert.encoder.layer.8.crossattention.self.key.bias', 'bert.encoder.layer.8.crossattention.self.key.weight', 'bert.encoder.layer.8.crossattention.self.query.bias', 'bert.encoder.layer.8.crossattention.self.query.weight', 'bert.encoder.layer.8.crossattention.self.value.bias', 'bert.encoder.layer.8.crossattention.self.value.weight', 'bert.encoder.layer.9.crossattention.output.LayerNorm.bias', 'bert.encoder.layer.9.crossattention.output.LayerNorm.weight', 'bert.encoder.layer.9.crossattention.output.dense.bias', 'bert.encoder.layer.9.crossattention.output.dense.weight', 'bert.encoder.layer.9.crossattention.self.key.bias', 'bert.encoder.layer.9.crossattention.self.key.weight', 'bert.encoder.layer.9.crossattention.self.query.bias', 'bert.encoder.layer.9.crossattention.self.query.weight', 'bert.encoder.layer.9.crossattention.self.value.bias', 'bert.encoder.layer.9.crossattention.self.value.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
            "The following encoder weights were not tied to the decoder ['bert/pooler']\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using EncoderDecoderModel, hence bert-base-uncased.\n",
            "Encoder info:\n",
            "trainable params: 1,327,104 || all params: 139,798,842 || trainable%: 0.9492954169105349\n"
          ]
        }
      ],
      "source": [
        "# Define the model\n",
        "#\n",
        "# 'model_name' = 'bert-base-uncased'  : The system will load the Bert-Uncased model\n",
        "# 'model_name' = 'google/flan-t5-base': The system will load the Flan-T5 model\n",
        "#\n",
        "# All the configs will be dynamically set according to the model chosen\n",
        "dsi_config['model_name'] = 'bert-base-uncased' # 'google/flan-t5-base' | 'bert-base-uncased'\n",
        "dsi_config['fine_tuning'] = False\n",
        "\n",
        "# ------------------------------------------------- #\n",
        "\n",
        "if(dsi_config['model_name'] != 'bert-base-uncased'):\n",
        "    model = FoundationModel(model_id=dsi_config['model_name'], encoder_id=None, decoder_id=None, MAX_LENGTH=dsi_config[\"MAX_LENGTH\"], fine_tuning=dsi_config['fine_tuning']) # Flan-T5 Model\n",
        "    dsi_config[\"hidden_size\"] = model.model.config.hidden_size\n",
        "else:\n",
        "    dsi_config['encoder_model'] = 'bert-base-uncased'\n",
        "    dsi_config['decoder_model'] = 'bert-base-uncased'\n",
        "    model = FoundationModel(model_id=None, encoder_id=dsi_config['encoder_model'], decoder_id=dsi_config['decoder_model'], MAX_LENGTH=dsi_config[\"MAX_LENGTH\"], fine_tuning=dsi_config['fine_tuning']) # Bert-Uncased Model\n",
        "    dsi_config[\"hidden_size\"] = model.model.encoder.config.hidden_size\n",
        "    dsi_config.pop(\"model_name\")\n",
        "\n",
        "# Create actual DSI Model\n",
        "dsi_config['wandb_configs']['enable'] = False # Disable WandB for the checkpoint model\n",
        "checkpoint_model = DSI_Model(dsi_config,\n",
        "                        model=model,\n",
        "                        query_doc_df=query_doc_df,\n",
        "                        training_dataloader_with_queries_filtered=train_dataloader_filtered,)\n",
        "\n",
        "# Convert the model to CUDA as expected by QLoRa\n",
        "checkpoint_model = checkpoint_model.to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QkjyhYyx2yjU",
        "outputId": "e0adb6ee-5309-434a-ac49-95ef837549b9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading checkpoint model from saved_models\\summarization-unstructured_atomic-10000_rows-checkpoint-runID_beautiful-caress-72\\epoch=6-step=1064.ckpt...\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 36,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Check the directory where the model is saved based on the student\n",
        "model_dir = 'summarization-unstructured_atomic-10000_rows-checkpoint-runID_beautiful-caress-72'\n",
        "\n",
        "if(usingColab):\n",
        "    if(dsi_config['student'] == \"Professor\"):\n",
        "        parent_path = f\"/content/saved_models/\"\n",
        "        DSI_URL = 'https://drive.google.com/uc?id=' + '1ehIC2ZsFalCwP1GfhE3vZa09rUBK_AO3' + '&export=download&confirm=t'\n",
        "        model_path = pathjoin(parent_path, model_dir)\n",
        "        if(not os.path.exists(model_path)):\n",
        "            os.makedirs(model_path)\n",
        "        gdown.download(DSI_URL, output=pathjoin(model_path, 'epoch=6-step=1064.ckpt'), quiet=False)\n",
        "    else:\n",
        "        parent_path = f\"/content/drive/MyDrive/{proj_dict[student][0]}/saved_models/\"\n",
        "else:\n",
        "    parent_path = f\"saved_models/\"\n",
        "\n",
        "# Retrieve the checkpoint path\n",
        "model_path = pathjoin(parent_path, model_dir, \"epoch=6-step=1064.ckpt\")\n",
        "\n",
        "# Load the state dict from PyTorch/Lightning API\n",
        "print(f\"Loading checkpoint model from {model_path}...\")\n",
        "#checkpoint_model.load_state_dict(torch.load(model_path))                     # ------> PyTorch API (Maybe expects .pt and not .ckpt)\n",
        "checkpoint_model.load_state_dict(torch.load(model_path)['state_dict']) # ------> PyTorch Lightning API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F262KBWtjej9"
      },
      "source": [
        "Make sure that test_dataloader is populated.<br>\n",
        "Also, make sure that the `obtain_top_100_df()` function is defined and loaded in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "referenced_widgets": [
            "e3b1b61433424bf9b7d28515f48de5d5"
          ]
        },
        "id": "R5ybUTbvjej-",
        "outputId": "db37c75a-b624-4164-fbc9-e92a1058be9b"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "GPU available: True (cuda), used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n",
            "HPU available: False, using: 0 HPUs\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
            "c:\\Users\\Gianmarco\\miniconda3\\Lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, test_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 12 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
            "  rank_zero_warn(\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e3b1b61433424bf9b7d28515f48de5d5",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Testing: 0it [00:00, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "| Epoch 0 | TEST\n",
            "\t- Total Accuracy => 0.0000\n"
          ]
        },
        {
          "data": {
            "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=\"font-weight: bold\">   Runningstage.testing    </span>┃<span style=\"font-weight: bold\">                           </span>┃\n",
              "┃<span style=\"font-weight: bold\">          metric           </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│<span style=\"color: #008080; text-decoration-color: #008080\">     avg_test_accuracy     </span>│<span style=\"color: #800080; text-decoration-color: #800080\">            0.0            </span>│\n",
              "│<span style=\"color: #008080; text-decoration-color: #008080\">       test_accuracy       </span>│<span style=\"color: #800080; text-decoration-color: #800080\">            0.0            </span>│\n",
              "└───────────────────────────┴───────────────────────────┘\n",
              "</pre>\n"
            ],
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1m  Runningstage.testing   \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m                           \u001b[0m┃\n",
              "┃\u001b[1m \u001b[0m\u001b[1m         metric          \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│\u001b[36m \u001b[0m\u001b[36m    avg_test_accuracy    \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m           0.0           \u001b[0m\u001b[35m \u001b[0m│\n",
              "│\u001b[36m \u001b[0m\u001b[36m      test_accuracy      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m           0.0           \u001b[0m\u001b[35m \u001b[0m│\n",
              "└───────────────────────────┴───────────────────────────┘\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In total, the Model was able to find 15.79% of the Documents, starting from its predictions.\n"
          ]
        }
      ],
      "source": [
        "checkpoint_model.test_type = \"accuracy\"     # \"accuracy\" or \"top_k\"\n",
        "checkpoint_model.model_found_label = 0      # Counter to simply keep track of at least how many Document can our Model retrieve\n",
        "if(checkpoint_model.top_100_df is None) and (checkpoint_model.test_type == \"top_k\"):\n",
        "    print(\"Obtaining Top 100 Documents for each Query ID in the Training Set...\")\n",
        "    checkpoint_model.top_100_df = obtain_top_100_df()\n",
        "\n",
        "trainer = pl.Trainer(accelerator = 'auto', fast_dev_run=False)\n",
        "trainer.test(checkpoint_model, dataloaders=test_dataloader)\n",
        "\n",
        "print(f\"In total, the Model was able to find {checkpoint_model.model_found_label/((len(test_dataloader)*dsi_config['batch_size'])*checkpoint_model.num_return_sequences)*100:.2f}% of the Documents, starting from its predictions.\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "VWJOpDVyK604",
        "mu86BnTiLIU_",
        "85QLEQCUJfBl",
        "GoKHVwm3LW69"
      ],
      "gpuType": "T4",
      "provenance": []
    },
    "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.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}