M003_retrieve_system_BiEncoder_finetune_model_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# About\n",
    "\n",
    "This notebook builds an IR retrieve system. It finetues a Biencoder model using a pretrianed BERT model and Train and Evaluation sets. The evaluator used for finetuning is EmbeddingSimilarityEvaluator developed by the sentence transformer package. The retrieve system is evaluated by the top 10 ranked passages for each query in the Test set.\n",
    "\n",
    "reference: https://www.sbert.net/examples/applications/retrieve_rerank/README.html\n",
    "https://www.sbert.net/docs/usage/semantic_textual_similarity.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import logging\n",
    "from sentence_transformers import LoggingHandler, util, SentenceTransformer\n",
    "from datetime import datetime\n",
    "#from sentence_transformers.cross_encoder import CrossEncoder\n",
    "from torch.utils.data import DataLoader\n",
    "#from sentence_transformers.cross_encoder.evaluation import CERerankingEvaluator, CEBinaryClassificationEvaluator\n",
    "from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator\n",
    "# import gzip\n",
    "# import os\n",
    "# import tarfile\n",
    "from sentence_transformers import InputExample, losses\n",
    "from datetime import datetime\n",
    "# import csv\n",
    "# import pickle\n",
    "# import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Just some code to print debug information to stdout\n",
    "logging.basicConfig(\n",
    "    format=\"%(asctime)s - %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\", level=logging.INFO, handlers=[LoggingHandler()]\n",
    ")\n",
    "#### /print debug information to stdout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configure the Training Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_batch_size = 30\n",
    "num_epochs = 1\n",
    "warmup_steps = 5000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_to_list(query_df, passage_df):\n",
    "    \"\"\"\n",
    "    # this function create train, test data\n",
    "    # it reads in query_df and passage_df, then for loop to create a list of list\n",
    "    # each query, passage pair is a list of 4 elements: [query, passage, score, rel]\n",
    "    # the ouput is [[query, passage, score, rel], [query, passage, score, rel], ...]\n",
    "    \"\"\"\n",
    "    list_query_passage_score_rel = []\n",
    "    for index, row in query_df.iterrows():\n",
    "        list = []\n",
    "        list.append(row['query'])\n",
    "        list.append(passage_df[passage_df['passage_id'] == row['passage_id']][\"passage\"].iloc[0])\n",
    "        list.append(row[\"score\"])\n",
    "        list.append(row[\"rel\"])\n",
    "        list_query_passage_score_rel.append(list)\n",
    "    return list_query_passage_score_rel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read in Train, Validation, Test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# query\n",
    "query_train = pd.read_csv(\"./output/query_train_set_with_passage_info.csv\")\n",
    "query_val =pd.read_csv(\"./output/query_val_set_with_passage_info.csv\")\n",
    "query_test =pd.read_csv(\"./output/query_test_set_with_passage_info.csv\")\n",
    "\n",
    "# passage\n",
    "passage_train = pd.read_csv(\"./output/train_passage_id_content_cleaned.csv\")\n",
    "passage_val = pd.read_csv(\"./output/val_passage_id_content_cleaned.csv\")\n",
    "passage_test = pd.read_csv(\"./output/test_passage_id_content_cleaned.csv\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(20000, 9)\n",
      "(20000, 9)\n",
      "(20000, 9)\n"
     ]
    }
   ],
   "source": [
    "print(query_train.shape)\n",
    "print(query_val.shape)\n",
    "print(query_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>query</th>\n",
       "      <th>used</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>username</th>\n",
       "      <th>rel</th>\n",
       "      <th>query_cleaned</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>568182</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_62_547459701</td>\n",
       "      <td>1</td>\n",
       "      <td>16.102699</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>568182</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_21_588232716</td>\n",
       "      <td>2</td>\n",
       "      <td>14.784700</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>568182</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_11_97323294</td>\n",
       "      <td>3</td>\n",
       "      <td>13.865800</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>568182</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_68_54887603</td>\n",
       "      <td>4</td>\n",
       "      <td>13.865799</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>568182</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_62_547466749</td>\n",
       "      <td>5</td>\n",
       "      <td>13.793600</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what are the characteristics of wool fibres</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   query_id                                        query used  \\\n",
       "0    568182  what are the characteristics of wool fibres   Q0   \n",
       "1    568182  what are the characteristics of wool fibres   Q0   \n",
       "2    568182  what are the characteristics of wool fibres   Q0   \n",
       "3    568182  what are the characteristics of wool fibres   Q0   \n",
       "4    568182  what are the characteristics of wool fibres   Q0   \n",
       "\n",
       "                     passage_id  rank      score  username  rel  \\\n",
       "0  msmarco_passage_62_547459701     1  16.102699  Anserini    1   \n",
       "1  msmarco_passage_21_588232716     2  14.784700  Anserini    1   \n",
       "2   msmarco_passage_11_97323294     3  13.865800  Anserini    1   \n",
       "3   msmarco_passage_68_54887603     4  13.865799  Anserini    1   \n",
       "4  msmarco_passage_62_547466749     5  13.793600  Anserini    1   \n",
       "\n",
       "                                 query_cleaned  \n",
       "0  what are the characteristics of wool fibres  \n",
       "1  what are the characteristics of wool fibres  \n",
       "2  what are the characteristics of wool fibres  \n",
       "3  what are the characteristics of wool fibres  \n",
       "4  what are the characteristics of wool fibres  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>query</th>\n",
       "      <th>used</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>username</th>\n",
       "      <th>rel</th>\n",
       "      <th>query_cleaned</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>791126</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_48_489095876</td>\n",
       "      <td>1</td>\n",
       "      <td>9.5709</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>791126</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_30_653845189</td>\n",
       "      <td>2</td>\n",
       "      <td>9.5522</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>791126</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_04_540229875</td>\n",
       "      <td>3</td>\n",
       "      <td>9.4101</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>791126</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_27_425453773</td>\n",
       "      <td>4</td>\n",
       "      <td>9.1409</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>791126</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_59_586798337</td>\n",
       "      <td>5</td>\n",
       "      <td>8.8821</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>what is respiration controlled by</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   query_id                              query used  \\\n",
       "0    791126  what is respiration controlled by   Q0   \n",
       "1    791126  what is respiration controlled by   Q0   \n",
       "2    791126  what is respiration controlled by   Q0   \n",
       "3    791126  what is respiration controlled by   Q0   \n",
       "4    791126  what is respiration controlled by   Q0   \n",
       "\n",
       "                     passage_id  rank   score  username  rel  \\\n",
       "0  msmarco_passage_48_489095876     1  9.5709  Anserini    1   \n",
       "1  msmarco_passage_30_653845189     2  9.5522  Anserini    1   \n",
       "2  msmarco_passage_04_540229875     3  9.4101  Anserini    1   \n",
       "3  msmarco_passage_27_425453773     4  9.1409  Anserini    1   \n",
       "4  msmarco_passage_59_586798337     5  8.8821  Anserini    1   \n",
       "\n",
       "                       query_cleaned  \n",
       "0  what is respiration controlled by  \n",
       "1  what is respiration controlled by  \n",
       "2  what is respiration controlled by  \n",
       "3  what is respiration controlled by  \n",
       "4  what is respiration controlled by  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_val.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>query</th>\n",
       "      <th>used</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>username</th>\n",
       "      <th>rel</th>\n",
       "      <th>query_cleaned</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_49_115778700</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2671</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_26_450017756</td>\n",
       "      <td>2</td>\n",
       "      <td>11.8537</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_19_202396335</td>\n",
       "      <td>3</td>\n",
       "      <td>11.2738</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_28_623516004</td>\n",
       "      <td>4</td>\n",
       "      <td>11.0361</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_02_47181174</td>\n",
       "      <td>5</td>\n",
       "      <td>10.9476</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   query_id                                     query used  \\\n",
       "0    203324  him functions to the paper health record   Q0   \n",
       "1    203324  him functions to the paper health record   Q0   \n",
       "2    203324  him functions to the paper health record   Q0   \n",
       "3    203324  him functions to the paper health record   Q0   \n",
       "4    203324  him functions to the paper health record   Q0   \n",
       "\n",
       "                     passage_id  rank    score  username  rel  \\\n",
       "0  msmarco_passage_49_115778700     1  13.2671  Anserini    1   \n",
       "1  msmarco_passage_26_450017756     2  11.8537  Anserini    1   \n",
       "2  msmarco_passage_19_202396335     3  11.2738  Anserini    1   \n",
       "3  msmarco_passage_28_623516004     4  11.0361  Anserini    1   \n",
       "4   msmarco_passage_02_47181174     5  10.9476  Anserini    1   \n",
       "\n",
       "                              query_cleaned  \n",
       "0  him functions to the paper health record  \n",
       "1  him functions to the paper health record  \n",
       "2  him functions to the paper health record  \n",
       "3  him functions to the paper health record  \n",
       "4  him functions to the paper health record  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>passage_id</th>\n",
       "      <th>passage</th>\n",
       "      <th>passage_cleaned</th>\n",
       "      <th>passage_cleaned_lemmatized</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>msmarco_passage_62_547459701</td>\n",
       "      <td>Table of Contents. Growth. Harvesting. Grading...</td>\n",
       "      <td>table of contents growth harvesting grading of...</td>\n",
       "      <td>table content growth harvesting grade wool fib...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>msmarco_passage_21_588232716</td>\n",
       "      <td>A micron ( micrometre) is the measurement used...</td>\n",
       "      <td>a micron micrometre is the measurement used to...</td>\n",
       "      <td>micron micrometre measurement express diameter...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>msmarco_passage_11_97323294</td>\n",
       "      <td>Objective measurements include diameter (micro...</td>\n",
       "      <td>objective measurements include diameter micron...</td>\n",
       "      <td>objective measurement include diameter micron ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>msmarco_passage_68_54887603</td>\n",
       "      <td>Objective measurements include diameter (micro...</td>\n",
       "      <td>objective measurements include diameter micron...</td>\n",
       "      <td>objective measurement include diameter micron ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>msmarco_passage_62_547466749</td>\n",
       "      <td>Summary of Characteristics of Wool Fibers. Woo...</td>\n",
       "      <td>summary of characteristics of wool fibers wool...</td>\n",
       "      <td>summary characteristic wool fiber wool protein...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                     passage_id  \\\n",
       "0  msmarco_passage_62_547459701   \n",
       "1  msmarco_passage_21_588232716   \n",
       "2   msmarco_passage_11_97323294   \n",
       "3   msmarco_passage_68_54887603   \n",
       "4  msmarco_passage_62_547466749   \n",
       "\n",
       "                                             passage  \\\n",
       "0  Table of Contents. Growth. Harvesting. Grading...   \n",
       "1  A micron ( micrometre) is the measurement used...   \n",
       "2  Objective measurements include diameter (micro...   \n",
       "3  Objective measurements include diameter (micro...   \n",
       "4  Summary of Characteristics of Wool Fibers. Woo...   \n",
       "\n",
       "                                     passage_cleaned  \\\n",
       "0  table of contents growth harvesting grading of...   \n",
       "1  a micron micrometre is the measurement used to...   \n",
       "2  objective measurements include diameter micron...   \n",
       "3  objective measurements include diameter micron...   \n",
       "4  summary of characteristics of wool fibers wool...   \n",
       "\n",
       "                          passage_cleaned_lemmatized  \n",
       "0  table content growth harvesting grade wool fib...  \n",
       "1  micron micrometre measurement express diameter...  \n",
       "2  objective measurement include diameter micron ...  \n",
       "3  objective measurement include diameter micron ...  \n",
       "4  summary characteristic wool fiber wool protein...  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare model input for Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[['what are the characteristics of wool fibres',\n",
       "  'Table of Contents. Growth. Harvesting. Grading of Wool Fibers. Properties of Wool Fibers. Application of Wool Fibers. Characteristics of Wool Fibers and Products. Summary of Characteristics of Wool Fibers. Of the major apparel fibres, wool is the most reusable and recyclable fibre on the planet.',\n",
       "  16.102699,\n",
       "  1],\n",
       " ['what are the characteristics of wool fibres',\n",
       "  'A micron ( micrometre) is the measurement used to express the diameter of wool fibre. Fine wool fibers have a low micron value. Fibre diameter is the most important characteristic of wool in determining its value.',\n",
       "  14.7847,\n",
       "  1],\n",
       " ['what are the characteristics of wool fibres',\n",
       "  'Objective measurements include diameter (micron), length, strength, position of break, vegetable matter and colour. AWEX-ID covers subjective characteristics. Diameter. Mean fibre diameter is a measurement in micrometres (microns) of the average diameter of wool fibres in a sale lot. Fibre diameter is responsible for 70-80 per cent of the greasy wool price over the long term.',\n",
       "  13.8658,\n",
       "  1]]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_query_passage_score_rel_train = df_to_list(query_train, passage_train)\n",
    "print(len(list_query_passage_score_rel_train))\n",
    "list_query_passage_score_rel_train[0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare model input for Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[['what is respiration controlled by',\n",
       "  'Explain what respirable crystalline silica is and the health hazards associated with it. 3. Know the specific tasks in the workplace that could result in exposure to respirable crystalline silica. 4. Know the types of control measures and respiratory protection for respirable crystalline silica. 5.',\n",
       "  9.5709,\n",
       "  1],\n",
       " ['what is respiration controlled by',\n",
       "  'soft and elastic. which part of the brain controls the involuntary action of respiration. medulla oblongata. when assessing respirations, what should be noted. abnormalitis in characteristics such as rate, rythm, and depth. what is the normal respiratory rate. 12-20.',\n",
       "  9.5522,\n",
       "  1],\n",
       " ['what is respiration controlled by',\n",
       "  'What Part of the Brain Controls Breathing. Let’s get to know What is Breathing and Respiration, before we get into ‘what part of the brain controls breathing’. We say we are ALIVE.',\n",
       "  9.4101,\n",
       "  1]]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_eval = df_to_list(query_val, passage_val)\n",
    "print(len(list_eval))\n",
    "list_eval[0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load in pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output/bi-encoder-finetuneall-MiniLM-L6-v2-2024-04-22_16-23-41\n"
     ]
    }
   ],
   "source": [
    "\n",
    "model_name = \"all-MiniLM-L6-v2\"\n",
    "model = SentenceTransformer(model_name)\n",
    "num_candidates = 20\n",
    "\n",
    "model_save_path = (\n",
    "    \"output/bi-encoder-finetune\"\n",
    "    + model_name.replace(\"/\", \"-\")\n",
    "    + \"-\"\n",
    "    + datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\n",
    ")\n",
    "print(model_save_path)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine tune the model using Train and Validation sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "667\n"
     ]
    }
   ],
   "source": [
    "train_sample = []\n",
    "for list in list_query_passage_score_rel_train:\n",
    "    # use the score as Target, as I am to use cosine similarity for evaluator\n",
    "    train_sample.append(InputExample(texts=[list[0], list[1]], label = list[2]))\n",
    "train_dataloader = DataLoader(train_sample, shuffle=True, batch_size=train_batch_size)\n",
    "print(len(train_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_sample = []\n",
    "for list in list_eval:\n",
    "    # use the score as Target, as I am to use cosine similarity for evaluator\n",
    "    val_sample.append(InputExample(texts=[list[0], list[1]], label = list[2])) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<sentence_transformers.evaluation.EmbeddingSimilarityEvaluator.EmbeddingSimilarityEvaluator at 0x7fc398c46810>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluator = EmbeddingSimilarityEvaluator.from_input_examples(val_sample, name=\"train-eval-binary\")\n",
    "evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-04-21 20:45:01 - Warmup-steps: 5000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b02c2912ecd54f01aa0e427d449fab5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d7a66113b9f426b884b39907618eafc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Iteration:   0%|          | 0/667 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-04-22 12:14:17 - EmbeddingSimilarityEvaluator: Evaluating the model on train-eval-binary dataset after epoch 0:\n",
      "2024-04-22 12:16:04 - Cosine-Similarity :\tPearson: 0.3018\tSpearman: 0.3901\n",
      "2024-04-22 12:16:04 - Manhattan-Distance:\tPearson: 0.3390\tSpearman: 0.3881\n",
      "2024-04-22 12:16:04 - Euclidean-Distance:\tPearson: 0.3406\tSpearman: 0.3901\n",
      "2024-04-22 12:16:04 - Dot-Product-Similarity:\tPearson: 0.3018\tSpearman: 0.3901\n",
      "2024-04-22 12:16:04 - Save model to output/bi-encoder-finetuneall-MiniLM-L6-v2-2024-04-21_20-27-45\n"
     ]
    }
   ],
   "source": [
    "# reference: https://sbert.net/docs/training/overview.html#training-data\n",
    "\n",
    "\n",
    "# Configure the training\n",
    "\n",
    "logging.info(\"Warmup-steps: {}\".format(warmup_steps))\n",
    "train_loss = losses.CosineSimilarityLoss(model)\n",
    "\n",
    "# Train the model\n",
    "model.fit(\n",
    "    train_objectives=[(train_dataloader, train_loss)],\n",
    "    evaluator=evaluator,\n",
    "    epochs=num_epochs,\n",
    "    evaluation_steps=10000,\n",
    "    warmup_steps=warmup_steps,\n",
    "    output_path=model_save_path,\n",
    "    use_amp=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [skip] Print evaluation result using Sbert function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model and eval on test set\n",
    "model_save_path = \"output/bi-encoder-finetuneall-MiniLM-L6-v2-2024-04-21_20-27-45\"\n",
    "\n",
    "model = SentenceTransformer(model_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20000\n",
      "20000\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<sentence_transformers.readers.InputExample.InputExample at 0x7f43c878f210>,\n",
       " <sentence_transformers.readers.InputExample.InputExample at 0x7f43c87c2c10>,\n",
       " <sentence_transformers.readers.InputExample.InputExample at 0x7f43ca0bd650>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_query_passage_score_rel_test = df_to_list(query_test, passage_test)\n",
    "print(len(list_query_passage_score_rel_test))\n",
    "\n",
    "test_sample = []\n",
    "for list in list_query_passage_score_rel_test:\n",
    "    test_sample.append(InputExample(texts=[list[0], list[1]], label = list[3]))\n",
    "\n",
    "print(len(test_sample))\n",
    "test_sample[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2858092725749748"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "\n",
    "evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_sample, name=\"sts-test-binary\")\n",
    "evaluator(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare input for Test Set's Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_data_for_inference(query_df, passage_df):\n",
    "    \"\"\"\n",
    "    # this function for loop each query, create list_query_passage (a list of list) for every query: [[query1, passage1], [query1, passage2] ....]\n",
    "    # then pass this query's list_query_passage to model to make prediction, each query will return a score (list): [score1, score2, score3 ....]\n",
    "    # the ouput dict_qid_all_prediction is a dictionary that append all query's list_score: {qid1: [score1, score2, score3 ....]}\n",
    "    # the ouput dict_qid_pid_true is a dictionary that append all query's True ranked qassage {qid1: [pid1, pid2, ..., pid of topn]}\n",
    "    \"\"\"\n",
    "\n",
    "    # create a list of list for all queries [[query1, passage1], [query1, passage2]..., [query2, passage1], [query2, passage2]..., ]\n",
    "    # This would be used to fit into model to make prediction\n",
    "    list_query_passage = []\n",
    "\n",
    "\n",
    "    # create a list of list for all queries [[qid1, pid1], [qid1, pid2]..., [qid2, pid1], [qid2, pid2]..., ]\n",
    "    # This would be used to collect and sort model prediction\n",
    "    list_qid_all_pid = []\n",
    "\n",
    "    # create a dict of dict for all queries {qid1: {pid1: rel, pid2: rel, ...}, qid2: {pid1: rel, pid2: rel, ...}, }\n",
    "    # This would be used to evaluate model prediction\n",
    "    dict_of_dict_qid_pid_rel = {}\n",
    "\n",
    "    # iterate over all qid\n",
    "    list_qid = query_df[\"query_id\"].unique()\n",
    "    for qid in list_qid:\n",
    "\n",
    "        \n",
    "        # get the rows that contains the qid, pid, score, rel for this qid\n",
    "        query_df_subset = query_df[query_df['query_id'] == qid]\n",
    "        \n",
    "        # iterate over all pid\n",
    "        for index, row in query_df_subset.iterrows():\n",
    "            \n",
    "            # create a list of 2 elements [query, passage]\n",
    "            list = []\n",
    "            list.append(row['query'])\n",
    "\n",
    "            # get the only 1 row that contains the passage text in passage_df and whose pid matches the row's pid\n",
    "            list.append(passage_df[passage_df['passage_id'] == row['passage_id']][\"passage\"].iloc[0]) # get passage value\n",
    "\n",
    "            # append all lists of 2 elements of all queries into one list\n",
    "            list_query_passage.append(list)\n",
    "\n",
    "            # append the sublist. Each sublist is the [qid, pid] of the query\n",
    "            list_qid_all_pid.append([row['query_id'], passage_df[passage_df['passage_id'] == row['passage_id']]['passage_id'].iloc[0]])\n",
    "\n",
    "        \n",
    "        \n",
    "        # for this query, create a dictionary of Y true - each key is a pid, each value is the relavence (1 or 0) of the passage {pid1: 1, pid2, 0, ....}\n",
    "        dict_pid_rel = pd.Series(query_df_subset['rel'].values, index = query_df_subset['passage_id']).to_dict()\n",
    "\n",
    "        # append the dict of this query to the dict for all queries \n",
    "        dict_of_dict_qid_pid_rel[qid] = dict_pid_rel\n",
    "\n",
    "    return list_query_passage, list_qid_all_pid, dict_of_dict_qid_pid_rel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_query_passage, list_qid_all_pid, dict_of_dict_qid_pid_rel = create_data_for_inference(query_test, passage_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20000\n",
      "['him functions to the paper health record', 'This information can be either paper-based, a combination of paper and digital (hybrid) or as is more often the case, a fully electronic health record (EHR). HIM professionals are using advanced technology as they perform most HIM-related functions.']\n",
      "['what does the word element cyte mean', 'What does -cyte mean? The combining form -cyte is used like a suffix meaning “cell.”. It is often used in medical and scientific terms, especially to name different types of cells. The form -cyte comes from the Greek kýtos, meaning “container,” “receptacle,” “body.”.']\n",
      "20000\n",
      "2\n",
      "[203324, 'msmarco_passage_49_115778700']\n",
      "1000\n",
      "20\n",
      "{'msmarco_passage_49_115778700': 1, 'msmarco_passage_26_450017756': 1, 'msmarco_passage_19_202396335': 1, 'msmarco_passage_28_623516004': 1, 'msmarco_passage_02_47181174': 1, 'msmarco_passage_02_47511835': 1, 'msmarco_passage_11_636220699': 1, 'msmarco_passage_41_133447691': 1, 'msmarco_passage_58_723159152': 1, 'msmarco_passage_24_583711786': 1, 'msmarco_passage_18_714845189': 0, 'msmarco_passage_59_733040871': 0, 'msmarco_passage_65_238890762': 0, 'msmarco_passage_17_178380607': 0, 'msmarco_passage_61_103079366': 0, 'msmarco_passage_07_257313185': 0, 'msmarco_passage_53_560231782': 0, 'msmarco_passage_16_326265100': 0, 'msmarco_passage_16_326272763': 0, 'msmarco_passage_02_835063630': 0}\n"
     ]
    }
   ],
   "source": [
    "# check\n",
    "print(len(list_query_passage))\n",
    "print(list_query_passage[0])\n",
    "print(list_query_passage[1000])\n",
    "\n",
    "print(len(list_qid_all_pid))\n",
    "print(len(list_qid_all_pid[0]))\n",
    "print(list_qid_all_pid[0])\n",
    "\n",
    "print(len(dict_of_dict_qid_pid_rel))\n",
    "print(len(dict_of_dict_qid_pid_rel[203324]))\n",
    "print(dict_of_dict_qid_pid_rel[203324])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict on Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_score = []\n",
    "for query, passage in list_query_passage:\n",
    "    query_embedding = model.encode(query, convert_to_tensor=True, show_progress_bar=False)\n",
    "    passage_embedding = model.encode(passage, convert_to_tensor=True, show_progress_bar=False)\n",
    "    cosine_score = util.cos_sim(query_embedding, passage_embedding)\n",
    "    list_score.append(cosine_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20000\n",
      "0.9623302817344666\n"
     ]
    }
   ],
   "source": [
    "# take value out of tensor\n",
    "list_score_ = []\n",
    "for i in range(len(list_score)):\n",
    "    list_score_.append(list_score[i][0][0].item())\n",
    "\n",
    "print(len(list_score_))\n",
    "print(list_score_[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(20000, 3)\n"
     ]
    }
   ],
   "source": [
    "test_with_pred = pd.DataFrame(list_qid_all_pid, columns = ['query_id', 'passage_id'])\n",
    "test_with_pred['predicted_score'] = list_score_\n",
    "test_with_pred.sort_values(by=['query_id', 'predicted_score'], ascending = False, inplace = True)\n",
    "test_with_pred.to_csv(\"./output/test_prediction_biencoder_refit_all-MiniLM-L6-v2_2024-04-22.csv\"\n",
    "                                       , index = False)\n",
    "print(test_with_pred.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>predicted_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>10378</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_37_349325721</td>\n",
       "      <td>0.969984</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10363</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_25_322137693</td>\n",
       "      <td>0.961158</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10364</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_25_322151896</td>\n",
       "      <td>0.961158</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10365</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_25_322155605</td>\n",
       "      <td>0.961158</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10368</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_27_266644928</td>\n",
       "      <td>0.960574</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10377</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_36_754635695</td>\n",
       "      <td>0.960268</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10375</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_48_763784628</td>\n",
       "      <td>0.958300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10379</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_15_272642271</td>\n",
       "      <td>0.957656</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10374</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_48_721956358</td>\n",
       "      <td>0.957487</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10370</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_39_120773224</td>\n",
       "      <td>0.954245</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10361</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_01_536296949</td>\n",
       "      <td>0.953881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10362</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_06_770443306</td>\n",
       "      <td>0.953881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10366</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_48_764103270</td>\n",
       "      <td>0.953881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10367</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_48_800555103</td>\n",
       "      <td>0.953881</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10372</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_03_2856732</td>\n",
       "      <td>0.952447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10369</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_48_449276269</td>\n",
       "      <td>0.951552</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10360</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_64_307501210</td>\n",
       "      <td>0.950863</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10371</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_01_493126868</td>\n",
       "      <td>0.949798</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10373</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_09_689048460</td>\n",
       "      <td>0.949683</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10376</th>\n",
       "      <td>1183224</td>\n",
       "      <td>msmarco_passage_61_527867091</td>\n",
       "      <td>0.949683</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       query_id                    passage_id  predicted_score\n",
       "10378   1183224  msmarco_passage_37_349325721         0.969984\n",
       "10363   1183224  msmarco_passage_25_322137693         0.961158\n",
       "10364   1183224  msmarco_passage_25_322151896         0.961158\n",
       "10365   1183224  msmarco_passage_25_322155605         0.961158\n",
       "10368   1183224  msmarco_passage_27_266644928         0.960574\n",
       "10377   1183224  msmarco_passage_36_754635695         0.960268\n",
       "10375   1183224  msmarco_passage_48_763784628         0.958300\n",
       "10379   1183224  msmarco_passage_15_272642271         0.957656\n",
       "10374   1183224  msmarco_passage_48_721956358         0.957487\n",
       "10370   1183224  msmarco_passage_39_120773224         0.954245\n",
       "10361   1183224  msmarco_passage_01_536296949         0.953881\n",
       "10362   1183224  msmarco_passage_06_770443306         0.953881\n",
       "10366   1183224  msmarco_passage_48_764103270         0.953881\n",
       "10367   1183224  msmarco_passage_48_800555103         0.953881\n",
       "10372   1183224    msmarco_passage_03_2856732         0.952447\n",
       "10369   1183224  msmarco_passage_48_449276269         0.951552\n",
       "10360   1183224  msmarco_passage_64_307501210         0.950863\n",
       "10371   1183224  msmarco_passage_01_493126868         0.949798\n",
       "10373   1183224  msmarco_passage_09_689048460         0.949683\n",
       "10376   1183224  msmarco_passage_61_527867091         0.949683"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# check\n",
    "display(test_with_pred[test_with_pred[\"query_id\"] == 1183224])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot distribution of predicted score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([1.000e+00, 0.000e+00, 0.000e+00, 0.000e+00, 1.000e+00, 2.000e+00,\n",
       "        1.000e+00, 2.000e+00, 1.000e+00, 1.000e+00, 3.000e+00, 2.000e+00,\n",
       "        5.000e+00, 5.000e+00, 4.000e+00, 3.000e+00, 2.000e+00, 2.000e+00,\n",
       "        7.000e+00, 4.000e+00, 1.100e+01, 4.000e+00, 6.000e+00, 1.400e+01,\n",
       "        1.400e+01, 1.200e+01, 2.100e+01, 1.800e+01, 3.200e+01, 3.600e+01,\n",
       "        5.600e+01, 6.900e+01, 9.000e+01, 1.060e+02, 1.430e+02, 1.900e+02,\n",
       "        2.200e+02, 3.030e+02, 3.050e+02, 4.640e+02, 6.120e+02, 8.660e+02,\n",
       "        1.102e+03, 1.482e+03, 1.883e+03, 2.470e+03, 3.025e+03, 3.275e+03,\n",
       "        2.537e+03, 5.880e+02]),\n",
       " array([0.61433995, 0.62196314, 0.62958633, 0.63720952, 0.64483272,\n",
       "        0.65245591, 0.6600791 , 0.66770229, 0.67532548, 0.68294868,\n",
       "        0.69057187, 0.69819506, 0.70581825, 0.71344144, 0.72106464,\n",
       "        0.72868783, 0.73631102, 0.74393421, 0.7515574 , 0.7591806 ,\n",
       "        0.76680379, 0.77442698, 0.78205017, 0.78967337, 0.79729656,\n",
       "        0.80491975, 0.81254294, 0.82016613, 0.82778933, 0.83541252,\n",
       "        0.84303571, 0.8506589 , 0.85828209, 0.86590529, 0.87352848,\n",
       "        0.88115167, 0.88877486, 0.89639805, 0.90402125, 0.91164444,\n",
       "        0.91926763, 0.92689082, 0.93451401, 0.94213721, 0.9497604 ,\n",
       "        0.95738359, 0.96500678, 0.97262998, 0.98025317, 0.98787636,\n",
       "        0.99549955]),\n",
       " <BarContainer object of 50 artists>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjEAAAGdCAYAAADjWSL8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAs2ElEQVR4nO3dfXSU9Z3//9eQmyHQMBLSTJISKG0JQYNsFzQJ1oIFE9iN0dIWuyz5Qg8irgpmgfWInK7pHiEuPQu0obCURaIEjF1drJ5KNJwqSsONpvCrkAio2EnW3DhumIBMJ4jX748eruOQG5iYm/kMz8c51znMdb3nyufNR+HFNddnLodlWZYAAAAMM2igBwAAANAThBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJGiB3oAfeXzzz/XRx99pPj4eDkcjoEeDgAAuAqWZens2bNKTU3VoEHdX2uJ2BDz0UcfKS0tbaCHAQAAeqC+vl4jR47stiZiQ0x8fLykv/4mDBs2bIBHAwAArkZbW5vS0tLsv8e7E7Eh5tJHSMOGDSPEAABgmKu5FYQbewEAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYKWKfYg0AwLXI4/HI6/V2W5OYmKhRo0b104j6DiEGAIAI4fF4lDE+Q/7z/m7r4obE6d26d40PMoQYAAAihNfrlf+8X/O2zJM73d1pTfPJZpUvLpfX6yXEAACA8OJOdyttYtpAD6PPcWMvAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASCGFmM2bN+vGG2/UsGHDNGzYMOXk5GjPnj32ccuyVFxcrNTUVMXFxWnatGk6fvx40DkCgYCWLFmixMREDR06VAUFBWpoaAiqaW1tVWFhoVwul1wulwoLC3XmzJmedwkAQATweDz64x//2OVWV1c30EPsV9GhFI8cOVJPPPGEvvWtb0mSnnrqKd155506cuSIbrjhBq1du1br1q1TWVmZ0tPT9fjjj+v222/XiRMnFB8fL0kqKirSSy+9pIqKCo0YMULLly9Xfn6+ampqFBUVJUmaO3euGhoaVFlZKUm69957VVhYqJdeeqk3ewcAwBgej0cZ4zPkP+8f6KGEjZBCzB133BH0evXq1dq8ebMOHjyo66+/Xhs2bNCqVas0e/ZsSX8NOW63W7t27dLixYvl8/m0bds27dixQzNmzJAklZeXKy0tTXv37lVeXp7q6upUWVmpgwcPKisrS5K0detW5eTk6MSJExo3blxv9A0AgFG8Xq/85/2at2We3OnuTmtq99Zqz+o9nR6LRD2+J+bixYuqqKjQp59+qpycHJ0+fVpNTU3Kzc21a5xOp6ZOnarq6mpJUk1NjS5cuBBUk5qaqszMTLvmwIEDcrlcdoCRpOzsbLlcLrumM4FAQG1tbUEbAACRxp3uVtrEtE63EaNGDPTw+lXIIeadd97RV77yFTmdTt13333avXu3rr/+ejU1NUmS3O7gdOh2u+1jTU1Nio2N1fDhw7utSUpK6vBzk5KS7JrOlJSU2PfQuFwupaWlhdoaAAAwSMghZty4cTp69KgOHjyof/qnf9L8+fNVW1trH3c4HEH1lmV12He5y2s6q7/SeVauXCmfz2dv9fX1V9sSAAAwUMghJjY2Vt/61rc0efJklZSUaOLEifrFL36h5ORkSepwtaSlpcW+OpOcnKz29na1trZ2W9Pc3Nzh53788ccdrvJ8kdPptFdNXdoAAEDk+tLfE2NZlgKBgMaMGaPk5GRVVVXZx9rb27Vv3z5NmTJFkjRp0iTFxMQE1TQ2NurYsWN2TU5Ojnw+nw4fPmzXHDp0SD6fz64BAAAIaXXSo48+qlmzZiktLU1nz55VRUWFXn/9dVVWVsrhcKioqEhr1qzR2LFjNXbsWK1Zs0ZDhgzR3LlzJUkul0sLFy7U8uXLNWLECCUkJGjFihWaMGGCvVpp/PjxmjlzphYtWqQtW7ZI+usS6/z8fFYmAQAAW0ghprm5WYWFhWpsbJTL5dKNN96oyspK3X777ZKkhx9+WH6/X/fff79aW1uVlZWlV1991f6OGElav369oqOjNWfOHPn9fk2fPl1lZWX2d8RI0s6dO7V06VJ7FVNBQYE2btzYG/0CAIAIEVKI2bZtW7fHHQ6HiouLVVxc3GXN4MGDVVpaqtLS0i5rEhISVF5eHsrQAADANYZnJwEAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRQgoxJSUluummmxQfH6+kpCTdddddOnHiRFDNggUL5HA4grbs7OygmkAgoCVLligxMVFDhw5VQUGBGhoagmpaW1tVWFgol8sll8ulwsJCnTlzpmddAgCAIHV1dfrjH//Y5ebxeAZ6iFcUHUrxvn379MADD+imm27SZ599plWrVik3N1e1tbUaOnSoXTdz5kxt377dfh0bGxt0nqKiIr300kuqqKjQiBEjtHz5cuXn56umpkZRUVGSpLlz56qhoUGVlZWSpHvvvVeFhYV66aWXetwsAADXurbmNjkGOTRv3rxu6+KGxOndunc1atSofhpZ6EIKMZcCxSXbt29XUlKSampq9N3vftfe73Q6lZyc3Ok5fD6ftm3bph07dmjGjBmSpPLycqWlpWnv3r3Ky8tTXV2dKisrdfDgQWVlZUmStm7dqpycHJ04cULjxo0LqUkAAPBXfp9f1ueW5m2ZJ3e6u9Oa5pPNKl9cLq/XGzkh5nI+n0+SlJCQELT/9ddfV1JSkq677jpNnTpVq1evVlJSkiSppqZGFy5cUG5url2fmpqqzMxMVVdXKy8vTwcOHJDL5bIDjCRlZ2fL5XKpurqaEAMAwJfkTncrbWLaQA/jS+lxiLEsS8uWLdN3vvMdZWZm2vtnzZqlH/3oRxo9erROnz6tn/70p/re976nmpoaOZ1ONTU1KTY2VsOHDw86n9vtVlNTkySpqanJDj1flJSUZNdcLhAIKBAI2K/b2tp62hoAADBAj0PMgw8+qD/96U/av39/0P67777b/nVmZqYmT56s0aNH63e/+51mz57d5fksy5LD4bBff/HXXdV8UUlJiX72s5+F2gYAADBUj5ZYL1myRC+++KJee+01jRw5stvalJQUjR49WqdOnZIkJScnq729Xa2trUF1LS0tcrvddk1zc3OHc3388cd2zeVWrlwpn89nb/X19T1pDQAAGCKkKzGWZWnJkiXavXu3Xn/9dY0ZM+aK7/nkk09UX1+vlJQUSdKkSZMUExOjqqoqzZkzR5LU2NioY8eOae3atZKknJwc+Xw+HT58WDfffLMk6dChQ/L5fJoyZUqnP8fpdMrpdIbSDgAAYcPj8cjr9XZ5vK6urh9HY4aQQswDDzygXbt26be//a3i4+Pt+1NcLpfi4uJ07tw5FRcX6wc/+IFSUlL04Ycf6tFHH1ViYqK+//3v27ULFy7U8uXLNWLECCUkJGjFihWaMGGCvVpp/PjxmjlzphYtWqQtW7ZI+usS6/z8fG7qBQBEHI/Ho4zxGfKf9w/0UIwSUojZvHmzJGnatGlB+7dv364FCxYoKipK77zzjp5++mmdOXNGKSkpuu222/Tss88qPj7erl+/fr2io6M1Z84c+f1+TZ8+XWVlZfZ3xEjSzp07tXTpUnsVU0FBgTZu3NjTPgEACFter1f+8/5ulz3X7q3VntV7+nlk4S3kj5O6ExcXp1deeeWK5xk8eLBKS0tVWlraZU1CQoLKy8tDGR4AAEbrbtlz88mO94pe63h2EgAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIwUUogpKSnRTTfdpPj4eCUlJemuu+7SiRMngmosy1JxcbFSU1MVFxenadOm6fjx40E1gUBAS5YsUWJiooYOHaqCggI1NDQE1bS2tqqwsFAul0sul0uFhYU6c+ZMz7oEAAARJ6QQs2/fPj3wwAM6ePCgqqqq9Nlnnyk3N1effvqpXbN27VqtW7dOGzdu1FtvvaXk5GTdfvvtOnv2rF1TVFSk3bt3q6KiQvv379e5c+eUn5+vixcv2jVz587V0aNHVVlZqcrKSh09elSFhYW90DIAAIgE0aEUV1ZWBr3evn27kpKSVFNTo+9+97uyLEsbNmzQqlWrNHv2bEnSU089JbfbrV27dmnx4sXy+Xzatm2bduzYoRkzZkiSysvLlZaWpr179yovL091dXWqrKzUwYMHlZWVJUnaunWrcnJydOLECY0bN643egcAAAb7UvfE+Hw+SVJCQoIk6fTp02pqalJubq5d43Q6NXXqVFVXV0uSampqdOHChaCa1NRUZWZm2jUHDhyQy+WyA4wkZWdny+Vy2TWXCwQCamtrC9oAAEDk6nGIsSxLy5Yt03e+8x1lZmZKkpqamiRJbrc7qNbtdtvHmpqaFBsbq+HDh3dbk5SU1OFnJiUl2TWXKykpse+fcblcSktL62lrAADAAD0OMQ8++KD+9Kc/6ZlnnulwzOFwBL22LKvDvstdXtNZfXfnWblypXw+n73V19dfTRsAAMBQPQoxS5Ys0YsvvqjXXntNI0eOtPcnJydLUoerJS0tLfbVmeTkZLW3t6u1tbXbmubm5g4/9+OPP+5wlecSp9OpYcOGBW0AACByhRRiLMvSgw8+qP/5n//R73//e40ZMybo+JgxY5ScnKyqqip7X3t7u/bt26cpU6ZIkiZNmqSYmJigmsbGRh07dsyuycnJkc/n0+HDh+2aQ4cOyefz2TUAAODaFtLqpAceeEC7du3Sb3/7W8XHx9tXXFwul+Li4uRwOFRUVKQ1a9Zo7NixGjt2rNasWaMhQ4Zo7ty5du3ChQu1fPlyjRgxQgkJCVqxYoUmTJhgr1YaP368Zs6cqUWLFmnLli2SpHvvvVf5+fmsTAIAAJJCDDGbN2+WJE2bNi1o//bt27VgwQJJ0sMPPyy/36/7779fra2tysrK0quvvqr4+Hi7fv369YqOjtacOXPk9/s1ffp0lZWVKSoqyq7ZuXOnli5daq9iKigo0MaNG3vSIwAAiEAhhRjLsq5Y43A4VFxcrOLi4i5rBg8erNLSUpWWlnZZk5CQoPLy8lCGBwAAriE8OwkAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJFCWp0EAABC5/F45PV6uzxeV1fXj6OJHIQYAAD6kMfjUcb4DPnP+wd6KBGHEAMAQB/yer3yn/dr3pZ5cqd3/vy/2r212rN6Tz+PzHyEGAAA+oE73a20iWmdHms+2fGhx7gybuwFAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACOFHGLeeOMN3XHHHUpNTZXD4dALL7wQdHzBggVyOBxBW3Z2dlBNIBDQkiVLlJiYqKFDh6qgoEANDQ1BNa2trSosLJTL5ZLL5VJhYaHOnDkTcoMAACAyhRxiPv30U02cOFEbN27ssmbmzJlqbGy0t5dffjnoeFFRkXbv3q2Kigrt379f586dU35+vi5evGjXzJ07V0ePHlVlZaUqKyt19OhRFRYWhjpcAAAQoaJDfcOsWbM0a9asbmucTqeSk5M7Pebz+bRt2zbt2LFDM2bMkCSVl5crLS1Ne/fuVV5enurq6lRZWamDBw8qKytLkrR161bl5OToxIkTGjduXKjDBgAAEaZP7ol5/fXXlZSUpPT0dC1atEgtLS32sZqaGl24cEG5ubn2vtTUVGVmZqq6ulqSdODAAblcLjvASFJ2drZcLpddc7lAIKC2tragDQAARK5eDzGzZs3Szp079fvf/17/8R//obfeekvf+973FAgEJElNTU2KjY3V8OHDg97ndrvV1NRk1yQlJXU4d1JSkl1zuZKSEvv+GZfLpbS0tF7uDAAAhJOQP066krvvvtv+dWZmpiZPnqzRo0frd7/7nWbPnt3l+yzLksPhsF9/8ddd1XzRypUrtWzZMvt1W1sbQQYA0Oc8Ho+8Xm+Xx+vq6vpxNNeWXg8xl0tJSdHo0aN16tQpSVJycrLa29vV2toadDWmpaVFU6ZMsWuam5s7nOvjjz+W2+3u9Oc4nU45nc4+6AAAgM55PB5ljM+Q/7x/oIdyTerzEPPJJ5+ovr5eKSkpkqRJkyYpJiZGVVVVmjNnjiSpsbFRx44d09q1ayVJOTk58vl8Onz4sG6++WZJ0qFDh+Tz+eygAwDAQPN6vfKf92velnlyp3f+j+zavbXas3pPP4/s2hByiDl37pzee+89+/Xp06d19OhRJSQkKCEhQcXFxfrBD36glJQUffjhh3r00UeVmJio73//+5Ikl8ulhQsXavny5RoxYoQSEhK0YsUKTZgwwV6tNH78eM2cOVOLFi3Sli1bJEn33nuv8vPzWZkEAAg77nS30iZ2fgtD88mOnyygd4QcYt5++23ddttt9utL96HMnz9fmzdv1jvvvKOnn35aZ86cUUpKim677TY9++yzio+Pt9+zfv16RUdHa86cOfL7/Zo+fbrKysoUFRVl1+zcuVNLly61VzEVFBR0+900AADg2hJyiJk2bZosy+ry+CuvvHLFcwwePFilpaUqLS3tsiYhIUHl5eWhDg8AAFwjeHYSAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIwUPdADAAAgXHk8Hnm93i6P19XV9eNocDlCDAAAnfB4PMoYnyH/ef9ADwVdIMQAANAJr9cr/3m/5m2ZJ3e6u9Oa2r212rN6Tz+PDJcQYgAA6IY73a20iWmdHms+2dzPo8EXcWMvAAAwEiEGAAAYiRADAACMRIgBAABGCjnEvPHGG7rjjjuUmpoqh8OhF154Iei4ZVkqLi5Wamqq4uLiNG3aNB0/fjyoJhAIaMmSJUpMTNTQoUNVUFCghoaGoJrW1lYVFhbK5XLJ5XKpsLBQZ86cCblBAAAQmUIOMZ9++qkmTpyojRs3dnp87dq1WrdunTZu3Ki33npLycnJuv3223X27Fm7pqioSLt371ZFRYX279+vc+fOKT8/XxcvXrRr5s6dq6NHj6qyslKVlZU6evSoCgsLe9AiAACIRCEvsZ41a5ZmzZrV6THLsrRhwwatWrVKs2fPliQ99dRTcrvd2rVrlxYvXiyfz6dt27Zpx44dmjFjhiSpvLxcaWlp2rt3r/Ly8lRXV6fKykodPHhQWVlZkqStW7cqJydHJ06c0Lhx43raLwAAiBC9ek/M6dOn1dTUpNzcXHuf0+nU1KlTVV1dLUmqqanRhQsXgmpSU1OVmZlp1xw4cEAul8sOMJKUnZ0tl8tl11wuEAiora0taAMAAJGrV0NMU1OTJMntDv5mQ7fbbR9rampSbGyshg8f3m1NUlJSh/MnJSXZNZcrKSmx759xuVxKS+v8i4kAAEBk6JPVSQ6HI+i1ZVkd9l3u8prO6rs7z8qVK+Xz+eytvr6+ByMHAACm6NUQk5ycLEkdrpa0tLTYV2eSk5PV3t6u1tbWbmuamzt+lfPHH3/c4SrPJU6nU8OGDQvaAABA5OrVEDNmzBglJyerqqrK3tfe3q59+/ZpypQpkqRJkyYpJiYmqKaxsVHHjh2za3JycuTz+XT48GG75tChQ/L5fHYNAAC4toW8OuncuXN677337NenT5/W0aNHlZCQoFGjRqmoqEhr1qzR2LFjNXbsWK1Zs0ZDhgzR3LlzJUkul0sLFy7U8uXLNWLECCUkJGjFihWaMGGCvVpp/PjxmjlzphYtWqQtW7ZIku69917l5+ezMgkAAEjqQYh5++23ddttt9mvly1bJkmaP3++ysrK9PDDD8vv9+v+++9Xa2ursrKy9Oqrryo+Pt5+z/r16xUdHa05c+bI7/dr+vTpKisrU1RUlF2zc+dOLV261F7FVFBQ0OV30wAAgGtPyCFm2rRpsiyry+MOh0PFxcUqLi7usmbw4MEqLS1VaWlplzUJCQkqLy8PdXgAAOAawbOTAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMFPISawAAIoHH45HX6+3yeF1dXT+OBj1BiAEAXHM8Ho8yxmfIf94/0EPBl0CIAQBcc7xer/zn/Zq3ZZ7c6Z0/WLh2b632rN7TzyNDKAgxAIBrljvdrbSJaZ0eaz7Z3M+jQai4sRcAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAwEiEGAAAYiRADAACMRIgBAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGCk6IEeAAAAvcnj8cjr9XZbU1dX10+jQV8ixAAAIobH41HG+Az5z/sHeijoB4QYAEDE8Hq98p/3a96WeXKnu7usq91bqz2r9/TjyNAXCDEAgIjjTncrbWJal8ebTzb342jMdaWP3RITEzVq1Kh+Gk1HhBgAABCkrblNjkEOzZs3r9u6uCFxerfu3QELMoQYAAAQxO/zy/rc6vZjueaTzSpfXC6v10uIAQAA4eVKH8sNNL4nBgAAGIkQAwAAjNTrIaa4uFgOhyNoS05Oto9blqXi4mKlpqYqLi5O06ZN0/Hjx4POEQgEtGTJEiUmJmro0KEqKChQQ0NDbw8VAAAYrE+uxNxwww1qbGy0t3feecc+tnbtWq1bt04bN27UW2+9peTkZN1+++06e/asXVNUVKTdu3eroqJC+/fv17lz55Sfn6+LFy/2xXABAICB+uTG3ujo6KCrL5dYlqUNGzZo1apVmj17tiTpqaeektvt1q5du7R48WL5fD5t27ZNO3bs0IwZMyRJ5eXlSktL0969e5WXl9cXQwYAAIbpkysxp06dUmpqqsaMGaMf//jH+uCDDyRJp0+fVlNTk3Jzc+1ap9OpqVOnqrq6WpJUU1OjCxcuBNWkpqYqMzPTrgEAAOj1KzFZWVl6+umnlZ6erubmZj3++OOaMmWKjh8/rqamJkmS2x285tztduvPf/6zJKmpqUmxsbEaPnx4h5pL7+9MIBBQIBCwX7e1tfVWSwAAIAz1eoiZNWuW/esJEyYoJydH3/zmN/XUU08pOztbkuRwOILeY1lWh32Xu1JNSUmJfvazn32JkQMAAJP0+RLroUOHasKECTp16pR9n8zlV1RaWlrsqzPJyclqb29Xa2trlzWdWblypXw+n73V19f3cicAACCc9Pk39gYCAdXV1enWW2/VmDFjlJycrKqqKn3729+WJLW3t2vfvn3693//d0nSpEmTFBMTo6qqKs2ZM0eS1NjYqGPHjmnt2rVd/hyn0ymn09nX7QAABpDH45HX6+3y+JUeWIjI0ushZsWKFbrjjjs0atQotbS06PHHH1dbW5vmz58vh8OhoqIirVmzRmPHjtXYsWO1Zs0aDRkyRHPnzpUkuVwuLVy4UMuXL9eIESOUkJCgFStWaMKECfZqJQDAtcfj8ShjfIb85/0DPRSEiV4PMQ0NDfqHf/gHeb1effWrX1V2drYOHjyo0aNHS5Iefvhh+f1+3X///WptbVVWVpZeffVVxcfH2+dYv369oqOjNWfOHPn9fk2fPl1lZWWKiorq7eECAAzh9XrlP+/v9qGEtXtrtWf1nn4eGQZKr4eYioqKbo87HA4VFxeruLi4y5rBgwertLRUpaWlvTw6AIDpunsoYfPJ5n4eDQYSz04CAABGIsQAAAAjEWIAAICRCDEAAMBIhBgAAGAkQgwAADASIQYAABiJEAMAAIxEiAEAAEYixAAAACMRYgAAgJEIMQAAwEiEGAAAYKRef4o1AAA94fF45PV6uzxeV1fXj6OBCQgxAIA+d6WA0tjYqB/+6If6i/8v/TgqmI4QAwDoUx6PRxnjM+Q/779i7bwt8+ROd3d6rHZvrfas3tPbw4PBCDEAgD7l9XrlP++/qoDiTncrbWJapzXNJ5v7cpgwECEGANAvCCjobaxOAgAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJFYnAQC+FL5pFwOFEAMA6LFQvsgO6G2EGABAj4XyRXZAbyPEAAC6dLUfFfFFdhgIhBgAQKf4qAjhjhADAOgUHxUh3BFiAADd4qMihCu+JwYAABiJKzEAcI3i+11gOkIMAFyDuGkXkYAQAwDXIG7aRSQgxADANYybdmEyQgwARCDud8G1gBADABGG+11wrSDEAIBhruYqC/e74FpAiAEAg4RylYX7XRDpCDEAEEa4ygJcPUIMAPSTKwWUxsZG/fBHP9Rf/H+54rm4ygIQYgCgX4TyMRBXWYCrQ4gBgF7Qmx8DcZUFuDphH2I2bdqkn//852psbNQNN9ygDRs26NZbbx3oYQG4hvAxEBCewjrEPPvssyoqKtKmTZt0yy23aMuWLZo1a5Zqa2s1atSogR4egDB3pfAhSYFAQE6ns8vjoQQUPgYC+ldYh5h169Zp4cKFuueeeyRJGzZs0CuvvKLNmzerpKRkgEcHRI6r+cs+MTGxV/7x0BvB4mpqrjZ8OAY5ZH1udVsjXV1A4SoL0L/CNsS0t7erpqZGjzzySND+3NxcVVdXd6gPBAIKBAL2a5/PJ0lqa2vrk/E1NTWpqamp25pBgwbp888/p4aasK5pbm5W4f8rVOAvgS5rJMk52KkdT++Q2935X+S9+bPkkHSlXHE1NZK+9+D3dN3I6zo95vmjR2//5u2rqmn3tyvwaefjvhC4IEmq///qu6xpOtVETZjUhOOYTKxpea9FknTu3Lle/bv20rks6yr+B7fC1P/+7/9akqw//OEPQftXr15tpaend6h/7LHHLP31jzQ2NjY2NjY2w7f6+vorZoWwvRJzicPhCHptWVaHfZK0cuVKLVu2zH79+eef6//+7/80YsSIoPq2tjalpaWpvr5ew4YN67uBhwF6jUz0GrmupX7pNTL1Rq+WZens2bNKTU29Ym3YhpjExERFRUV1+MimpaWl08vZTqezw+fj1113XZfnHzZsWMT/x3QJvUYmeo1c11K/9BqZvmyvLpfrquoG9fgn9LHY2FhNmjRJVVVVQfurqqo0ZcqUARoVAAAIF2F7JUaSli1bpsLCQk2ePFk5OTn69a9/LY/Ho/vuu2+ghwYAAAZYWIeYu+++W5988on+7d/+TY2NjcrMzNTLL7+s0aNH9/icTqdTjz322BWXb0YCeo1M9Bq5rqV+6TUy9XevDsu6mjVMAAAA4SVs74kBAADoDiEGAAAYiRADAACMRIgBAABGMj7EbNq0SWPGjNHgwYM1adIkvfnmm93WBwIBrVq1SqNHj5bT6dQ3v/lNPfnkk0E1zz//vK6//no5nU5df/312r17d1+2cNV6u9eysjI5HI4O21/+cuWn9fa1UHpdsGBBp33ccMMNQXXhOq9S7/cbKXMrSTt37tTEiRM1ZMgQpaSk6Cc/+Yk++eSToJpwndve7jWS5vVXv/qVxo8fr7i4OI0bN05PP/10h5pImdcr9Rqu8/rGG2/ojjvuUGpqqhwOh1544YUrvmffvn2aNGmSBg8erG984xv6z//8zw41vTqvvfKgowFSUVFhxcTEWFu3brVqa2uthx56yBo6dKj15z//ucv3FBQUWFlZWVZVVZV1+vRp69ChQ0HPZ6qurraioqKsNWvWWHV1ddaaNWus6Oho6+DBg/3RUpf6otft27dbw4YNsxobG4O2gRZqr2fOnAkaf319vZWQkGA99thjdk24zqtl9U2/kTK3b775pjVo0CDrF7/4hfXBBx9Yb775pnXDDTdYd911l10TrnPbF71Gyrxu2rTJio+PtyoqKqz333/feuaZZ6yvfOUr1osvvmjXRMq8Xk2v4TqvL7/8srVq1Srr+eeftyRZu3fv7rb+gw8+sIYMGWI99NBDVm1trbV161YrJibGeu655+ya3p5Xo0PMzTffbN13331B+zIyMqxHHnmk0/o9e/ZYLpfL+uSTT7o855w5c6yZM2cG7cvLy7N+/OMff/kBfwl90ev27dstl8vVm8PsFaH2erndu3dbDofD+vDDD+194TqvltU3/UbK3P785z+3vvGNbwTt++Uvf2mNHDnSfh2uc9sXvUbKvObk5FgrVqwI2vfQQw9Zt9xyi/06Uub1anoN13n9oqsJMQ8//LCVkZERtG/x4sVWdna2/bq359XYj5Pa29tVU1Oj3NzcoP25ubmqrq7u9D0vvviiJk+erLVr1+prX/ua0tPTtWLFCvn9frvmwIEDHc6Zl5fX5Tn7Q1/1Kv31EeqjR4/WyJEjlZ+fryNHjvRZH1ejJ71ebtu2bZoxY0bQlyKG47xKfdevFBlzO2XKFDU0NOjll1+WZVlqbm7Wc889p7//+7+3a8JxbvuqVyky5jUQCGjw4MFB++Li4nT48GFduHBBUuTM69X0KoXfvPZEV3P29ttv99m8GhtivF6vLl682OFhkG63u8NDIy/54IMPtH//fh07dky7d+/Whg0b9Nxzz+mBBx6wa5qamkI6Z3/oq14zMjJUVlamF198Uc8884wGDx6sW265RadOnerTfrrTk16/qLGxUXv27NE999wTtD8c51Xqu34jZW6nTJminTt36u6771ZsbKySk5N13XXXqbS01K4Jx7ntq14jZV7z8vL0X//1X6qpqZFlWXr77bf15JNP6sKFC/J6vZIiZ16vptdwnNee6GrOPvvssz6bV2NDzCUOhyPotWVZHfZd8vnnn8vhcGjnzp26+eab9Xd/93dat26dysrKgq5QhHLO/tTbvWZnZ2vevHmaOHGibr31Vv3mN79Renp60B+aA6Wnc1BWVqbrrrtOd911V6+dsz/0dr+RMre1tbVaunSp/vVf/1U1NTWqrKzU6dOnOzw/LVzntrd7jZR5/elPf6pZs2YpOztbMTExuvPOO7VgwQJJUlRUVI/O2Z96u9dwntdQdfZ7c/n+3pxXY0NMYmKioqKiOqS3lpaWDinvkpSUFH3ta18LesT3+PHjZVmWGhoaJEnJyckhnbM/9FWvlxs0aJBuuummAU3/Pen1Esuy9OSTT6qwsFCxsbFBx8JxXqW+6/dyps5tSUmJbrnlFv3Lv/yLbrzxRuXl5WnTpk168skn1djYKCk857aver2cqfMaFxenJ598UufPn9eHH34oj8ejr3/964qPj1diYqKkyJnXq+n1cuEwrz3R1ZxFR0drxIgR3db0dF6NDTGxsbGaNGmSqqqqgvZXVVVpypQpnb7nlltu0UcffaRz587Z+06ePKlBgwZp5MiRkqScnJwO53z11Ve7PGd/6KteL2dZlo4ePaqUlJTeG3yIetLrJfv27dN7772nhQsXdjgWjvMq9V2/lzN1bs+fP69Bg4L/mLr0r9dL/8ILx7ntq14vZ+q8XhITE6ORI0cqKipKFRUVys/Pt38PImVeL+mu18uFw7z2RFdzNnnyZMXExHRb0+N57dHtwGHi0lK3bdu2WbW1tVZRUZE1dOhQe5XGI488YhUWFtr1Z8+etUaOHGn98Ic/tI4fP27t27fPGjt2rHXPPffYNX/4wx+sqKgo64knnrDq6uqsJ554IqyW9fVmr8XFxVZlZaX1/vvvW0eOHLF+8pOfWNHR0dahQ4f6vb8vCrXXS+bNm2dlZWV1es5wnVfL6pt+I2Vut2/fbkVHR1ubNm2y3n//fWv//v3W5MmTrZtvvtmuCde57YteI2VeT5w4Ye3YscM6efKkdejQIevuu++2EhISrNOnT9s1kTKvV9NruM7r2bNnrSNHjlhHjhyxJFnr1q2zjhw5Yi8nv7zXS0us//mf/9mqra21tm3b1mGJdW/Pq9EhxrIs61e/+pU1evRoKzY21vrbv/1ba9++ffax+fPnW1OnTg2qr6urs2bMmGHFxcVZI0eOtJYtW2adP38+qOa///u/rXHjxlkxMTFWRkaG9fzzz/dHK1fU270WFRVZo0aNsmJjY62vfvWrVm5urlVdXd1f7XQr1F7PnDljxcXFWb/+9a+7PGe4zqtl9X6/kTS3v/zlL63rr7/eiouLs1JSUqx//Md/tBoaGoJqwnVue7vXSJnX2tpa62/+5m+suLg4a9iwYdadd95pvfvuux3OGQnzejW9huu8vvbaa5akDtv8+fMty+r8v+HXX3/d+va3v23FxsZaX//6163Nmzd3OG9vzqvDsrq4TgkAABDGjL0nBgAAXNsIMQAAwEiEGAAAYCRCDAAAMBIhBgAAGIkQAwAAjESIAQAARiLEAAAAIxFiAACAkQgxAADASIQYAABgJEIMAAAw0v8P1ArwWnncXYYAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(test_with_pred[\"predicted_score\"], color='lightgreen', ec='black', bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate using MRR@10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MRR@10:  0.9151666666666675\n"
     ]
    }
   ],
   "source": [
    "dict_qid_all_pid_prediction = {}\n",
    "\n",
    "for qid, group in test_with_pred.groupby('query_id'):\n",
    "    dict_qid_all_pid_prediction[qid] = group['passage_id'].values.tolist()\n",
    "\n",
    "\n",
    "topn = 10\n",
    "mrr = 0.0\n",
    "qids = []\n",
    "rrs = []\n",
    "\n",
    "# iterate over all the qid (keys in the dictionary)\n",
    "for qid in dict_qid_all_pid_prediction:\n",
    "\n",
    "    rr = 0.0\n",
    "    \n",
    "    # iterate over all the element of topn pid [pid1, pid2, ... pid topn]\n",
    "    for i, pid in enumerate(dict_qid_all_pid_prediction[qid][: topn + 1]):\n",
    "\n",
    "        # check if the pid has relenvace > 0\n",
    "        if qid in dict_of_dict_qid_pid_rel and pid in dict_of_dict_qid_pid_rel[qid] and dict_of_dict_qid_pid_rel[qid][pid] > 0:\n",
    "        \n",
    "            rr = 1 / (i+1)\n",
    "        \n",
    "            break\n",
    "    qids.append(qid)\n",
    "    rrs.append(rr)\n",
    "    mrr += rr\n",
    "mrr /= len(dict_qid_all_pid_prediction)\n",
    "print(\"MRR@10: \", mrr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate using MAP@10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(20000, 9)\n",
      "(20000, 3)\n",
      "(20000, 10)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>query</th>\n",
       "      <th>used</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>username</th>\n",
       "      <th>rel</th>\n",
       "      <th>query_cleaned</th>\n",
       "      <th>predicted_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_49_115778700</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2671</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.962330</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_26_450017756</td>\n",
       "      <td>2</td>\n",
       "      <td>11.8537</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.948867</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_19_202396335</td>\n",
       "      <td>3</td>\n",
       "      <td>11.2738</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.954807</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_28_623516004</td>\n",
       "      <td>4</td>\n",
       "      <td>11.0361</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.954931</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_02_47181174</td>\n",
       "      <td>5</td>\n",
       "      <td>10.9476</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.959981</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   query_id                                     query used  \\\n",
       "0    203324  him functions to the paper health record   Q0   \n",
       "1    203324  him functions to the paper health record   Q0   \n",
       "2    203324  him functions to the paper health record   Q0   \n",
       "3    203324  him functions to the paper health record   Q0   \n",
       "4    203324  him functions to the paper health record   Q0   \n",
       "\n",
       "                     passage_id  rank    score  username  rel  \\\n",
       "0  msmarco_passage_49_115778700     1  13.2671  Anserini    1   \n",
       "1  msmarco_passage_26_450017756     2  11.8537  Anserini    1   \n",
       "2  msmarco_passage_19_202396335     3  11.2738  Anserini    1   \n",
       "3  msmarco_passage_28_623516004     4  11.0361  Anserini    1   \n",
       "4   msmarco_passage_02_47181174     5  10.9476  Anserini    1   \n",
       "\n",
       "                              query_cleaned  predicted_score  \n",
       "0  him functions to the paper health record         0.962330  \n",
       "1  him functions to the paper health record         0.948867  \n",
       "2  him functions to the paper health record         0.954807  \n",
       "3  him functions to the paper health record         0.954931  \n",
       "4  him functions to the paper health record         0.959981  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(query_test.shape)\n",
    "print(test_with_pred.shape)\n",
    "\n",
    "result = query_test.merge(test_with_pred, how = \"inner\", on = [\"query_id\", \"passage_id\"])\n",
    "print(result.shape)\n",
    "display(result.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function for calculating average precision for a query\n",
    "def average_precision(qid, topn = 10):\n",
    "    \n",
    "    # ========= for each query_id, do the following:\n",
    "    #print(f\"Started to work on query_id {qid}\")\n",
    "    # get passage_id (each query_id would have upto 10 passage_id)\n",
    "    qresult = result.loc[result[\"query_id\"] == qid, [\"passage_id\", \"rel\", \"predicted_score\"]]\n",
    "\n",
    "\n",
    "    \n",
    "    # sort rows so that passages with the highest similarity ranks as top \n",
    "    qresult.sort_values(by=\"predicted_score\", ascending = False, inplace=True)\n",
    "\n",
    "    # # Taking Top 10 documents for the evaluation\n",
    "    ranking = qresult.head(topn)['rel'].values\n",
    "    #display(qresult.head(topn))\n",
    "    \n",
    "    # # Calculating precision\n",
    "    precision=[]\n",
    "    for i in range(1,11):\n",
    "      if ranking[i-1]:\n",
    "        precision.append(np.sum(ranking[:i])/i)\n",
    "    \n",
    "    # If no relevant document in list then return 0\n",
    "    if precision==[]:\n",
    "      return 0\n",
    "    #print(len(precision))\n",
    "    #print(precision)\n",
    "    avg_precision = np.mean(precision)\n",
    "    #print(\"Avg precision: {}\".format(avg_precision))\n",
    "    return avg_precision#, len(precision), qresult.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "result[\"avg_precision\"] = result.apply(lambda x: average_precision(x['query_id'])\n",
    "                                    ,axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_id</th>\n",
       "      <th>query</th>\n",
       "      <th>used</th>\n",
       "      <th>passage_id</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>username</th>\n",
       "      <th>rel</th>\n",
       "      <th>query_cleaned</th>\n",
       "      <th>predicted_score</th>\n",
       "      <th>avg_precision</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_49_115778700</td>\n",
       "      <td>1</td>\n",
       "      <td>13.2671</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.962330</td>\n",
       "      <td>0.77381</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_26_450017756</td>\n",
       "      <td>2</td>\n",
       "      <td>11.8537</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.948867</td>\n",
       "      <td>0.77381</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_19_202396335</td>\n",
       "      <td>3</td>\n",
       "      <td>11.2738</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.954807</td>\n",
       "      <td>0.77381</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_28_623516004</td>\n",
       "      <td>4</td>\n",
       "      <td>11.0361</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.954931</td>\n",
       "      <td>0.77381</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>203324</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>Q0</td>\n",
       "      <td>msmarco_passage_02_47181174</td>\n",
       "      <td>5</td>\n",
       "      <td>10.9476</td>\n",
       "      <td>Anserini</td>\n",
       "      <td>1</td>\n",
       "      <td>him functions to the paper health record</td>\n",
       "      <td>0.959981</td>\n",
       "      <td>0.77381</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   query_id                                     query used  \\\n",
       "0    203324  him functions to the paper health record   Q0   \n",
       "1    203324  him functions to the paper health record   Q0   \n",
       "2    203324  him functions to the paper health record   Q0   \n",
       "3    203324  him functions to the paper health record   Q0   \n",
       "4    203324  him functions to the paper health record   Q0   \n",
       "\n",
       "                     passage_id  rank    score  username  rel  \\\n",
       "0  msmarco_passage_49_115778700     1  13.2671  Anserini    1   \n",
       "1  msmarco_passage_26_450017756     2  11.8537  Anserini    1   \n",
       "2  msmarco_passage_19_202396335     3  11.2738  Anserini    1   \n",
       "3  msmarco_passage_28_623516004     4  11.0361  Anserini    1   \n",
       "4   msmarco_passage_02_47181174     5  10.9476  Anserini    1   \n",
       "\n",
       "                              query_cleaned  predicted_score  avg_precision  \n",
       "0  him functions to the paper health record         0.962330        0.77381  \n",
       "1  him functions to the paper health record         0.948867        0.77381  \n",
       "2  him functions to the paper health record         0.954807        0.77381  \n",
       "3  him functions to the paper health record         0.954931        0.77381  \n",
       "4  him functions to the paper health record         0.959981        0.77381  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean Average Precision=> 0.836973914399093\n"
     ]
    }
   ],
   "source": [
    "print('Mean Average Precision=>', result[\"avg_precision\"].mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.to_csv(\"./output/test_prediction_biencoder_refit_avg_precision_MiniLM-L6-v2_2024-04-22.csv\"\n",
    "                                       , index = False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}