DevRev_Inference_with_Final_Pipeline_GPU.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ris27hav/devrevs_domain_specific_qa/blob/main/DevRev_Inference_with_Final_Pipeline_GPU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load packages and import libraries"
      ],
      "metadata": {
        "id": "M9jJuSCs26SW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install packages\n",
        "!pip install --upgrade --no-cache-dir gdown\n",
        "!pip install -U sentence-transformers\n",
        "!pip install -U faiss-gpu\n",
        "!pip install transformers sentencepiece\n",
        "!pip install optimum[onnxruntime-gpu]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1UjR7j2KkIIZ",
        "outputId": "18539307-7af9-4b22-ccc9-2eb44b817ced"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: gdown in /usr/local/lib/python3.8/dist-packages (4.4.0)\n",
            "Collecting gdown\n",
            "  Downloading gdown-4.6.0-py3-none-any.whl (14 kB)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from gdown) (4.64.1)\n",
            "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.8/dist-packages (from gdown) (4.6.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from gdown) (3.9.0)\n",
            "Requirement already satisfied: requests[socks] in /usr/local/lib/python3.8/dist-packages (from gdown) (2.25.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.8/dist-packages (from gdown) (1.15.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (2022.12.7)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (2.10)\n",
            "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (1.7.1)\n",
            "Installing collected packages: gdown\n",
            "  Attempting uninstall: gdown\n",
            "    Found existing installation: gdown 4.4.0\n",
            "    Uninstalling gdown-4.4.0:\n",
            "      Successfully uninstalled gdown-4.4.0\n",
            "Successfully installed gdown-4.6.0\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting sentence-transformers\n",
            "  Downloading sentence-transformers-2.2.2.tar.gz (85 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.0/86.0 KB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting transformers<5.0.0,>=4.6.0\n",
            "  Downloading transformers-4.26.0-py3-none-any.whl (6.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.3/6.3 MB\u001b[0m \u001b[31m55.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (4.64.1)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.13.1+cu116)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (0.14.1+cu116)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.21.6)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.0.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.7.3)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (3.7)\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.97-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m37.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting huggingface-hub>=0.4.0\n",
            "  Downloading huggingface_hub-0.12.0-py3-none-any.whl (190 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.3/190.3 KB\u001b[0m \u001b[31m5.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (3.9.0)\n",
            "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (23.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (4.4.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (6.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (2.25.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (2022.6.2)\n",
            "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
            "  Downloading tokenizers-0.13.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.6/7.6 MB\u001b[0m \u001b[31m45.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from nltk->sentence-transformers) (1.2.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.8/dist-packages (from nltk->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->sentence-transformers) (3.1.0)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.8/dist-packages (from torchvision->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (4.0.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (2.10)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (2022.12.7)\n",
            "Building wheels for collected packages: sentence-transformers\n",
            "  Building wheel for sentence-transformers (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sentence-transformers: filename=sentence_transformers-2.2.2-py3-none-any.whl size=125938 sha256=db487e4230e47737a9c14e85f31477f7b9c271012b92696471f69134afc46f04\n",
            "  Stored in directory: /root/.cache/pip/wheels/5e/6f/8c/d88aec621f3f542d26fac0342bef5e693335d125f4e54aeffe\n",
            "Successfully built sentence-transformers\n",
            "Installing collected packages: tokenizers, sentencepiece, huggingface-hub, transformers, sentence-transformers\n",
            "Successfully installed huggingface-hub-0.12.0 sentence-transformers-2.2.2 sentencepiece-0.1.97 tokenizers-0.13.2 transformers-4.26.0\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting faiss-cpu\n",
            "  Downloading faiss_cpu-1.7.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m17.0/17.0 MB\u001b[0m \u001b[31m16.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: faiss-cpu\n",
            "Successfully installed faiss-cpu-1.7.3\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.8/dist-packages (4.26.0)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.8/dist-packages (0.1.97)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers) (3.9.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (2022.6.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (1.21.6)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.8/dist-packages (from transformers) (0.12.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from transformers) (2.25.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.8/dist-packages (from transformers) (4.64.1)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from transformers) (23.0)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers) (0.13.2)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.4.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (1.24.3)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting optimum[onnxruntime]\n",
            "  Downloading optimum-1.6.3-py3-none-any.whl (227 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m227.4/227.4 KB\u001b[0m \u001b[31m5.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting coloredlogs\n",
            "  Downloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 KB\u001b[0m \u001b[31m5.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: huggingface-hub>=0.8.0 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (0.12.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (23.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.7.1)\n",
            "Requirement already satisfied: numpy<1.24.0 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.21.6)\n",
            "Requirement already satisfied: transformers[sentencepiece]>=4.26.0 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (4.26.0)\n",
            "Requirement already satisfied: torch>=1.9 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.13.1+cu116)\n",
            "Collecting evaluate\n",
            "  Downloading evaluate-0.4.0-py3-none-any.whl (81 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m81.4/81.4 KB\u001b[0m \u001b[31m10.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting onnxruntime>=1.9.0\n",
            "  Downloading onnxruntime-1.13.1-cp38-cp38-manylinux_2_27_x86_64.whl (4.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.5/4.5 MB\u001b[0m \u001b[31m46.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting protobuf==3.20.1\n",
            "  Downloading protobuf-3.20.1-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m48.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting datasets>=1.2.1\n",
            "  Downloading datasets-2.9.0-py3-none-any.whl (462 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m462.8/462.8 KB\u001b[0m \u001b[31m33.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting onnx\n",
            "  Downloading onnx-1.13.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (13.5 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.5/13.5 MB\u001b[0m \u001b[31m59.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: dill<0.3.7 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (0.3.6)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (1.3.5)\n",
            "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (4.64.1)\n",
            "Collecting xxhash\n",
            "  Downloading xxhash-3.2.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (213 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m213.0/213.0 KB\u001b[0m \u001b[31m20.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (2.25.1)\n",
            "Collecting responses<0.19\n",
            "  Downloading responses-0.18.0-py3-none-any.whl (38 kB)\n",
            "Requirement already satisfied: fsspec[http]>=2021.11.1 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (2023.1.0)\n",
            "Requirement already satisfied: pyarrow>=6.0.0 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (9.0.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (6.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (3.8.3)\n",
            "Collecting multiprocess\n",
            "  Downloading multiprocess-0.70.14-py38-none-any.whl (132 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m132.0/132.0 KB\u001b[0m \u001b[31m12.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.8.0->optimum[onnxruntime]) (4.4.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.8.0->optimum[onnxruntime]) (3.9.0)\n",
            "Requirement already satisfied: flatbuffers in /usr/local/lib/python3.8/dist-packages (from onnxruntime>=1.9.0->optimum[onnxruntime]) (1.12)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers[sentencepiece]>=4.26.0->optimum[onnxruntime]) (0.13.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers[sentencepiece]>=4.26.0->optimum[onnxruntime]) (2022.6.2)\n",
            "Requirement already satisfied: sentencepiece!=0.1.92,>=0.1.91 in /usr/local/lib/python3.8/dist-packages (from transformers[sentencepiece]>=4.26.0->optimum[onnxruntime]) (0.1.97)\n",
            "Collecting humanfriendly>=9.1\n",
            "  Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 KB\u001b[0m \u001b[31m11.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting onnx\n",
            "  Downloading onnx-1.12.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (13.1 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.1/13.1 MB\u001b[0m \u001b[31m71.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.8/dist-packages (from sympy->optimum[onnxruntime]) (1.2.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (22.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (1.3.3)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (6.0.4)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (1.3.1)\n",
            "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (2.1.1)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (1.8.2)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets>=1.2.1->optimum[onnxruntime]) (4.0.2)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets>=1.2.1->optimum[onnxruntime]) (2.10)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets>=1.2.1->optimum[onnxruntime]) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets>=1.2.1->optimum[onnxruntime]) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets>=1.2.1->optimum[onnxruntime]) (2022.12.7)\n",
            "Collecting urllib3<1.27,>=1.21.1\n",
            "  Downloading urllib3-1.26.14-py2.py3-none-any.whl (140 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m140.6/140.6 KB\u001b[0m \u001b[31m17.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets>=1.2.1->optimum[onnxruntime]) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets>=1.2.1->optimum[onnxruntime]) (2022.7.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7.3->pandas->datasets>=1.2.1->optimum[onnxruntime]) (1.15.0)\n",
            "Installing collected packages: xxhash, urllib3, protobuf, multiprocess, humanfriendly, onnx, coloredlogs, responses, onnxruntime, datasets, evaluate, optimum\n",
            "  Attempting uninstall: urllib3\n",
            "    Found existing installation: urllib3 1.24.3\n",
            "    Uninstalling urllib3-1.24.3:\n",
            "      Successfully uninstalled urllib3-1.24.3\n",
            "  Attempting uninstall: protobuf\n",
            "    Found existing installation: protobuf 3.19.6\n",
            "    Uninstalling protobuf-3.19.6:\n",
            "      Successfully uninstalled protobuf-3.19.6\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.9.2 requires protobuf<3.20,>=3.9.2, but you have protobuf 3.20.1 which is incompatible.\n",
            "tensorboard 2.9.1 requires protobuf<3.20,>=3.9.2, but you have protobuf 3.20.1 which is incompatible.\n",
            "googleapis-common-protos 1.58.0 requires protobuf!=3.20.0,!=3.20.1,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-translate 3.8.4 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-language 2.6.1 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-firestore 2.7.3 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-datastore 2.11.1 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-bigquery 3.4.2 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-cloud-bigquery-storage 2.18.1 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\n",
            "google-api-core 2.11.0 requires protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.19.5, but you have protobuf 3.20.1 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed coloredlogs-15.0.1 datasets-2.9.0 evaluate-0.4.0 humanfriendly-10.0 multiprocess-0.70.14 onnx-1.12.0 onnxruntime-1.13.1 optimum-1.6.3 protobuf-3.20.1 responses-0.18.0 urllib3-1.26.14 xxhash-3.2.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Import libraries\n",
        "import gdown\n",
        "import nltk\n",
        "import faiss\n",
        "import json\n",
        "import os\n",
        "import time\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import collections\n",
        "import json\n",
        "import re\n",
        "import string\n",
        "import timeit\n",
        "import tarfile\n",
        "import os\n",
        "\n",
        "from sentence_transformers import SentenceTransformer\n",
        "from transformers import AutoTokenizer, AutoModelForQuestionAnswering\n",
        "from optimum.onnxruntime import ORTModelForQuestionAnswering, ORTOptimizer\n",
        "from optimum.onnxruntime.configuration import OptimizationConfig\n",
        "\n",
        "from optimum.pipelines import pipeline\n",
        "from tqdm import tqdm\n",
        "from ast import literal_eval\n",
        "\n",
        "nltk.download('punkt')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lPHDMY7Xkpfh",
        "outputId": "c33905df-d1b9-4a7e-9a5f-71cbd7e06af0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
            "[nltk_data]   Unzipping tokenizers/punkt.zip.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper functions"
      ],
      "metadata": {
        "id": "vHtzxJq_pwHn"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UzHw72nd6_wA"
      },
      "source": [
        "### Sentence Encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_FarPEbyBsIh"
      },
      "source": [
        "For a given theme, break its paragraphs into sentences and store their paragraph id. Load sentence encoder and calculate embeddings for the sentences from paragraphs and the queries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5-iuMv8R7BPv"
      },
      "outputs": [],
      "source": [
        "def para_to_sentences(para):\n",
        "    \"\"\"Splits a paragraph into sentences.\"\"\"\n",
        "    para = para.replace('\\n', ' ').replace('\\t', ' ').replace('\\x00', ' ')\n",
        "    return nltk.sent_tokenize(para)\n",
        "\n",
        "def load_sents_from_para(paras):\n",
        "    \"\"\"Splits a list of paragraphs into sentences and returns the sentences\n",
        "    and their corresponding paragraph id\"\"\"\n",
        "    sents = []\n",
        "    para_id = []\n",
        "    for i,p in enumerate(paras):\n",
        "        new_sents = para_to_sentences(p['paragraph'])\n",
        "        sents += new_sents\n",
        "        para_id += [p['id']]*len(new_sents)\n",
        "    return sents, para_id"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fQKHrVVbDvwN"
      },
      "outputs": [],
      "source": [
        "def load_encoder():\n",
        "    \"\"\"Load mpnet-base-v2 Sentence Encoder\"\"\"\n",
        "    # model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
        "    gdown.load(\"https://drive.google.com/file/d/137tZvp-iTMR2xIogasglSH4jTTLW4_Sf/view\", fuzzy=True, use_cookies=False)\n",
        "    with ZipFile('/content/finetuned_mpnet_triplet.zip') as zobj:\n",
        "        zobj.extractall()\n",
        "    model = SentenceTransformer('/content/kaggle/working/finetuned_mpnet_triplet')\n",
        "    return model\n",
        "\n",
        "def get_embeddings(sents, model):\n",
        "    \"\"\"Generates embeddings for each sentence in the list of 768 dimesions\"\"\"\n",
        "    return model.encode(sents)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTEq0Khe7DM0"
      },
      "source": [
        "### Nearest Neighbour Search using FAISS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xo6xW9HsHoMl"
      },
      "source": [
        "Based on the embeddings calculated, indexes them based on L2 distance and then applies nearest neighbour search to get top k closest sentences for each query"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6SwA5Hhc7JMO"
      },
      "outputs": [],
      "source": [
        "def save_index(source_embeds, output_path):\n",
        "    \"\"\"Creates and saves the faiss L2 Index using source_embeds\"\"\"\n",
        "    index = faiss.IndexFlatL2(source_embeds.shape[1])\n",
        "    index.add(np.array(source_embeds))\n",
        "    faiss.write_index(index, output_path)\n",
        "\n",
        "def load_index(path):\n",
        "    \"\"\"Loads faiss index from the disk\"\"\"\n",
        "    index = faiss.read_index(path)\n",
        "    return index\n",
        "\n",
        "def get_k_nearest_neighbours(index, query_embeds, k = 10):\n",
        "    \"\"\"Returns k nearest neighbours of target_embeds in source_embeds\"\"\"\n",
        "    return index.search(np.array(query_embeds), k)\n",
        "\n",
        "def get_nearest_sentences(questions, theme):\n",
        "    \"\"\"Retrieve nearest sentences to the questions\"\"\"\n",
        "    ques_list = [q['question'] for q in questions]\n",
        "    ques_embed = get_embeddings(ques_list, sentence_encoder)\n",
        "    index = load_index(f'/content/indices/{theme}_para_l2_index')\n",
        "    return get_k_nearest_neighbours(index, ques_embed, k)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JyKtVXm661x3"
      },
      "source": [
        "### Load Existing QA and paragraphs data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xdk3n4f899fc"
      },
      "source": [
        "Load validation data for testing, based on missing data in the training data from squad 2.0 dataset. Round 1 data contains themes that are not present in training data. While, round 2 data contains themes that are present in training data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PDsZ3veewUjM"
      },
      "outputs": [],
      "source": [
        "def download_existing_data():\n",
        "    \"\"\"Download the test data (4 csv files)\"\"\"\n",
        "    ids = [\n",
        "        \"1LY3tKSdcVMb6Q38ZlgLEP9vYF4NPSP-Q\",\n",
        "        \"1feZeSoxc2zIBZ3_VPZtgtk1jUoml6J82\",\n",
        "        \"1deZKNy6oV3PSnfMpFP576m5XlFLiLWnT\",\n",
        "        \"1xHYf9vQbG_y9GCew0mavzdlDSeS9jOTP\"\n",
        "    ]\n",
        "    for id in ids:\n",
        "        url = f\"https://drive.google.com/u/1/uc?id={id}&export=download\"\n",
        "        gdown.download(url, quiet=True)\n",
        "\n",
        "\n",
        "def load_existing_data():\n",
        "    \"\"\"Load already answered questions and paragraphs, theme-wise.\n",
        "    Also breaks the paragraphs into sentences\"\"\"\n",
        "    paras, solved_ques = {}, {}\n",
        "    paragraphs = json.loads(pd.read_csv(\"input_paragraph.csv\").to_json(orient=\"records\"))\n",
        "    questions = json.loads(pd.read_csv(\"input_question.csv\").to_json(orient=\"records\"))\n",
        "    theme_intervals = json.loads(pd.read_csv(\"theme_interval.csv\").to_json(orient=\"records\"))\n",
        "    truth = pd.read_csv(\"ground_truth.csv\")\n",
        "    truth.fillna(value='', inplace=True)\n",
        "    truth.paragraph_id = truth.paragraph_id.apply(literal_eval)\n",
        "    truth.answers = truth.answers.apply(literal_eval)\n",
        "\n",
        "    for theme_interval in theme_intervals:\n",
        "        theme = theme_interval[\"theme\"]\n",
        "        theme_paras = [p for p in paragraphs if p[\"theme\"] == theme]\n",
        "        sents, para_id = load_sents_from_para(theme_paras)\n",
        "        paras[theme] = {\n",
        "            'id': para_id,\n",
        "            'sentences': sents\n",
        "        }\n",
        "\n",
        "        theme_ques = questions[int(theme_interval[\"start\"]) - 1: int(theme_interval[\"end\"])]\n",
        "        solved_ques[theme] = {\n",
        "            'id': [],\n",
        "            'question': [],\n",
        "            'paragraph_id': [],\n",
        "            'answers': []\n",
        "        }\n",
        "        for q in theme_ques:\n",
        "            truth_row = truth.loc[truth['question_id'] == q[\"id\"]].iloc[-1]\n",
        "            truth_paragraph_id = [ int(i) for i in truth_row[\"paragraph_id\"] ]\n",
        "            solved_ques[theme]['id'].append(q[\"id\"])\n",
        "            solved_ques[theme]['question'].append(q[\"question\"])\n",
        "            solved_ques[theme]['paragraph_id'].append(truth_paragraph_id)\n",
        "            solved_ques[theme]['answers'].append(truth_row[\"answers\"])\n",
        "\n",
        "    return paras, solved_ques\n",
        "\n",
        "\n",
        "def store_faiss_indices(paras, solved_ques, encoder):\n",
        "    \"\"\"Generates embeddings for paragraph sentences and queries. Then it creates\n",
        "    and saves the faiss index using them into disk\"\"\"\n",
        "    if not os.path.exists('/content/indices/'):\n",
        "        os.mkdir('/content/indices/')\n",
        "    for theme in paras:\n",
        "        theme_paras = paras[theme]\n",
        "        theme_ques = solved_ques[theme]\n",
        "\n",
        "        para_embeds = get_embeddings(theme_paras['sentences'], encoder)\n",
        "        output_path = f'/content/indices/{theme}_para_l2_index'\n",
        "        save_index(para_embeds, output_path)\n",
        "\n",
        "        ques_embeds = get_embeddings(theme_ques['question'], encoder)\n",
        "        output_path = f'/content/indices/{theme}_ques_l2_index'\n",
        "        save_index(ques_embeds, output_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uxnd1xuWqVeA"
      },
      "source": [
        "### Search previously answered queries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OSpVs8P8qoMh"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xccnbbAf7J1o"
      },
      "source": [
        "### Context Generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aJgYZNGdR5ni"
      },
      "source": [
        "Generates a context for a given query and its nearest neighbours. Also provides a method to get the paragraph id given the start idx of the answer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cdyftt2i7SMV"
      },
      "outputs": [],
      "source": [
        "def get_context(sents, para_ids, nearest_neighbours, distances):\n",
        "    \"\"\"Generate the context for a given query and store the para_id for\n",
        "    each sentence\"\"\"\n",
        "    context = \"\"\n",
        "    context_para_ids, sent_length = [], []\n",
        "    for sent_id, dist in zip(nearest_neighbours, distances):\n",
        "        if dist > distance_threshold*distances[0]:\n",
        "            break\n",
        "        context += sents[sent_id] + ' '\n",
        "        context_para_ids.append(para_ids[sent_id])\n",
        "        sent_length.append(len(sents[sent_id]))\n",
        "        if len(context.split()) >= context_length_threshold:\n",
        "            break\n",
        "    sum = -1\n",
        "    for i in range(len(sent_length)):\n",
        "        sum += sent_length[i] + 1\n",
        "        sent_length[i] = sum\n",
        "    return context.strip(), context_para_ids, sent_length\n",
        "\n",
        "\n",
        "def para_id_retriever(start_idx, sent_length, context_para_ids):\n",
        "    \"\"\"Given start index of the answer, return the id of the paragraph\n",
        "    in which the answer belongs\"\"\"\n",
        "    if start_idx == -1:\n",
        "        return -1\n",
        "    for j in range(len(sent_length)):\n",
        "        if start_idx <= sent_length[j]:\n",
        "            return context_para_ids[j]\n",
        "    return context_para_ids[-1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "raCs2XgbZpTG"
      },
      "source": [
        "### Load fine-tuned QA models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "actaNQ_CafqR"
      },
      "source": [
        "Given a theme, load the corresponding fine-tuned QA model and load the QA pipeline"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def download_fine_tuned_models():\n",
        "    \"\"\"Download and unzip cluster-wise fine-tuned QA models\"\"\"\n",
        "    urls = [\n",
        "        (\"1-7XfPhjfmUo8xz0iqmFHusbZ74q-SS3A\", \"zipped_0_11.tar.gz\"),\n",
        "        (\"1-BIhfqK992YZW1eWiG5yOCLiX8vOyrNI\", \"zipped_12_22.tar.gz\"),\n",
        "        (\"1-B8b2_s9i2pwTn7EPgzMg50nNM6Dp4B-\", \"zipped_23_34.tar.gz\"),\n",
        "        (\"1-KDxa6wWMGqrDR7ZJq-bYSyWaa_Zsikq\", \"zipped_35_42.tar.gz\")\n",
        "    ]\n",
        "    for url, filename in urls:\n",
        "        if not os.path.exists(filename):\n",
        "            link = f\"https://drive.google.com/u/1/uc?id={url}&export=download\"\n",
        "            gdown.download(link, quiet=True, use_cookies=False)\n",
        "            with tarfile.open(filename, 'r') as tar:\n",
        "                tar.extractall()\n",
        "            # os.remove(filename)\n",
        "\n",
        "def download_generic_model():\n",
        "    \"\"\"Download and optimize electra base model using onnx\"\"\"\n",
        "    model_id = 'PremalMatalia/electra-base-best-squad2'\n",
        "    save_path = \"/content/models/generic_model/\"\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "    ort_model = ORTModelForQuestionAnswering.from_pretrained(\n",
        "        model_id, from_transformers=True,\n",
        "        provider=\"CUDAExecutionProvider\"\n",
        "    )\n",
        "    optimizer = ORTOptimizer.from_pretrained(ort_model)\n",
        "    optimization_config = OptimizationConfig(optimization_level=99)\n",
        "    optimizer.optimize(save_dir=save_path, optimization_config=optimization_config)"
      ],
      "metadata": {
        "id": "uN7wGiaRuUn5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NNzi6Bz-mGED"
      },
      "outputs": [],
      "source": [
        "def load_models_mapping():\n",
        "    \"\"\"Loads map for checking cluster of a theme and vice versa\"\"\"\n",
        "    theme_to_cluster = {}\n",
        "    cluster_to_themes = {}\n",
        "    if not os.path.exists(\"clusters.json\"):\n",
        "        file_url = \"https://drive.google.com/file/d/1P6dp7f2m67-iPaUbaNZiDYTmTH7Mw9ec/view?usp=share_link\"\n",
        "        gdown.download(url=file_url, output='clusters.json', quiet=False, fuzzy=True)\n",
        "    with open('clusters.json') as fo:\n",
        "        map = json.load(fo)\n",
        "    for cluster, themes in map.items():\n",
        "        cluster = int(cluster)\n",
        "        if cluster not in cluster_to_themes:\n",
        "            cluster_to_themes[cluster] = []\n",
        "        for theme in themes:\n",
        "            theme_to_cluster[theme] = cluster\n",
        "            cluster_to_themes[cluster].append(theme)\n",
        "    return theme_to_cluster, cluster_to_themes\n",
        "\n",
        "\n",
        "def load_qa_model_pipeline(model_path):\n",
        "    \"\"\"Load QA model pipeline for a given cluster\"\"\"\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "    for i in range(5):\n",
        "        try:\n",
        "            model = ORTModelForQuestionAnswering.from_pretrained(\n",
        "                model_path, file_name=\"model_optimized.onnx\"\n",
        "            )\n",
        "        except:\n",
        "            continue\n",
        "        else:\n",
        "            break\n",
        "    optimum_qa = pipeline(\n",
        "        task = 'question-answering', model=model,\n",
        "        tokenizer=tokenizer, handle_impossible_answer=True\n",
        "    )\n",
        "    return optimum_qa"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Q_1CZAKZZ1w"
      },
      "source": [
        "## Execution"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Download sentence encoder model and fine-tuned QA models\n",
        "download_fine_tuned_models()\n",
        "download_generic_model()\n",
        "sentence_encoder = load_encoder()\n",
        "theme_to_cluster, cluster_to_themes = load_models_mapping()\n",
        "\n",
        "# Load existing QA pairs for themes and pre-process it\n",
        "download_existing_data()\n",
        "paras, solved_ques = load_existing_data()\n",
        "store_faiss_indices(paras, solved_ques, sentence_encoder)\n",
        "\n",
        "# Parameters for context generation\n",
        "k = 10\n",
        "distance_threshold = 1.5\n",
        "context_length_threshold = 205"
      ],
      "metadata": {
        "id": "YWnuk1j6qPuV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_theme_model(theme):\n",
        "    \"\"\"Load theme model if available, otherwise use generic model\"\"\"\n",
        "    if theme in theme_to_cluster:\n",
        "        cluster = theme_to_cluster[theme]\n",
        "        model_path = f'/content/models/electra-base-best-squad2-finetuned-squad-{cluster}'\n",
        "        if os.path.exists(model_path):\n",
        "            return load_qa_model_pipeline(model_path)\n",
        "    model_path = f'/content/models/generic_model'\n",
        "    return load_qa_model_pipeline(model_path)"
      ],
      "metadata": {
        "id": "p8RcyC4O2DFk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def pred_theme_ans(questions, theme_model, pred_out):\n",
        "    ann_inference_time, qna_inference_time = 0., 0.\n",
        "    theme = questions[0][\"theme\"]\n",
        "    print(f'Theme: {theme}')\n",
        "\n",
        "    # Nearest Neighbour Search\n",
        "    start_time = time.time()\n",
        "    D, I = get_nearest_sentences(questions, theme)\n",
        "    ann_inference_time = (time.time() - start_time)*1000.\n",
        "\n",
        "    # QA Model Prediction\n",
        "    start_time = time.time()\n",
        "    for i in tqdm(range(len(questions))):\n",
        "        q = questions[i]\n",
        "        # Context Generation\n",
        "        context, context_para_ids, sent_length = get_context(\n",
        "            paras[theme]['sentences'], paras[theme]['id'], I[i], D[i]\n",
        "        )\n",
        "        # Answer Prediction and Paragraph Retrieval\n",
        "        prediction = theme_model(question=q['question'], context=context)\n",
        "        ans = {\n",
        "            \"question_id\": q['id'],\n",
        "            \"answers\": prediction['answer'],\n",
        "            \"paragraph_id\": -1\n",
        "        }\n",
        "        if prediction['answer'] != \"\":\n",
        "            ans[\"paragraph_id\"] = para_id_retriever(\n",
        "                prediction['start'], sent_length, context_para_ids\n",
        "            )\n",
        "        pred_out.append(ans)\n",
        "\n",
        "    # Print Inference Time\n",
        "    qna_inference_time = (time.time() - start_time)*1000.\n",
        "    print(\n",
        "        f'Avg. ANN IT = {round(ann_inference_time/len(questions), 2)} ms, ' +\n",
        "        f'Avg. QnA IT = {round(qna_inference_time/len(questions),2)} ms\\n'\n",
        "    )"
      ],
      "metadata": {
        "id": "Aa4x0ljoIGpP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# NOT allowed to make changes.\n",
        "\n",
        "# All theme prediction.\n",
        "questions = json.loads(pd.read_csv(\"input_question.csv\").to_json(orient=\"records\"))\n",
        "theme_intervals = json.loads(pd.read_csv(\"theme_interval.csv\").to_json(orient=\"records\"))\n",
        "pred_out = []\n",
        "theme_inf_time = {}\n",
        "for theme_interval in theme_intervals[:2]:\n",
        "    theme_ques = questions[int(theme_interval[\"start\"]) - 1: int(theme_interval[\"end\"])]\n",
        "    theme = theme_ques[0][\"theme\"]\n",
        "    # Load model fine-tuned for this theme.\n",
        "    theme_model = get_theme_model(theme)\n",
        "    execution_time = timeit.timeit(lambda: pred_theme_ans(theme_ques, theme_model, pred_out), number=1)\n",
        "    theme_inf_time[theme_interval[\"theme\"]] = execution_time * 1000 # in milliseconds.\n",
        "pred_df = pd.DataFrame.from_records(pred_out)\n",
        "pred_df.fillna(value='', inplace=True)\n",
        "# Write prediction to a CSV file. Teams are required to submit this csv file.\n",
        "pred_df.to_csv('output_prediction.csv', index=False)"
      ],
      "metadata": {
        "id": "V4MGougaIImX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f3bc82db-fa0b-4cb0-837a-45859d6c285f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Theme: IPod\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 326/326 [03:23<00:00,  1.60it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Avg. ANN IT = 63.86 ms, Avg. QnA IT = 624.17 ms\n",
            "\n",
            "Theme: 2008_Sichuan_earthquake\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 521/521 [06:33<00:00,  1.33it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Avg. ANN IT = 91.31 ms, Avg. QnA IT = 754.41 ms\n",
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# NOT allowed to make changes.\n",
        "\n",
        "def normalize_answer(s):\n",
        "    \"\"\"Lower text and remove punctuation, articles and extra whitespace.\"\"\"\n",
        "    def remove_articles(text):\n",
        "        regex = re.compile(r'\\b(a|an|the)\\b', re.UNICODE)\n",
        "        return re.sub(regex, ' ', text)\n",
        "    def white_space_fix(text):\n",
        "        return ' '.join(text.split())\n",
        "    def remove_punc(text):\n",
        "        exclude = set(string.punctuation)\n",
        "        return ''.join(ch for ch in text if ch not in exclude)\n",
        "    def lower(text):\n",
        "        return text.lower()\n",
        "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
        "\n",
        "def get_tokens(s):\n",
        "    if not s: return []\n",
        "    return normalize_answer(s).split()\n",
        "\n",
        "def calc_f1(a_gold, a_pred):\n",
        "    gold_toks = get_tokens(a_gold)\n",
        "    pred_toks = get_tokens(a_pred)\n",
        "    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)\n",
        "    num_same = sum(common.values())\n",
        "    if len(gold_toks) == 0 or len(pred_toks) == 0:\n",
        "        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise\n",
        "        return int(gold_toks == pred_toks)\n",
        "    if num_same == 0:\n",
        "        return 0\n",
        "    precision = 1.0 * num_same / len(pred_toks)\n",
        "    recall = 1.0 * num_same / len(gold_toks)\n",
        "    f1 = (2 * precision * recall) / (precision + recall)\n",
        "    return f1\n",
        "\n",
        "def calc_max_f1(predicted, ground_truths):\n",
        "    max_f1 = 0\n",
        "    if len(ground_truths) == 0:\n",
        "        return len(predicted) == 0\n",
        "    for ground_truth in ground_truths:\n",
        "        f1 = calc_f1(predicted, ground_truth)\n",
        "        max_f1 = max(max_f1, f1)\n",
        "    return max_f1"
      ],
      "metadata": {
        "id": "lempoIKsIJ_G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# NOT allowed to make changes.\n",
        "\n",
        "# Evaluation methodology.\n",
        "metrics = {}\n",
        "pred = pd.read_csv(\"output_prediction.csv\")\n",
        "pred.fillna(value='', inplace=True)\n",
        "truth = pd.read_csv(\"ground_truth.csv\")\n",
        "truth.fillna(value='', inplace=True)\n",
        "truth.paragraph_id = truth.paragraph_id.apply(literal_eval)\n",
        "truth.answers = truth.answers.apply(literal_eval)\n",
        "questions = pd.read_csv(\"input_question.csv\")\n",
        "for idx in pred.index:\n",
        "    q_id = pred[\"question_id\"][idx]\n",
        "    q_rows = questions.loc[questions['id'] == q_id].iloc[-1]\n",
        "    theme = q_rows[\"theme\"]\n",
        "    predicted_paragraph = pred[\"paragraph_id\"][idx]\n",
        "    predicted_ans = pred[\"answers\"][idx]\n",
        "\n",
        "    if theme not in metrics.keys():\n",
        "        metrics[theme] = {\"true_positive\": 0, \"true_negative\": 0, \"total_predictions\": 0, \"f1_sum\": 0}\n",
        "\n",
        "    truth_row = truth.loc[truth['question_id'] == q_id].iloc[-1]\n",
        "    truth_paragraph_id = [ int(i) for i in truth_row[\"paragraph_id\"] ]\n",
        "    if predicted_paragraph in truth_paragraph_id:\n",
        "        # Increase TP for that theme.\n",
        "        metrics[theme][\"true_positive\"] = metrics[theme][\"true_positive\"] + 1\n",
        "    # -1 prediction in case there is no paragraph which can answer the query.\n",
        "    if predicted_paragraph == -1 and truth_row[\"paragraph_id\"] == []:\n",
        "        # Increase TN.\n",
        "        metrics[theme][\"true_negative\"] = metrics[theme][\"true_negative\"] + 1\n",
        "    # Increase total predictions for that theme.\n",
        "    metrics[theme][\"total_predictions\"] = metrics[theme][\"total_predictions\"] + 1\n",
        "    f1 = calc_max_f1(predicted_ans, truth_row[\"answers\"])\n",
        "    metrics[theme][\"f1_sum\"] = metrics[theme][\"f1_sum\"] + f1"
      ],
      "metadata": {
        "id": "-HA5KB3RIL4z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# NOT allowed to make changes.\n",
        "\n",
        "# Final score.\n",
        "inf_time_threshold = 1000.0 # milliseconds.\n",
        "final_para_score = 0.0\n",
        "final_qa_score = 0.0\n",
        "# Weight would stay hidden from teams.\n",
        "theme_weights = {\"Kubernetes\": 0.5, \"ChatGPT\": 0.4, \"Football world cup\": 0.1}\n",
        "for theme in metrics:\n",
        "    inf_time_score = 1.0\n",
        "    metric = metrics[theme]\n",
        "    para_score = (metric[\"true_positive\"] + metric[\"true_negative\"]) / metric[\"total_predictions\"]\n",
        "    qa_score = metric[\"f1_sum\"] / metric[\"total_predictions\"]\n",
        "    avg_inf_time = theme_inf_time[theme] / metric[\"total_predictions\"]\n",
        "    if avg_inf_time > inf_time_threshold:\n",
        "        inf_time_score = inf_time_threshold / avg_inf_time\n",
        "    final_qa_score += 1. * inf_time_score * qa_score\n",
        "    final_para_score += 1. * inf_time_score * para_score\n",
        "print (final_para_score/len(metrics))\n",
        "print (final_qa_score/len(metrics))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tuojd-v8INyd",
        "outputId": "3d5b07e3-2605-4904-dc39-65183186b9a8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.8739652390989485\n",
            "0.8408593320896819\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "CI7eZjF8JCSa"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "M9jJuSCs26SW",
        "UzHw72nd6_wA",
        "BTEq0Khe7DM0",
        "JyKtVXm661x3",
        "Uxnd1xuWqVeA",
        "raCs2XgbZpTG"
      ],
      "provenance": [],
      "authorship_tag": "ABX9TyPUI+dxbEBQPsE1TXb8wVQP",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}