Homework 3 Part 3 - Deep Learning Models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9egyhuvU3_GI"
   },
   "source": [
    "# SI 650 / EECS 549: Homework 3 Part 3\n",
    "\n",
    "Homework 3 Part 3 will have you working with deep learning models in a variety of ways. You will likely need to run this on Great Lakes unless you have access to a GPU elsewhere (or be prepared to wait a long time). You should have completed Parts 1 and 2 before attempting this notebook to familiarize yourself with how PyTerrier works.\n",
    "\n",
    "In Part 3, you'll try the following tasks:\n",
    " - Use a large language model to re-rank content\n",
    " - Use a text-to-text model to perform query augmentation\n",
    " - Train a deep learning IR model and compare its performance.\n",
    " \n",
    "The first two of these tasks will rely on models that we've pretrained for you. However, we've also provided code for how to train these. In the third task, you'll do one simple training and evaluate its results.\n",
    "\n",
    "For the first two tasks, we've provided most of the code. *You are expected to submit results showing that you have successfully executed it*. You'll need to understand some of the code to complete task 3, which requires you to write new code.\n",
    "\n",
    "As with the past notebooks, you'll need to have `JAVA_HOME` set, which will need to be run on Great Lakes. You can potentially set it in the notebook with a Jupyter command like\n",
    "```\n",
    "!export JAVA_HOME=/fill/in/the/path/to/the/JDK/here\n",
    "```\n",
    "by first figuring out where the JDK is installed. (This is setting the `JAVA_HOME` environment variable in the unix way). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mV0C6jJvqhMR"
   },
   "source": [
    "### Install the PyTerrier extensions\n",
    "\n",
    "You'll need two extensions for [OpenNIR](https://opennir.net/) and [doc2query](https://github.com/terrierteam/pyterrier_doc2query). We've provided the package install commands in comments below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python 3.8.8\n"
     ]
    }
   ],
   "source": [
    "# !python==3.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "AkIR_PXdet7R"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting git+https://github.com/Georgetown-IR-Lab/OpenNIR\n",
      "  Cloning https://github.com/Georgetown-IR-Lab/OpenNIR to /tmp/pip-req-build-cty2kiwx\n",
      "  Running command git clone -q https://github.com/Georgetown-IR-Lab/OpenNIR /tmp/pip-req-build-cty2kiwx\n",
      "  fatal: unable to access 'https://github.com/Georgetown-IR-Lab/OpenNIR/': Could not resolve proxy: http_proxy=http\n",
      "\u001b[33mWARNING: Discarding git+https://github.com/Georgetown-IR-Lab/OpenNIR. Command errored out with exit status 128: git clone -q https://github.com/Georgetown-IR-Lab/OpenNIR /tmp/pip-req-build-cty2kiwx Check the logs for full command output.\u001b[0m\n",
      "\u001b[31mERROR: Command errored out with exit status 128: git clone -q https://github.com/Georgetown-IR-Lab/OpenNIR /tmp/pip-req-build-cty2kiwx Check the logs for full command output.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# !pip install --upgrade git+https://github.com/Georgetown-IR-Lab/OpenNIR\n",
    "# !pip install --upgrade git+https://github.com/terrierteam/pyterrier_doc2query.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A-nQrpNP5pN7"
   },
   "source": [
    "## Getting started\n",
    "\n",
    "Start PyTerrier as we have in past notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "6FegcyWr5lja"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTerrier 0.7.1 has loaded Terrier 5.6 (built by craigmacdonald on 2021-09-17 13:27)\n",
      "Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.\n"
     ]
    }
   ],
   "source": [
    "import pyterrier as pt\n",
    "if not pt.started():\n",
    "    pt.init(tqdm='notebook')\n",
    "import onir_pt\n",
    "import pyterrier_doc2query\n",
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hPVjr448rIPc"
   },
   "source": [
    "### [TREC-COVID19](https://ir.nist.gov/covidSubmit/) Dataset download\n",
    "\n",
    "The following cell downloads the [TREC-COVID19](https://ir.nist.gov/covidSubmit/) dataset that we will use periodically throughout this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "IJMHFRfArd7O"
   },
   "outputs": [],
   "source": [
    "dataset = pt.datasets.get_dataset('irds:cord19/trec-covid')\n",
    "topics = dataset.get_topics(variant='description')\n",
    "qrels = dataset.get_qrels()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OF3HIPhtrqOH"
   },
   "source": [
    "# Task 1: Build the inverted index for the TREC-COVID19 collection. (2 points)\n",
    "\n",
    "Build the index for the TREC Covid-19 (CORD19) data like we have in past notebooks but without any fancy options (e.g., no positional indexing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "cord19 = pt.datasets.get_dataset('irds:cord19/trec-covid')\n",
    "pt_index_path = './terrier_cord19'\n",
    "\n",
    "if not os.path.exists(pt_index_path + \"/data.properties\"):\n",
    "    # TODO\n",
    "    \n",
    "    # create the index, using the IterDictIndexer indexer \n",
    "    indexer = pt.IterDictIndexer(pt_index_path, overwrite=True)\n",
    "\n",
    "    # we give the dataset get_corpus_iter() directly to the indexer\n",
    "    # while specifying the fields to index and the metadata to record\n",
    "    index_ref = indexer.index(cord19.get_corpus_iter(), fields=('title', 'abstract'))\n",
    "\n",
    "else:\n",
    "    # TODO\n",
    "    \n",
    "    # if you already have the index, use it.\n",
    "    index_ref = pt.IndexRef.of(pt_index_path +\"/data.properties\")\n",
    "\n",
    "index = pt.IndexFactory.of(index_ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rwDams5M7g6c"
   },
   "source": [
    "## Using an untuned Re-rankers\n",
    "\n",
    "This notebook will have you work with a few neural re-ranking methods that you've used in class. We can build them from scratch using `onir_pt.reranker` or load them from pretrained models. The models we load from scratch won't have been trained to do IR (yet), however.\n",
    "\n",
    "And OpenNIR reranking model consists of:\n",
    " - `ranker` (e.g., `drmm`, `knrm`, or `pacrr`). This defines the neural ranking architecture. We discussed the `knrm` approach in class.\n",
    " - `vocab` (e.g., `wordvec_hash`, or `bert`). This defines how text is encoded by the model. This approach makes it easy to swap out different text representations. \n",
    " \n",
    "Let's start with the `knrm` method we discussed in class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "F0O79K2K6fvn"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "config file not found: config\n",
      "[2021-11-24 00:58:54,450][WordvecHashVocab][DEBUG] [starting] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p\n",
      "[2021-11-24 00:59:20,072][WordvecHashVocab][DEBUG] [finished] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p [25.67s]\n"
     ]
    }
   ],
   "source": [
    "knrm = onir_pt.reranker('knrm', 'wordvec_hash', text_field='abstract')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x1avVTpxDORN"
   },
   "source": [
    "Let's look at how well this model work at ranking compared with our default `BatchRetrieve`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "4FWUIN577v1O"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 00:59:35,213][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 01:00:10,650][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/1250 [16ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:00:13,072][onir_pt][DEBUG] [finished] batches: [2.42s] [1250it] [516.71it/s]\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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut.10</th>\n",
       "      <th>P.10</th>\n",
       "      <th>mrt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DPH</td>\n",
       "      <td>0.073907</td>\n",
       "      <td>0.173291</td>\n",
       "      <td>0.641104</td>\n",
       "      <td>0.684</td>\n",
       "      <td>81.301257</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>DPH &gt;&gt; KNRM</td>\n",
       "      <td>0.057971</td>\n",
       "      <td>0.151021</td>\n",
       "      <td>0.372247</td>\n",
       "      <td>0.470</td>\n",
       "      <td>871.655583</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name       map      ndcg  ndcg_cut.10   P.10         mrt\n",
       "0          DPH  0.073907  0.173291     0.641104  0.684   81.301257\n",
       "1  DPH >> KNRM  0.057971  0.151021     0.372247  0.470  871.655583"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "br = pt.BatchRetrieve(index) % 100\n",
    "pipeline = br >> pt.text.get_text(dataset, 'abstract') >> knrm\n",
    "pt.Experiment(\n",
    "    [br, pipeline],\n",
    "    topics,\n",
    "    qrels,\n",
    "    names=['DPH', 'DPH >> KNRM'],\n",
    "    eval_metrics=[\"map\", \"ndcg\", 'ndcg_cut.10', 'P.10', 'mrt']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MhquobQypVNJ"
   },
   "source": [
    "The `knrm` models' performance is lower! The mode doesn't work very well because it hasn't yet been trained for IR; it's using random weights to combine the scores from the similarity matrix--but this is at least a start."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WLv7quA43yAP"
   },
   "source": [
    "## Loading a trained re-ranker\n",
    "\n",
    "You can train re-ranking models in PyTerrier using the `fit` method. Here's an example of how to train the `knrm` model on the MS MARCO dataset, which is a large IR collection.\n",
    "\n",
    "```python\n",
    "# transfer training signals from a medical sample of MS MARCO\n",
    "from sklearn.model_selection import train_test_split\n",
    "train_ds = pt.datasets.get_dataset('irds:msmarco-passage/train/medical')\n",
    "train_topics, valid_topics = train_test_split(train_ds.get_topics(), test_size=50, random_state=42) # split into training and validation sets\n",
    "\n",
    "# Index MS MARCO\n",
    "indexer = pt.index.IterDictIndexer('./terrier_msmarco-passage')\n",
    "tr_index_ref = indexer.index(train_ds.get_corpus_iter(), fields=('text',), meta=('docno',))\n",
    "\n",
    "pipeline = (pt.BatchRetrieve(tr_index_ref) % 100 # get top 100 results\n",
    "            >> pt.text.get_text(train_ds, 'text') # fetch the document text\n",
    "            >> pt.apply.generic(lambda df: df.rename(columns={'text': 'abstract'})) # rename columns\n",
    "            >> knrm) # apply neural re-ranker\n",
    "\n",
    "pipeline.fit(\n",
    "    train_topics,\n",
    "    train_ds.get_qrels(),\n",
    "    valid_topics,\n",
    "    train_ds.get_qrels())\n",
    "```\n",
    "\n",
    "Training deep learning models takes a bit of time (especially for large datasets like MS MARCO), so we've provided a model that's already been trained for you to download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "Yk7FBOgvDa8V"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:00:13,224][onir_pt][INFO] using cached checkpoint: /home/seffendi/data/onir/model_checkpoints/e40f18b6f6e8851dd92f4d1196a43875\n",
      "[2021-11-24 01:00:13,231][WordvecHashVocab][DEBUG] [starting] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p\n",
      "[2021-11-24 01:00:14,711][WordvecHashVocab][DEBUG] [finished] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p [1.48s]\n"
     ]
    }
   ],
   "source": [
    "del knrm # free up the memory before loading a new version of the ranker (helpful for the GPU)\n",
    "knrm = onir_pt.reranker.from_checkpoint('http://jurgens.people.si.umich.edu/ir/knrm.medmarco.tar.gz', text_field='abstract', \n",
    "                                        expected_md5=\"d70b1d4f899690dae51161537e69ed5a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "1BQQKv8lL0Ta"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:00:24,907][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 01:00:25,220][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/1250 [16ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:00:27,355][onir_pt][DEBUG] [finished] batches: [2.13s] [1250it] [586.00it/s]\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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>P.10</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut.10</th>\n",
       "      <th>mrt</th>\n",
       "      <th>map +</th>\n",
       "      <th>map -</th>\n",
       "      <th>map p-value</th>\n",
       "      <th>P.10 +</th>\n",
       "      <th>P.10 -</th>\n",
       "      <th>P.10 p-value</th>\n",
       "      <th>ndcg +</th>\n",
       "      <th>ndcg -</th>\n",
       "      <th>ndcg p-value</th>\n",
       "      <th>ndcg_cut.10 +</th>\n",
       "      <th>ndcg_cut.10 -</th>\n",
       "      <th>ndcg_cut.10 p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DPH</td>\n",
       "      <td>0.073907</td>\n",
       "      <td>0.684</td>\n",
       "      <td>0.173291</td>\n",
       "      <td>0.641104</td>\n",
       "      <td>52.085205</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>DPH &gt;&gt; KNRM</td>\n",
       "      <td>0.068051</td>\n",
       "      <td>0.614</td>\n",
       "      <td>0.164650</td>\n",
       "      <td>0.538218</td>\n",
       "      <td>97.554703</td>\n",
       "      <td>15.0</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0.006176</td>\n",
       "      <td>13.0</td>\n",
       "      <td>24.0</td>\n",
       "      <td>0.023178</td>\n",
       "      <td>13.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>0.00124</td>\n",
       "      <td>15.0</td>\n",
       "      <td>34.0</td>\n",
       "      <td>0.001311</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name       map   P.10      ndcg  ndcg_cut.10        mrt  map +  \\\n",
       "0          DPH  0.073907  0.684  0.173291     0.641104  52.085205    NaN   \n",
       "1  DPH >> KNRM  0.068051  0.614  0.164650     0.538218  97.554703   15.0   \n",
       "\n",
       "   map -  map p-value  P.10 +  P.10 -  P.10 p-value  ndcg +  ndcg -  \\\n",
       "0    NaN          NaN     NaN     NaN           NaN     NaN     NaN   \n",
       "1   35.0     0.006176    13.0    24.0      0.023178    13.0    37.0   \n",
       "\n",
       "   ndcg p-value  ndcg_cut.10 +  ndcg_cut.10 -  ndcg_cut.10 p-value  \n",
       "0           NaN            NaN            NaN                  NaN  \n",
       "1       0.00124           15.0           34.0             0.001311  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline2 = br >> pt.text.get_text(dataset, 'abstract') >> knrm\n",
    "pt.Experiment(\n",
    "    [br, pipeline2],\n",
    "    topics,\n",
    "    qrels,\n",
    "    names=['DPH', 'DPH >> KNRM'],\n",
    "    baseline=0,\n",
    "    eval_metrics=[\"map\", \"ndcg\", 'ndcg_cut.10', 'P.10', 'mrt']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aI1_8O8rtXKK"
   },
   "source": [
    "The tuned performance is a little better than before, but `knrm` still underperforms our first-stage ranking model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "79l1jn0pRQEY"
   },
   "source": [
    "## Reranking with BERT\n",
    "\n",
    "Large language models such as [BERT](https://arxiv.org/abs/1810.04805) are much more powerful neural models that have been shown to be effective for ranking like we discussed in class. \n",
    "\n",
    "Like with `knrm`, we'll start by using BERT for re-ranking with its initial parameters. These parameters have been turned for the masked language modeling (i.e., filling a word in the blank) and predicting the next sentence--but have not been tuned for IR at all."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "-qlXPHqN3iO0"
   },
   "outputs": [],
   "source": [
    "del knrm # clear out memory from KNRM (useful for GPU)\n",
    "bert = onir_pt.reranker('vanilla_transformer', 'bert', text_field='abstract', vocab_config={'train': True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "progrVwaunrn"
   },
   "source": [
    "Let's see how this non-IR trained model does on CORD10 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "PkasovrjQjy0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:00:42,605][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 01:00:42,843][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/1250 [16ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:01:40,385][onir_pt][DEBUG] [finished] batches: [57.54s] [1250it] [21.72it/s]\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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>P.10</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut.10</th>\n",
       "      <th>mrt</th>\n",
       "      <th>map +</th>\n",
       "      <th>map -</th>\n",
       "      <th>map p-value</th>\n",
       "      <th>P.10 +</th>\n",
       "      <th>P.10 -</th>\n",
       "      <th>P.10 p-value</th>\n",
       "      <th>ndcg +</th>\n",
       "      <th>ndcg -</th>\n",
       "      <th>ndcg p-value</th>\n",
       "      <th>ndcg_cut.10 +</th>\n",
       "      <th>ndcg_cut.10 -</th>\n",
       "      <th>ndcg_cut.10 p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DPH</td>\n",
       "      <td>0.073907</td>\n",
       "      <td>0.684</td>\n",
       "      <td>0.173291</td>\n",
       "      <td>0.641104</td>\n",
       "      <td>52.394214</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>DPH &gt;&gt; VBERT</td>\n",
       "      <td>0.060708</td>\n",
       "      <td>0.478</td>\n",
       "      <td>0.152852</td>\n",
       "      <td>0.397904</td>\n",
       "      <td>1204.979833</td>\n",
       "      <td>3.0</td>\n",
       "      <td>47.0</td>\n",
       "      <td>0.000007</td>\n",
       "      <td>6.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>8.344970e-07</td>\n",
       "      <td>6.0</td>\n",
       "      <td>44.0</td>\n",
       "      <td>0.000003</td>\n",
       "      <td>8.0</td>\n",
       "      <td>41.0</td>\n",
       "      <td>5.312061e-09</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name       map   P.10      ndcg  ndcg_cut.10          mrt  map +  \\\n",
       "0           DPH  0.073907  0.684  0.173291     0.641104    52.394214    NaN   \n",
       "1  DPH >> VBERT  0.060708  0.478  0.152852     0.397904  1204.979833    3.0   \n",
       "\n",
       "   map -  map p-value  P.10 +  P.10 -  P.10 p-value  ndcg +  ndcg -  \\\n",
       "0    NaN          NaN     NaN     NaN           NaN     NaN     NaN   \n",
       "1   47.0     0.000007     6.0    37.0  8.344970e-07     6.0    44.0   \n",
       "\n",
       "   ndcg p-value  ndcg_cut.10 +  ndcg_cut.10 -  ndcg_cut.10 p-value  \n",
       "0           NaN            NaN            NaN                  NaN  \n",
       "1      0.000003            8.0           41.0         5.312061e-09  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline3 = br % 100 >> pt.text.get_text(dataset, 'abstract') >> bert\n",
    "pt.Experiment(\n",
    "    [br, pipeline3],\n",
    "    topics,\n",
    "    qrels,\n",
    "    names=['DPH', 'DPH >> VBERT'],\n",
    "    baseline=0,\n",
    "    eval_metrics=[\"map\", \"ndcg\", 'ndcg_cut.10', 'P.10', 'mrt']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wBrfNZ_1u_pD"
   },
   "source": [
    "As we see, although the ERT model is pre-trained for recognizing language, it doesn't do very well at ranking on our benchmark. To get better performance, we'll need to tune for the task of relevance ranking.\n",
    "\n",
    "We can train the model for ranking (as shown above for KNRM) or we can download a trained model. Here, we will use the [SLEDGE](https://arxiv.org/abs/2010.05987) model, which is a BERT model trained on scientific text and tuned on medical queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "VsXQKNyYSXOj"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:01:40,544][onir_pt][INFO] using cached checkpoint: /home/seffendi/data/onir/model_checkpoints/1f28b1381d23b0062a219452a9258032\n"
     ]
    }
   ],
   "source": [
    "bert = onir_pt.reranker.from_checkpoint('http://jurgens.people.si.umich.edu/ir/scibert-medmarco.tar.gz', \n",
    "                                         text_field='abstract', expected_md5=\"854966d0b61543ffffa44cea627ab63b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run another experiment to see how this new model trained for IR does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "dUH-daNJSoNy"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:02:09,782][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 01:02:09,916][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/1250 [16ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 01:03:03,958][onir_pt][DEBUG] [finished] batches: [54.04s] [1250it] [23.13it/s]\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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>P.10</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut.10</th>\n",
       "      <th>mrt</th>\n",
       "      <th>map +</th>\n",
       "      <th>map -</th>\n",
       "      <th>map p-value</th>\n",
       "      <th>P.10 +</th>\n",
       "      <th>P.10 -</th>\n",
       "      <th>P.10 p-value</th>\n",
       "      <th>ndcg +</th>\n",
       "      <th>ndcg -</th>\n",
       "      <th>ndcg p-value</th>\n",
       "      <th>ndcg_cut.10 +</th>\n",
       "      <th>ndcg_cut.10 -</th>\n",
       "      <th>ndcg_cut.10 p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>DPH</td>\n",
       "      <td>0.073907</td>\n",
       "      <td>0.684</td>\n",
       "      <td>0.173291</td>\n",
       "      <td>0.641104</td>\n",
       "      <td>59.672703</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>DPH &gt;&gt; Trained-BERT</td>\n",
       "      <td>0.079776</td>\n",
       "      <td>0.774</td>\n",
       "      <td>0.177709</td>\n",
       "      <td>0.698749</td>\n",
       "      <td>1133.077568</td>\n",
       "      <td>36.0</td>\n",
       "      <td>14.0</td>\n",
       "      <td>0.011571</td>\n",
       "      <td>30.0</td>\n",
       "      <td>13.0</td>\n",
       "      <td>0.014143</td>\n",
       "      <td>29.0</td>\n",
       "      <td>21.0</td>\n",
       "      <td>0.131269</td>\n",
       "      <td>29.0</td>\n",
       "      <td>21.0</td>\n",
       "      <td>0.128519</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  name       map   P.10      ndcg  ndcg_cut.10          mrt  \\\n",
       "0                  DPH  0.073907  0.684  0.173291     0.641104    59.672703   \n",
       "1  DPH >> Trained-BERT  0.079776  0.774  0.177709     0.698749  1133.077568   \n",
       "\n",
       "   map +  map -  map p-value  P.10 +  P.10 -  P.10 p-value  ndcg +  ndcg -  \\\n",
       "0    NaN    NaN          NaN     NaN     NaN           NaN     NaN     NaN   \n",
       "1   36.0   14.0     0.011571    30.0    13.0      0.014143    29.0    21.0   \n",
       "\n",
       "   ndcg p-value  ndcg_cut.10 +  ndcg_cut.10 -  ndcg_cut.10 p-value  \n",
       "0           NaN            NaN            NaN                  NaN  \n",
       "1      0.131269           29.0           21.0             0.128519  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline4 = br % 100 >> pt.text.get_text(dataset, 'abstract') >> bert\n",
    "pt.Experiment(\n",
    "    [br, pipeline4],\n",
    "    topics,\n",
    "    qrels,\n",
    "    names=['DPH', 'DPH >> Trained-BERT'],\n",
    "    baseline=0,\n",
    "    eval_metrics=[\"map\", \"ndcg\", 'ndcg_cut.10', 'P.10', 'mrt']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OtAxDQHyv4ON"
   },
   "source": [
    "Training helped a lot! We're able to improve upon the initial ranking from `BatchRetrieve`. However, from looking at `mrt` we can see that this is pretty slow to run--and this was using a GPU! This performance time underscores the trade-off in using large language models at retrieval time: they may perform better, but could be much slower."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep learning at indexing time: doc2query\n",
    "\n",
    "Instead of using our large language models to rerank, another option is to use them at _indexing time_ to augment our documents. In class, we discussed one such option, doc2query, that augments an inverted index structure by predicting queries that may be used to search for the document, and appending those to the document text.\n",
    "\n",
    "We can use doc2query using the `pyterrier_doc2query` package, which was loaded at the top."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading a pre-trained model\n",
    "\n",
    "We'll start by using a version of the doc2query model released by the authors that is trained on the MS MARCO collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"t5-base.zip\"):\n",
    "  !wget http://jurgens.people.si.umich.edu/ir/t5-base.zip\n",
    "  !unzip t5-base.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can load the model weights by specifying the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Doc2query using cuda\n"
     ]
    }
   ],
   "source": [
    "doc2query = pyterrier_doc2query.Doc2Query('model.ckpt-1004000', batch_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running doc2queries on sample text\n",
    "\n",
    "Let's see what queries it predicts for the sample document that we've made up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The University of Michigan School of Information (UMSI) delivers innovative, elegant and ethical solutions connecting people, information and technology. The school was one of the first iSchools in the nation and is the premier institution studying and using technology to improve human computer interactions.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame([{\"docno\" : \"d1\", \"text\" :\"The University of Michigan School of Information (UMSI) delivers innovative, elegant and ethical solutions connecting people, information and technology. The school was one of the first iSchools in the nation and is the premier institution studying and using technology to improve human computer interactions.\"}])\n",
    "df.iloc[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'what is umsi university of michigan school of information where is umsi'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc2query_df = doc2query(df)\n",
    "doc2query_df.iloc[0].querygen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not too bad, though the questions are somewhat generic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading an index of doc2query documents\n",
    "\n",
    "Let's see how doc2query does on improving the performance in the TREC COVID data. Since indexing with doc2query takes a while (due to needing to run the deep learning models), we've provided an index with the text already added.\n",
    "\n",
    "If you would like to index the collection with doc2query yourself (or use doc2query for your course project), you can use the following code:\n",
    "\n",
    "```python\n",
    "dataset = pt.get_dataset(\"irds:cord19/trec-covid\")\n",
    "indexer = (\n",
    "  pyterrier_doc2query.Doc2Query('model.ckpt-1004000', doc_attr='abstract', batch_size=8, append=True) # aply doc2query on abstracts and append\n",
    "  >> pt.apply.generic(lambda df: df.rename(columns={'abstract': 'text'}) # rename \"abstract\" column to \"text\" for indexing\n",
    "  >> pt.IterDictIndexer(\"./doc2query_index_path\")) # index the expanded documents\n",
    "indexref = indexer.index(dataset.get_corpus_iter())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('doc2query_marco_cord19.zip'):\n",
    "  !wget http://jurgens.people.si.umich.edu/ir/doc2query_marco_cord19.zip\n",
    "  !unzip doc2query_marco_cord19.zip\n",
    "doc2query_indexref = pt.IndexRef.of('./doc2query_index_path/data.properties')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at the results on TREC COVID by first merging the scores with the rankings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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>qid</th>\n",
       "      <th>docid</th>\n",
       "      <th>docno</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>query</th>\n",
       "      <th>title</th>\n",
       "      <th>label</th>\n",
       "      <th>iteration</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>101299</td>\n",
       "      <td>jwmrgy5d</td>\n",
       "      <td>0</td>\n",
       "      <td>8.427298</td>\n",
       "      <td>coronavirus origin</td>\n",
       "      <td>COVID-19 in the heart and the lungs: could we ...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>182167</td>\n",
       "      <td>g8grcy5j</td>\n",
       "      <td>0</td>\n",
       "      <td>13.922648</td>\n",
       "      <td>coronavirus response to weather changes</td>\n",
       "      <td>The Stirling Protocol – Putting the environmen...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>85678</td>\n",
       "      <td>tl30wlpy</td>\n",
       "      <td>0</td>\n",
       "      <td>7.224180</td>\n",
       "      <td>coronavirus immunity</td>\n",
       "      <td>Receptor-dependent coronavirus infection of de...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>145871</td>\n",
       "      <td>l5fxswfz</td>\n",
       "      <td>0</td>\n",
       "      <td>12.773362</td>\n",
       "      <td>how do people die from the coronavirus</td>\n",
       "      <td>Analysis on 54 Mortality Cases of Coronavirus ...</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>180990</td>\n",
       "      <td>3sepefqa</td>\n",
       "      <td>0</td>\n",
       "      <td>12.995980</td>\n",
       "      <td>animal models of covid 19</td>\n",
       "      <td>Current global vaccine and drug efforts agains...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  qid   docid     docno  rank      score  \\\n",
       "0   1  101299  jwmrgy5d     0   8.427298   \n",
       "1   2  182167  g8grcy5j     0  13.922648   \n",
       "2   3   85678  tl30wlpy     0   7.224180   \n",
       "3   4  145871  l5fxswfz     0  12.773362   \n",
       "4   5  180990  3sepefqa     0  12.995980   \n",
       "\n",
       "                                     query  \\\n",
       "0                       coronavirus origin   \n",
       "1  coronavirus response to weather changes   \n",
       "2                     coronavirus immunity   \n",
       "3   how do people die from the coronavirus   \n",
       "4                animal models of covid 19   \n",
       "\n",
       "                                               title  label iteration  \n",
       "0  COVID-19 in the heart and the lungs: could we ...    0.0         5  \n",
       "1  The Stirling Protocol – Putting the environmen...    0.0         4  \n",
       "2  Receptor-dependent coronavirus infection of de...    NaN       NaN  \n",
       "3  Analysis on 54 Mortality Cases of Coronavirus ...    2.0       1.5  \n",
       "4  Current global vaccine and drug efforts agains...    0.0         4  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = pt.get_dataset('irds:cord19/trec-covid')\n",
    "pipeline = pt.BatchRetrieve(doc2query_indexref) % 1 >> pt.text.get_text(dataset, 'title')\n",
    "res = pipeline(dataset.get_topics('title'))\n",
    "res.merge(dataset.get_qrels(), how='left').head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What kind of queries does doc2query generate for the CORD19 documents?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88e00f95121e46a18a12df3426e0f41b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "cord19/trec-covid documents:   0%|          | 0/192509 [15ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l5fxswfz\n",
      "what is the population of coronavirus coronavirus death toll coronavirus worldwide\n",
      "Since the identification of the first case of coronavirus disease 2019 (COVID-19), the global number of confirmed cases as of March 15, 2020, is 156,400, with total death in 5,833 (3.7%) worldwide. Here, we summarize the morality data from February 19 when the first mortality occurred to 0 am, March 10, 2020, in Korea with comparison to other countries. The overall case fatality rate of COVID-19 in Korea was 0.7% as of 0 am, March 10, 2020.\n",
      "3sepefqa\n",
      "what is the cure for copid what is the treatment of covid what is a cure for covid 19?\n",
      "COVID-19 has become one of the biggest health concern, along with huge economic burden. With no clear remedies to treat the disease, doctors are repurposing drugs like chloroquine and remdesivir to treat COVID-19 patients. In parallel, research institutes in collaboration with biotech companies have identified strategies to use viral proteins as vaccine candidates for COVID-19. Although this looks promising, they still need to pass the test of challenge studies in animal models. As various models for SARS-CoV-2 are under testing phase, biotech companies have bypassed animal studies and moved to Phase I clinical trials. In view of the present outbreak, this looks a justified approach, but the problem is that in the absence of animal studies, we can never predict the outcomes in humans. Since animal models are critical for vaccine development and SARS-CoV-2 has different transmission dynamics, in this review we compare different animal models of SARS-CoV-2 with humans for their pathogenic, immune response and transmission dynamics that make them ideal models for vaccine testing for COVID-19. Another issue of using animal model is the ethics of using animals for research; thus, we also discuss the pros and cons of using animals for vaccine development studies.\n"
     ]
    }
   ],
   "source": [
    "df = pd.DataFrame(doc for doc in dataset.get_corpus_iter() if doc['docno'] in ('3sepefqa', 'l5fxswfz'))\n",
    "df = df.rename(columns={'abstract': 'text'})\n",
    "doc2query_df = doc2query(df)\n",
    "for querygen, docno, text in zip(doc2query_df['querygen'], doc2query_df['docno'], df['text']):\n",
    "    print(docno)\n",
    "    print(querygen)\n",
    "    print(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the effects of doc2query\n",
    "\n",
    "Here, we'll change our evaluation setup a bit from what we did before. Rather than compare two models for the same index, we'll instead compare the same model (BM25) with two different ways of indexing (two indices)! Our baseline will be an index of CORD19 without the doc2query additions.\n",
    "\n",
    "Let's load a copy of the CORD19 index that we used earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexref = pt.IndexRef.of('./terrier_cord19/data.properties')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 2: Write the `Experiment` to compare indices (3 points)\n",
    "Run an `Experiment` using a `BM25` ranker that compares the indices `indexref` and `doc2query_indexref`. Compare your models using MAP, NDCG, and NDCG@10. Note that our doc2query model was trained on MS MARCO, which isn't the same kind of collection as CORD19, so this performance tells us how well that model can transfer to a new setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm25_indexref = pt.BatchRetrieve(indexref, wmodel=\"BM25\")\n",
    "bm25_doc2query_indexref = pt.BatchRetrieve(doc2query_indexref, wmodel=\"BM25\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut_10</th>\n",
       "      <th>map +</th>\n",
       "      <th>map -</th>\n",
       "      <th>map p-value</th>\n",
       "      <th>ndcg +</th>\n",
       "      <th>ndcg -</th>\n",
       "      <th>ndcg p-value</th>\n",
       "      <th>ndcg_cut_10 +</th>\n",
       "      <th>ndcg_cut_10 -</th>\n",
       "      <th>ndcg_cut_10 p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>BR(BM25)</td>\n",
       "      <td>0.222260</td>\n",
       "      <td>0.451164</td>\n",
       "      <td>0.623256</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>BR(BM25)</td>\n",
       "      <td>0.196293</td>\n",
       "      <td>0.412673</td>\n",
       "      <td>0.623652</td>\n",
       "      <td>15.0</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0.000634</td>\n",
       "      <td>11.0</td>\n",
       "      <td>39.0</td>\n",
       "      <td>0.000019</td>\n",
       "      <td>26.0</td>\n",
       "      <td>23.0</td>\n",
       "      <td>0.98621</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       name       map      ndcg  ndcg_cut_10  map +  map -  map p-value  \\\n",
       "0  BR(BM25)  0.222260  0.451164     0.623256    NaN    NaN          NaN   \n",
       "1  BR(BM25)  0.196293  0.412673     0.623652   15.0   35.0     0.000634   \n",
       "\n",
       "   ndcg +  ndcg -  ndcg p-value  ndcg_cut_10 +  ndcg_cut_10 -  \\\n",
       "0     NaN     NaN           NaN            NaN            NaN   \n",
       "1    11.0    39.0      0.000019           26.0           23.0   \n",
       "\n",
       "   ndcg_cut_10 p-value  \n",
       "0                  NaN  \n",
       "1              0.98621  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pt.Experiment([bm25_indexref, bm25_doc2query_indexref],\n",
    "            topics,\n",
    "            qrels,\n",
    "            baseline=0,\n",
    "            eval_metrics=[\"map\", \"ndcg\", \"ndcg_cut_10\"]\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Task 3: Train a new model! (25 points)\n",
    "\n",
    "All of the prior exercises have had you working with either off-the-shelf models (not trained for IR) or models that someone else has trained for you. To give you a sense of how to train a model, your primary task in this notebook is to train a simple `knrm` model, which should be relatively efficient to train on a GPU. \n",
    "\n",
    "To keep things simple, we'll use the same setup for CORD19 that we did in Part 2 (30 queries in train, 5 in dev, 15 queries in test) which is still relatively small for training a deep learning model but will get you started on the process. \n",
    "\n",
    "Your tasks are the following:\n",
    "- Load the CORD19 dataset and split it into train, dev, and test\n",
    "- Create a new `knrm` ranker and a pipeline that uses it\n",
    "- Run an `Experiment` comparing four models:\n",
    "  - a default `BatchRetrieve`, filtering to the top 100 results\n",
    "  - BM25, filtering to the top 100 results\n",
    "  - a pipeline that feeds the top 100 results of the default `BatchRetrieve` to your `knrm` model\n",
    "  - a pipeline that feeds the top 100 results of BM25 to your `knrm` model\n",
    "  \n",
    "Your `Experiment` should evaluate models using MAP, NDCG, NDCG@10, Precision@10, and Mean Response Time.\n",
    "  \n",
    "We expect to see the `Experiment`'s results in the final cell. You are, of course, welcome to try training any of the fancier models to see how they do as well!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RANK_CUTOFF = 4\n",
    "SEED=42\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "tr_va_topics, test_topics = train_test_split(topics, test_size=30, random_state=SEED)\n",
    "train_topics, valid_topics =  train_test_split(tr_va_topics, test_size=5, random_state=SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 02:20:28,337][WordvecHashVocab][DEBUG] [starting] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p\n",
      "[2021-11-24 02:20:29,769][WordvecHashVocab][DEBUG] [finished] reading cached at /home/seffendi/data/onir/vocab/wordvec_hash/fasttext-wiki-news-300d-1M.p [1.43s]\n"
     ]
    }
   ],
   "source": [
    "del knrm # clear out memory from KNRM (useful for GPU)\n",
    "knrm = onir_pt.reranker('knrm', 'wordvec_hash', text_field='abstract')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "br = pt.BatchRetrieve(index)% 100\n",
    "bm25 = pt.BatchRetrieve(index, wmodel=\"BM25\")% 100\n",
    "br_knrm = (br% 100) >> pt.text.get_text(dataset, 'abstract') >> knrm\n",
    "bm25_knrm = (bm25% 100) >> pt.text.get_text(dataset, 'abstract') >> knrm\n",
    "# bm25_indexref = pt.BatchRetrieve(indexref, wmodel=\"BM25\")\n",
    "# bm25_doc2query_indexref = pt.BatchRetrieve(doc2query_indexref, wmodel=\"BM25\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 02:51:12,297][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 02:51:12,299][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/750 [15ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 02:51:13,593][onir_pt][DEBUG] [finished] batches: [1.29s] [750it] [580.30it/s]\n",
      "[2021-11-24 02:51:15,150][onir_pt][DEBUG] using GPU (deterministic)\n",
      "[2021-11-24 02:51:15,152][onir_pt][DEBUG] [starting] batches\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "batches:   0%|          | 0/750 [15ms<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-11-24 02:51:16,371][onir_pt][DEBUG] [finished] batches: [1.22s] [750it] [615.77it/s]\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>name</th>\n",
       "      <th>map</th>\n",
       "      <th>P.10</th>\n",
       "      <th>ndcg</th>\n",
       "      <th>ndcg_cut_10</th>\n",
       "      <th>mrt</th>\n",
       "      <th>map +</th>\n",
       "      <th>map -</th>\n",
       "      <th>map p-value</th>\n",
       "      <th>P.10 +</th>\n",
       "      <th>P.10 -</th>\n",
       "      <th>P.10 p-value</th>\n",
       "      <th>ndcg +</th>\n",
       "      <th>ndcg -</th>\n",
       "      <th>ndcg p-value</th>\n",
       "      <th>ndcg_cut_10 +</th>\n",
       "      <th>ndcg_cut_10 -</th>\n",
       "      <th>ndcg_cut_10 p-value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>RankCutoff(BR(DPH), 100)</td>\n",
       "      <td>0.067066</td>\n",
       "      <td>0.650000</td>\n",
       "      <td>0.161810</td>\n",
       "      <td>0.618921</td>\n",
       "      <td>46.670864</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>RankCutoff(BR(BM25), 100)</td>\n",
       "      <td>0.067096</td>\n",
       "      <td>0.633333</td>\n",
       "      <td>0.162287</td>\n",
       "      <td>0.616385</td>\n",
       "      <td>52.282603</td>\n",
       "      <td>10.0</td>\n",
       "      <td>20.0</td>\n",
       "      <td>9.938921e-01</td>\n",
       "      <td>8.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>4.931637e-01</td>\n",
       "      <td>11.0</td>\n",
       "      <td>19.0</td>\n",
       "      <td>9.333127e-01</td>\n",
       "      <td>14.0</td>\n",
       "      <td>14.0</td>\n",
       "      <td>9.061608e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Compose(Compose(RankCutoff(RankCutoff(BR(DPH),...</td>\n",
       "      <td>0.051141</td>\n",
       "      <td>0.410000</td>\n",
       "      <td>0.137866</td>\n",
       "      <td>0.316237</td>\n",
       "      <td>90.812105</td>\n",
       "      <td>1.0</td>\n",
       "      <td>29.0</td>\n",
       "      <td>5.676918e-08</td>\n",
       "      <td>3.0</td>\n",
       "      <td>24.0</td>\n",
       "      <td>2.247270e-06</td>\n",
       "      <td>0.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>8.700451e-09</td>\n",
       "      <td>1.0</td>\n",
       "      <td>29.0</td>\n",
       "      <td>1.303282e-09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Compose(Compose(RankCutoff(RankCutoff(BR(BM25)...</td>\n",
       "      <td>0.050528</td>\n",
       "      <td>0.343333</td>\n",
       "      <td>0.136713</td>\n",
       "      <td>0.266706</td>\n",
       "      <td>90.037235</td>\n",
       "      <td>4.0</td>\n",
       "      <td>26.0</td>\n",
       "      <td>2.201511e-04</td>\n",
       "      <td>2.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>3.121184e-07</td>\n",
       "      <td>6.0</td>\n",
       "      <td>24.0</td>\n",
       "      <td>5.259529e-04</td>\n",
       "      <td>3.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>4.090351e-09</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                name       map      P.10  \\\n",
       "0                           RankCutoff(BR(DPH), 100)  0.067066  0.650000   \n",
       "1                          RankCutoff(BR(BM25), 100)  0.067096  0.633333   \n",
       "2  Compose(Compose(RankCutoff(RankCutoff(BR(DPH),...  0.051141  0.410000   \n",
       "3  Compose(Compose(RankCutoff(RankCutoff(BR(BM25)...  0.050528  0.343333   \n",
       "\n",
       "       ndcg  ndcg_cut_10        mrt  map +  map -   map p-value  P.10 +  \\\n",
       "0  0.161810     0.618921  46.670864    NaN    NaN           NaN     NaN   \n",
       "1  0.162287     0.616385  52.282603   10.0   20.0  9.938921e-01     8.0   \n",
       "2  0.137866     0.316237  90.812105    1.0   29.0  5.676918e-08     3.0   \n",
       "3  0.136713     0.266706  90.037235    4.0   26.0  2.201511e-04     2.0   \n",
       "\n",
       "   P.10 -  P.10 p-value  ndcg +  ndcg -  ndcg p-value  ndcg_cut_10 +  \\\n",
       "0     NaN           NaN     NaN     NaN           NaN            NaN   \n",
       "1    12.0  4.931637e-01    11.0    19.0  9.333127e-01           14.0   \n",
       "2    24.0  2.247270e-06     0.0    30.0  8.700451e-09            1.0   \n",
       "3    25.0  3.121184e-07     6.0    24.0  5.259529e-04            3.0   \n",
       "\n",
       "   ndcg_cut_10 -  ndcg_cut_10 p-value  \n",
       "0            NaN                  NaN  \n",
       "1           14.0         9.061608e-01  \n",
       "2           29.0         1.303282e-09  \n",
       "3           27.0         4.090351e-09  "
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pt.Experiment([br, bm25, br_knrm, bm25_knrm],\n",
    "            test_topics,\n",
    "            qrels,\n",
    "            baseline=0,\n",
    "            eval_metrics=[\"map\", \"ndcg\", \"ndcg_cut_10\", \"P.10\", \"mrt\"]\n",
    "            )"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "neural ir.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}