notebooks/pruned-indices.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "racial-consultancy",
   "metadata": {},
   "source": [
    "# Table 1 - Full Index Runs\n",
    "\n",
    "This notebook reproduces the ColBERT runs in Table 1 which uses the newly generated indices.\n",
    "BM25 figures are obtained from the [Terrier Data Repository](http://data.terrier.org/msmarco_passage.dataset.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "julian-factory",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PyTerrier 0.8.0 has loaded Terrier 5.6 (built by craigmacdonald on 2021-09-17 13:27)\n",
      "\n",
      "No etc/terrier.properties, using terrier.default.properties for bootstrap configuration.\n"
     ]
    }
   ],
   "source": [
    "import pyterrier as pt\n",
    "pt.init(tqdm='notebook')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "efficient-arrival",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.10.0+cu113'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "amber-heritage",
   "metadata": {},
   "source": [
    "## load the stops removed index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "historical-charles",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing ColBERT: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing ColBERT from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing ColBERT from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of ColBERT were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['linear.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 15, 21:40:48] #> Loading model checkpoint.\n",
      "[Feb 15, 21:40:48] #> Loading checkpoint http://www.dcs.gla.ac.uk/~craigm/ecir2021-tutorial/colbert_model_checkpoint.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/colbert38/lib/python3.8/site-packages/torch/hub.py:513: UserWarning: Falling back to the old format < 1.6. This support will be deprecated in favor of default zipfile format introduced in 1.6. Please redo torch.save() to save it in the new zipfile format.\n",
      "  warnings.warn('Falling back to the old format < 1.6. This support will be '\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 15, 21:40:59] #> checkpoint['epoch'] = 0\n",
      "[Feb 15, 21:40:59] #> checkpoint['batch'] = 44500\n"
     ]
    }
   ],
   "source": [
    "import pyterrier_colbert\n",
    "from pyterrier_colbert.ranking import *\n",
    "checkpoint=\"http://www.dcs.gla.ac.uk/~craigm/ecir2021-tutorial/colbert_model_checkpoint.zip\"\n",
    "index=(\"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/\",\"stops\")\n",
    "factory = pyterrier_colbert.ranking.ColBERTFactory(checkpoint, *index, )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "protected-television",
   "metadata": {},
   "outputs": [],
   "source": [
    "factory.faiss_index_on_gpu = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "union-pilot",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading reranking index, memtype=mem\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfb4864749bb429f9f1ef8f174208bec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading index shards to memory:   0%|          | 0/30 [00:00<?, ?shard/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[15449104,\n",
       " 15627119,\n",
       " 15646684,\n",
       " 15674365,\n",
       " 15689361,\n",
       " 15670286,\n",
       " 15757845,\n",
       " 15532714,\n",
       " 15453631,\n",
       " 15518354,\n",
       " 15539412,\n",
       " 15540951,\n",
       " 15621620,\n",
       " 15633490,\n",
       " 15590560,\n",
       " 15631383,\n",
       " 15648003,\n",
       " 15692722,\n",
       " 16015586,\n",
       " 16162576,\n",
       " 17950180,\n",
       " 16334153,\n",
       " 15961213,\n",
       " 15340614,\n",
       " 15133805,\n",
       " 15173588,\n",
       " 15015812,\n",
       " 15057601,\n",
       " 15250116,\n",
       " 7106533]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[x.mmap.shape[0] for x in factory._rrm().part_mmap]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "noted-egypt",
   "metadata": {},
   "source": [
    "## eval setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "tribal-technology",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyterrier.measures import *\n",
    "MEASURES = [AP(rel=2)@1000, nDCG@10,nDCG@20,nDCG@100, RR(rel=2)@10,RR(rel=2),RR@10, RR ,R(rel=2)@1000,\"mrt\", \"num_ret\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "reduced-audio",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "def do_exp(pipelines, topics='2019', names=None, baseline=None, save_dir=None) -> pd.DataFrame:\n",
    "    TQ = None\n",
    "    topics_map = { '2019' : 'test-2019', '2020' : 'test-2020', 'dev' : 'dev.small' }\n",
    "    TQ = (\n",
    "       pt.get_dataset(\"msmarco_passage\").get_topics(topics_map[topics]), \n",
    "       pt.get_dataset(\"msmarco_passage\").get_qrels(topics_map[topics])\n",
    "    )\n",
    "    if save_dir is not None:\n",
    "        save_dir = save_dir % str(topics)\n",
    "        os.makedirs(save_dir, exist_ok=True)\n",
    "    df = pt.Experiment(\n",
    "        pipelines, \n",
    "        *TQ,\n",
    "        MEASURES,\n",
    "        batch_size=20,\n",
    "        filter_by_qrels=True,\n",
    "        verbose=True,\n",
    "        names=names,\n",
    "        save_dir=save_dir\n",
    "    )\n",
    "    if save_dir:\n",
    "        df.to_pickle(save_dir + \"/results_df.pkl\")\n",
    "        \n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accepting-chess",
   "metadata": {},
   "source": [
    "# load in results from the Original index \n",
    "\n",
    "This corresponds to the set of documents to be re-ranked in the \"Orig. Faiss\" rows of Table 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "forward-likelihood",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_df = pd.concat([\n",
    "    pt.io.read_results(\"/nfs/staticPruning/results/repro/baseline_%s/approx_baseline.res.gz\" %t) for t in ['2019', '2020', 'dev']\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vital-match",
   "metadata": {},
   "source": [
    "## Perform runs using the Stopwords Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "received-canadian",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 15, 21:44:34] #> Loading the FAISS index from /nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/stops/ivfpq.faiss ..\n",
      "[Feb 15, 21:44:58] #> Building the emb2pid mapping..\n",
      "[Feb 15, 21:45:23] len(self.emb2pid) = 461419381\n"
     ]
    }
   ],
   "source": [
    "baselineApprox = (\n",
    "    (factory.ann_retrieve_score() % 1000) \n",
    "    >> factory.index_scorer(query_encoded=True)\n",
    ")\n",
    "\n",
    "def _add_docno(df):\n",
    "    df = df.copy()\n",
    "    df[\"docid\"] = df.docno.astype(int)\n",
    "    return df\n",
    "\n",
    "orig_faiss = (\n",
    "    pt.Transformer.from_df(src_df) \n",
    "    >> pt.apply.generic(_add_docno)\n",
    "    >> factory.index_scorer()\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abstract-mission",
   "metadata": {},
   "source": [
    "test that one topic works using the orig_faiss pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "criminal-surface",
   "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>docno</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>1037798</td>\n",
       "      <td>1597</td>\n",
       "      <td>13.247594</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        qid docno      score\n",
       "45  1037798  1597  13.247594"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_topic = pt.get_dataset(\"msmarco_passage\").get_topics('test-2019')\n",
    "test_topic = test_topic[test_topic.qid == \"1037798\"]\n",
    "\n",
    "\n",
    "orig_faiss(test_topic).head(1)[['qid', 'docno', 'score']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "narrative-district",
   "metadata": {},
   "source": [
    "calculate the results for the stopwords index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "elect-telephone",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2019\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b57b5fd003284699bdf97eece4cec7ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/6 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "      <th>mrt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.399892</td>\n",
       "      <td>0.686446</td>\n",
       "      <td>0.638233</td>\n",
       "      <td>0.572319</td>\n",
       "      <td>0.825581</td>\n",
       "      <td>0.82591</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.672651</td>\n",
       "      <td>43000.0</td>\n",
       "      <td>225.552680</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.398735</td>\n",
       "      <td>0.686294</td>\n",
       "      <td>0.638128</td>\n",
       "      <td>0.571983</td>\n",
       "      <td>0.830233</td>\n",
       "      <td>0.83041</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>43000.0</td>\n",
       "      <td>155.356856</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline        0.399892  0.686446  0.638233  0.572319   \n",
       "1       orig_faiss        0.398735  0.686294  0.638128  0.571983   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000  num_ret  \\\n",
       "0      0.825581    0.82591  0.928295  0.928295       0.672651  43000.0   \n",
       "1      0.830233    0.83041  0.928295  0.928295       0.695360  43000.0   \n",
       "\n",
       "          mrt  \n",
       "0  225.552680  \n",
       "1  155.356856  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ac363b56f8b4818ac3db265028c7804",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/6 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.440025</td>\n",
       "      <td>0.681037</td>\n",
       "      <td>0.645387</td>\n",
       "      <td>0.591887</td>\n",
       "      <td>0.859899</td>\n",
       "      <td>0.861323</td>\n",
       "      <td>0.916226</td>\n",
       "      <td>0.917650</td>\n",
       "      <td>0.729086</td>\n",
       "      <td>54000.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.452856</td>\n",
       "      <td>0.687692</td>\n",
       "      <td>0.653272</td>\n",
       "      <td>0.604778</td>\n",
       "      <td>0.860340</td>\n",
       "      <td>0.861764</td>\n",
       "      <td>0.916667</td>\n",
       "      <td>0.918091</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>54000.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline        0.440025  0.681037  0.645387  0.591887   \n",
       "1       orig_faiss        0.452856  0.687692  0.653272  0.604778   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000  num_ret  \n",
       "0      0.859899   0.861323  0.916226  0.917650       0.729086  54000.0  \n",
       "1      0.860340   0.861764  0.916667  0.918091       0.750889  54000.0  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dev\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc100b4e66e941fb8e1b7a8a913bd384",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/698 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.386964</td>\n",
       "      <td>0.411081</td>\n",
       "      <td>0.438271</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.328412</td>\n",
       "      <td>0.339039</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6980000.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.391132</td>\n",
       "      <td>0.416078</td>\n",
       "      <td>0.444745</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.331263</td>\n",
       "      <td>0.342262</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6980000.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline             0.0  0.386964  0.411081  0.438271   \n",
       "1       orig_faiss             0.0  0.391132  0.416078  0.444745   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000    num_ret  \n",
       "0           0.0        0.0  0.328412  0.339039            0.0  6980000.0  \n",
       "1           0.0        0.0  0.331263  0.342262            0.0  6980000.0  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_res = []\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "for t in ['2019', '2020', 'dev']:\n",
    "    print(t)\n",
    "    df = do_exp(\n",
    "        [baselineApprox ,orig_faiss],\n",
    "        names=['approx_baseline', 'orig_faiss'],\n",
    "        topics=t,\n",
    "        save_dir='/nfs/staticPruning/results/repro/stopindex_%s/'\n",
    "    ) \n",
    "    display(df)\n",
    "    all_res.append(df)\n",
    "\n",
    "#pd.concat(all_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "studied-warrant",
   "metadata": {},
   "source": [
    "# Perform runs using the IDF_uniform_100 index.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "raised-imaging",
   "metadata": {},
   "outputs": [],
   "source": [
    "#this may not free sufficient memory to run the following cells.\n",
    "del(factory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "median-nelson",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing ColBERT: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing ColBERT from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing ColBERT from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of ColBERT were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['linear.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 15, 22:46:59] #> Loading model checkpoint.\n",
      "[Feb 15, 22:46:59] #> Loading checkpoint http://www.dcs.gla.ac.uk/~craigm/ecir2021-tutorial/colbert_model_checkpoint.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/colbert38/lib/python3.8/site-packages/torch/hub.py:513: UserWarning: Falling back to the old format < 1.6. This support will be deprecated in favor of default zipfile format introduced in 1.6. Please redo torch.save() to save it in the new zipfile format.\n",
      "  warnings.warn('Falling back to the old format < 1.6. This support will be '\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 15, 22:47:08] #> checkpoint['epoch'] = 0\n",
      "[Feb 15, 22:47:08] #> checkpoint['batch'] = 44500\n",
      "Loading reranking index, memtype=mem\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03512cd6019a48ffbf371d6cab98d897",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading index shards to memory:   0%|          | 0/30 [00:00<?, ?shard/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[12785415,\n",
       " 12940023,\n",
       " 12959797,\n",
       " 12993264,\n",
       " 12996879,\n",
       " 12973964,\n",
       " 13051655,\n",
       " 12899739,\n",
       " 12865110,\n",
       " 12914274,\n",
       " 12928676,\n",
       " 12925241,\n",
       " 12998281,\n",
       " 13006198,\n",
       " 12967598,\n",
       " 13001830,\n",
       " 13012426,\n",
       " 13041338,\n",
       " 13231898,\n",
       " 13345915,\n",
       " 14870935,\n",
       " 13493776,\n",
       " 13217827,\n",
       " 12717560,\n",
       " 12557106,\n",
       " 12597488,\n",
       " 12464258,\n",
       " 12498494,\n",
       " 12656402,\n",
       " 5918433]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pyterrier_colbert\n",
    "from pyterrier_colbert.ranking import *\n",
    "checkpoint=\"http://www.dcs.gla.ac.uk/~craigm/ecir2021-tutorial/colbert_model_checkpoint.zip\"\n",
    "index=(\"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/\",\"idf_uniform_100\")\n",
    "factory = pyterrier_colbert.ranking.ColBERTFactory(checkpoint, *index, )\n",
    "factory.faiss_index_on_gpu = False\n",
    "[x.mmap.shape[0] for x in factory._rrm().part_mmap]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "sunset-fruit",
   "metadata": {},
   "outputs": [],
   "source": [
    "baselineApprox = (\n",
    "    (factory.ann_retrieve_score() % 1000) \n",
    "    >> factory.index_scorer(query_encoded=True)\n",
    ")\n",
    "\n",
    "def _add_docno(df):\n",
    "    df = df.copy()\n",
    "    df[\"docid\"] = df.docno.astype(int)\n",
    "    return df\n",
    "\n",
    "orig_faiss = (\n",
    "    pt.Transformer.from_df(src_df) \n",
    "    >> pt.apply.generic(_add_docno)\n",
    "    >> factory.index_scorer()\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "prepared-serbia",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2019\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "75dd12b5fa744ec188ccd68c6b271c76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/6 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "      <th>mrt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.405467</td>\n",
       "      <td>0.687114</td>\n",
       "      <td>0.650281</td>\n",
       "      <td>0.576059</td>\n",
       "      <td>0.844961</td>\n",
       "      <td>0.845303</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.673932</td>\n",
       "      <td>43000.0</td>\n",
       "      <td>230.568591</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.402628</td>\n",
       "      <td>0.688892</td>\n",
       "      <td>0.650414</td>\n",
       "      <td>0.579728</td>\n",
       "      <td>0.845736</td>\n",
       "      <td>0.845922</td>\n",
       "      <td>0.932171</td>\n",
       "      <td>0.932171</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>43000.0</td>\n",
       "      <td>163.985808</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline        0.405467  0.687114  0.650281  0.576059   \n",
       "1       orig_faiss        0.402628  0.688892  0.650414  0.579728   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000  num_ret  \\\n",
       "0      0.844961   0.845303  0.936047  0.936047       0.673932  43000.0   \n",
       "1      0.845736   0.845922  0.932171  0.932171       0.695360  43000.0   \n",
       "\n",
       "          mrt  \n",
       "0  230.568591  \n",
       "1  163.985808  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0e0d894b2a2416f8c0ff80dd15f812f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/6 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.447345</td>\n",
       "      <td>0.682683</td>\n",
       "      <td>0.653719</td>\n",
       "      <td>0.605753</td>\n",
       "      <td>0.847399</td>\n",
       "      <td>0.850625</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.926066</td>\n",
       "      <td>0.735760</td>\n",
       "      <td>54000.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.457357</td>\n",
       "      <td>0.687942</td>\n",
       "      <td>0.658446</td>\n",
       "      <td>0.612369</td>\n",
       "      <td>0.850044</td>\n",
       "      <td>0.851728</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.927169</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>54000.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline        0.447345  0.682683  0.653719  0.605753   \n",
       "1       orig_faiss        0.457357  0.687942  0.658446  0.612369   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000  num_ret  \n",
       "0      0.847399   0.850625  0.922840  0.926066       0.735760  54000.0  \n",
       "1      0.850044   0.851728  0.925485  0.927169       0.750889  54000.0  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dev\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e8a94ba7871b4b259542005d9ecb38ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/698 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>AP(rel=2)@1000</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@20</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>RR(rel=2)@10</th>\n",
       "      <th>RR(rel=2)</th>\n",
       "      <th>RR@10</th>\n",
       "      <th>RR</th>\n",
       "      <th>R(rel=2)@1000</th>\n",
       "      <th>num_ret</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.391636</td>\n",
       "      <td>0.416701</td>\n",
       "      <td>0.443613</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.333394</td>\n",
       "      <td>0.344127</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6980000.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>orig_faiss</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.394758</td>\n",
       "      <td>0.420366</td>\n",
       "      <td>0.448454</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.335727</td>\n",
       "      <td>0.346748</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6980000.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  \\\n",
       "0  approx_baseline             0.0  0.391636  0.416701  0.443613   \n",
       "1       orig_faiss             0.0  0.394758  0.420366  0.448454   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000    num_ret  \n",
       "0           0.0        0.0  0.333394  0.344127            0.0  6980000.0  \n",
       "1           0.0        0.0  0.335727  0.346748            0.0  6980000.0  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_resIdf = []\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "for t in ['2019', '2020', 'dev']:\n",
    "    print(t)\n",
    "    df = do_exp(\n",
    "        [baselineApprox,orig_faiss ],\n",
    "        names=['approx_baseline', 'orig_faiss'],\n",
    "        topics=t,\n",
    "        save_dir='/nfs/staticPruning/results/repro/idf_uniform_100_index_%s/'\n",
    "    ) \n",
    "    display(df)\n",
    "    all_resIdf.append(df)\n",
    "\n",
    "#pd.concat(all_resIdf)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "colbert38",
   "language": "python",
   "name": "colbert38"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}