DevRev_Inference_with_Final_Pipeline.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.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-cpu\n",
        "!pip install transformers sentencepiece\n",
        "!pip install optimum[onnxruntime]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1UjR7j2KkIIZ",
        "outputId": "a7f41fb4-8f99-44c5-bbda-e4391f6b7cb4"
      },
      "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.6.0)\n",
            "Requirement already satisfied: requests[socks] in /usr/local/lib/python3.8/dist-packages (from gdown) (2.25.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: tqdm in /usr/local/lib/python3.8/dist-packages (from gdown) (4.64.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.8/dist-packages (from gdown) (1.15.0)\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: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (2.10)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (1.26.14)\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: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.8/dist-packages (from requests[socks]->gdown) (1.7.1)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: sentence-transformers in /usr/local/lib/python3.8/dist-packages (2.2.2)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (0.1.97)\n",
            "Requirement already satisfied: transformers<5.0.0,>=4.6.0 in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (4.26.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (4.64.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.7.3)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.0.2)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (0.14.1+cu116)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (3.7)\n",
            "Requirement already satisfied: huggingface-hub>=0.4.0 in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (0.12.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from sentence-transformers) (1.21.6)\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: 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: filelock in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.4.0->sentence-transformers) (3.9.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: 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: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (0.13.2)\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",
            "Requirement already satisfied: click in /usr/local/lib/python3.8/dist-packages (from nltk->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from nltk->sentence-transformers) (1.2.0)\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: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (1.26.14)\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",
            "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",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: faiss-cpu in /usr/local/lib/python3.8/dist-packages (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: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers) (3.9.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: regex!=2019.12.17 in /usr/local/lib/python3.8/dist-packages (from transformers) (2022.6.2)\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: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.8/dist-packages (from transformers) (0.12.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from transformers) (23.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: 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: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (1.26.14)\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: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers) (2022.12.7)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: optimum[onnxruntime] in /usr/local/lib/python3.8/dist-packages (1.6.3)\n",
            "Requirement 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: torch>=1.9 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.13.1+cu116)\n",
            "Requirement already satisfied: coloredlogs in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (15.0.1)\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: packaging in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (23.0)\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: sympy in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.7.1)\n",
            "Requirement already satisfied: onnx in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.12.0)\n",
            "Requirement already satisfied: datasets>=1.2.1 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (2.9.0)\n",
            "Requirement already satisfied: evaluate in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (0.4.0)\n",
            "Requirement already satisfied: onnxruntime>=1.9.0 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (1.13.1)\n",
            "Requirement already satisfied: protobuf==3.20.1 in /usr/local/lib/python3.8/dist-packages (from optimum[onnxruntime]) (3.20.1)\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",
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (0.70.14)\n",
            "Requirement 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: 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: pandas in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (1.3.5)\n",
            "Requirement already satisfied: responses<0.19 in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (0.18.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",
            "Requirement 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",
            "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: xxhash in /usr/local/lib/python3.8/dist-packages (from datasets>=1.2.1->optimum[onnxruntime]) (3.2.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: filelock in /usr/local/lib/python3.8/dist-packages (from huggingface-hub>=0.8.0->optimum[onnxruntime]) (3.9.0)\n",
            "Requirement 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: flatbuffers in /usr/local/lib/python3.8/dist-packages (from onnxruntime>=1.9.0->optimum[onnxruntime]) (1.12)\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: 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: 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",
            "Requirement already satisfied: humanfriendly>=9.1 in /usr/local/lib/python3.8/dist-packages (from coloredlogs->optimum[onnxruntime]) (10.0)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.8/dist-packages (from sympy->optimum[onnxruntime]) (1.2.1)\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: 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: 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: 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: 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: 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: 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: 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",
            "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.26.14)\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: 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: 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"
          ]
        }
      ]
    },
    {
      "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",
        "from optimum.pipelines import pipeline\n",
        "from tqdm import tqdm\n",
        "from ast import literal_eval\n",
        "from zipfile import ZipFile\n",
        "\n",
        "nltk.download('punkt')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lPHDMY7Xkpfh",
        "outputId": "66a1473f-c50b-4e90-bf98-c2c7f662c306"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
            "[nltk_data]   Package punkt is already up-to-date!\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "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.download(\n",
        "        \"https://drive.google.com/file/d/137tZvp-iTMR2xIogasglSH4jTTLW4_Sf/view\",\n",
        "        fuzzy=True, use_cookies=False, quiet=True\n",
        "    )\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_queries(ques_embed, theme):\n",
        "    \"\"\"Retrieve nearest already answered queries to the questions\"\"\"\n",
        "    index = load_index(f'/content/indices/{theme}_ques_l2_index')\n",
        "    return get_k_nearest_neighbours(index, ques_embed, 3)\n",
        "\n",
        "def get_nearest_sentences(ques_embed, theme):\n",
        "    \"\"\"Retrieve nearest sentences to the questions\"\"\"\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": [
        "def search_previously_answered_queries(q_id, dist, query_idx, solved_queries):\n",
        "    \"\"\"Search previously answered queries and return its answer if it exists\"\"\"\n",
        "    if dist > query_threshold:\n",
        "        return False, None\n",
        "    ans = {\n",
        "        \"question_id\": q_id,\n",
        "        \"answers\": solved_queries['answers'][query_idx][0] if len(solved_queries['answers'][query_idx]) > 0 else '',\n",
        "        \"paragraph_id\": solved_queries['paragraph_id'][query_idx]\n",
        "    }\n",
        "    return True, ans"
      ]
    },
    {
      "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",
        "    )\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)"
      ],
      "metadata": {
        "id": "YWnuk1j6qPuV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Parameters for context generation\n",
        "k = 8\n",
        "query_threshold = 0.3\n",
        "distance_threshold = 1.8\n",
        "context_length_threshold = 205"
      ],
      "metadata": {
        "id": "W0RjrWvFoIRq"
      },
      "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",
        "    solved_queries = solved_ques[theme]\n",
        "    print(f'Theme: {theme}')\n",
        "\n",
        "    # Nearest Neighbour Search\n",
        "    start_time = time.time()\n",
        "    ques_list = [q['question'] for q in questions]\n",
        "    ques_embed = get_embeddings(ques_list, sentence_encoder)\n",
        "    D_ques, I_ques = get_nearest_queries(ques_embed, theme)\n",
        "    D_sents, I_sents = get_nearest_sentences(ques_embed, 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",
        "        # Check previously answered queries\n",
        "        found, ans = search_previously_answered_queries(\n",
        "            q[\"id\"], D_ques[i][0], I_ques[i][0], solved_queries\n",
        "        )\n",
        "        if found:\n",
        "            pred_out.append(ans)\n",
        "            continue\n",
        "        # Context Generation\n",
        "        context, context_para_ids, sent_length = get_context(\n",
        "            paras[theme]['sentences'], paras[theme]['id'], I_sents[i], D_sents[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/",
          "height": 413
        },
        "outputId": "60babd15-978b-4cec-c03a-8eadb972805a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Theme: IPod\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "IndexError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-39-24e68e6d3e5c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;31m# Load model fine-tuned for this theme.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mtheme_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_theme_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m     \u001b[0mexecution_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtimeit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mpred_theme_ans\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheme_ques\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpred_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnumber\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m     \u001b[0mtheme_inf_time\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtheme_interval\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"theme\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mexecution_time\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;36m1000\u001b[0m \u001b[0;31m# in milliseconds.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0mpred_df\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_records\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.8/timeit.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(stmt, setup, timer, number, globals)\u001b[0m\n\u001b[1;32m    231\u001b[0m            number=default_number, globals=None):\n\u001b[1;32m    232\u001b[0m     \u001b[0;34m\"\"\"Convenience function to create Timer object and call timeit method.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 233\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mTimer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstmt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msetup\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mglobals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumber\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    234\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    235\u001b[0m def repeat(stmt=\"pass\", setup=\"pass\", timer=default_timer,\n",
            "\u001b[0;32m/usr/lib/python3.8/timeit.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m    175\u001b[0m         \u001b[0mgc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdisable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    176\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 177\u001b[0;31m             \u001b[0mtiming\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    178\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    179\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mgcold\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.8/timeit.py\u001b[0m in \u001b[0;36minner\u001b[0;34m(_it, _timer, _stmt)\u001b[0m\n",
            "\u001b[0;32m<ipython-input-39-24e68e6d3e5c>\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;31m# Load model fine-tuned for this theme.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mtheme_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_theme_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m     \u001b[0mexecution_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtimeit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mpred_theme_ans\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheme_ques\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpred_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnumber\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m     \u001b[0mtheme_inf_time\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtheme_interval\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"theme\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mexecution_time\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;36m1000\u001b[0m \u001b[0;31m# in milliseconds.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0mpred_df\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_records\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-38-364150689918>\u001b[0m in \u001b[0;36mpred_theme_ans\u001b[0;34m(questions, theme_model, pred_out)\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0mques_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'question'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mq\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mquestions\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0mques_embed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_embeddings\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mques_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msentence_encoder\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0mD_sents\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mI_sents\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_nearest_sentences\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mques_embed\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m     \u001b[0mD_sents\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mI_sents\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_nearest_sentences\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mques_embed\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mann_inference_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mstart_time\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;36m1000.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-7-acbcc489bad7>\u001b[0m in \u001b[0;36mget_nearest_sentences\u001b[0;34m(questions, theme)\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_nearest_sentences\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquestions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;34m\"\"\"Retrieve nearest sentences to the questions\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m     \u001b[0mques_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'question'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mq\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mquestions\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m     \u001b[0mques_embed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_embeddings\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mques_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msentence_encoder\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'/content/indices/{theme}_para_l2_index'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-7-acbcc489bad7>\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_nearest_sentences\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquestions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtheme\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;34m\"\"\"Retrieve nearest sentences to the questions\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m     \u001b[0mques_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'question'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mq\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mquestions\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     19\u001b[0m     \u001b[0mques_embed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_embeddings\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mques_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msentence_encoder\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'/content/indices/{theme}_para_l2_index'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mIndexError\u001b[0m: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices"
          ]
        }
      ]
    },
    {
      "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. * 1. * 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": "eb56cb3f-9540-4eb4-9d20-dea73b02b53f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.8898943749043251\n",
            "0.8689999198401515\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 0.8872037021772665\n",
        "# 0.8519277289714773\n",
        "\n",
        "# 0.8898943749043251\n",
        "# 0.8689999198401515"
      ],
      "metadata": {
        "id": "CI7eZjF8JCSa"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "M9jJuSCs26SW",
        "UzHw72nd6_wA",
        "BTEq0Khe7DM0",
        "JyKtVXm661x3",
        "Uxnd1xuWqVeA",
        "raCs2XgbZpTG"
      ],
      "provenance": [],
      "authorship_tag": "ABX9TyPV/Cs4M9QKPZMCHxiP07QN",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}