cikm2021-demos.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "polar-increase",
   "metadata": {},
   "source": [
    "# CIKM 2021 ColBERT Papers\n",
    "\n",
    "This notebook demonstrates the use of techniques proposed in our CIKM 2021 papers:\n",
    "\n",
    " - [Macdonald21a]: On Approximate Nearest Neighbour Selection for Multi-Stage Dense Retrieval. Craig Macdonald and Nicola Tonellotto. In Proceedings of CIKM 2021. https://arxiv.org/abs/2108.11480 \n",
    " - [Tonellotto21]: Query Embedding Pruning for Dense Retrieval Nicola Tonellotto and Craig Macdonald. In Proceedings of CIKM 2021. https://arxiv.org/abs/2108.10341"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "oriented-greene",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "Install pyt_colbert installs PyTerrier too. You also need to have [FAISS installed](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "following-reading",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/terrierteam/pyterrier_colbert.git@cikm2021\n",
      "  Cloning https://github.com/terrierteam/pyterrier_colbert.git (to revision cikm2021) to /tmp/pip-req-build-lufimuww\n",
      "  Running command git clone -q https://github.com/terrierteam/pyterrier_colbert.git /tmp/pip-req-build-lufimuww\n",
      "  Running command git checkout -b cikm2021 --track origin/cikm2021\n",
      "  Switched to a new branch 'cikm2021'\n",
      "  Branch 'cikm2021' set up to track remote branch 'cikm2021' from 'origin'.\n",
      "Building wheels for collected packages: pyterrier-colbert\n",
      "  Building wheel for pyterrier-colbert (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pyterrier-colbert: filename=pyterrier_colbert-0.0.1-py3-none-any.whl size=21094 sha256=e0d16832b0c075205ff59cb32875f1cede099f0e79ba1ad0ff095d3789c4c958\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-83tdy0gu/wheels/c3/e5/5b/24fce6cf44d216004312001eeeb43826e86b074f7c5222ad90\n",
      "Successfully built pyterrier-colbert\n",
      "Installing collected packages: pyterrier-colbert\n",
      "  Attempting uninstall: pyterrier-colbert\n",
      "    Found existing installation: pyterrier-colbert 0.0.1\n",
      "    Uninstalling pyterrier-colbert-0.0.1:\n",
      "      Successfully uninstalled pyterrier-colbert-0.0.1\n",
      "Successfully installed pyterrier-colbert-0.0.1\n"
     ]
    }
   ],
   "source": [
    "!/opt/conda/envs/colbert_cikm2021/bin/pip install --force-reinstall --no-deps git+https://github.com/terrierteam/pyterrier_colbert.git@cikm2021"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "female-status",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTerrier 0.7.0 has loaded Terrier 5.6 (built by craigmacdonald on 2021-09-17 13:27)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No etc/terrier.properties, using terrier.default.properties for bootstrap configuration.\n"
     ]
    }
   ],
   "source": [
    "import pyterrier as pt\n",
    "pt.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "assumed-belgium",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We have an existing index for the MSMARCO v1 Passage corpus, previously indexed using pyt_colbert (this adds the tokenids file, which is needed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "manual-finish",
   "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": [
      "[Sep 30, 09:04:55] #> Loading model checkpoint.\n",
      "[Sep 30, 09:04:55] #> Loading checkpoint http://www.dcs.gla.ac.uk/~craigm/colbert.dnn.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/colbert_cikm2021/lib/python3.8/site-packages/torch/hub.py:498: 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": [
      "[Sep 30, 09:05:04] #> checkpoint['epoch'] = 0\n",
      "[Sep 30, 09:05:04] #> checkpoint['batch'] = 44500\n"
     ]
    }
   ],
   "source": [
    "from pyterrier_colbert.ranking import ColBERTFactory\n",
    "\n",
    "factory = ColBERTFactory(\n",
    "    \"http://www.dcs.gla.ac.uk/~craigm/colbert.dnn.zip\",\n",
    "    \"/nfs/indices/colbert_passage/\",\"index_name3\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "active-patrick",
   "metadata": {},
   "source": [
    "## Baseline\n",
    "\n",
    "This is the default ColBERT dense retrieval setting - a set ANN retrieval from the FAISS index, followed an exact scoring using the large ColBERT index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "conceptual-karaoke",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Sep 30, 09:05:05] #> Loading the FAISS index from /nfs/indices/colbert_passage/index_name3/ivfpq.faiss ..\n",
      "[Sep 30, 09:05:33] #> Building the emb2pid mapping..\n",
      "[Sep 30, 09:06:04] len(self.emb2pid) = 687989391\n",
      "Loading reranking index, memtype=mem\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading index shards to memory: 100%|██████████| 24/24 [03:26<00:00,  8.59s/shard]\n"
     ]
    }
   ],
   "source": [
    "e2e = factory.end_to_end()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compound-credits",
   "metadata": {},
   "source": [
    "## CIKM pipelines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "critical-harassment",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Sep 30, 09:10:01] #> Building the emb2tid mapping..\n",
      "687989391\n",
      "Computing collection frequencies\n",
      "Done\n",
      "Loading doclens\n"
     ]
    }
   ],
   "source": [
    "import pyterrier_colbert.pruning\n",
    "\n",
    "#CIKM 2021 Approximate Scoring paper: only retrieve 200 candidates for exact re-ranking\n",
    "ann_pipe = (factory.ann_retrieve_score() % 200) >> factory.index_scorer(query_encoded=True)\n",
    "\n",
    "#CIKM 2021 query embeddings paper: only keep the 9 tokens with highest ICF\n",
    "qep_pipe5 = (factory.query_encoder() \n",
    "            >> pyterrier_colbert.pruning.query_embedding_pruning(factory, 5) \n",
    "            >> factory.set_retrieve(query_encoded=True)\n",
    "            >> factory.index_scorer(query_encoded=False)\n",
    ")\n",
    "qep_pipe9 = (factory.query_encoder() \n",
    "            >> pyterrier_colbert.pruning.query_embedding_pruning(factory, 9) \n",
    "            >> factory.set_retrieve(query_encoded=True)\n",
    "            >> factory.index_scorer(query_encoded=False)\n",
    ")\n",
    "\n",
    "# a QEP baseline that suppresses [Q], [CLS]] and [MASK] tokens in the query\n",
    "nocls_nomask_noQ = (factory.query_encoder() \n",
    "            >> pyterrier_colbert.pruning.query_embedding_pruning_special(Q=True, CLS=True, MASK=True)\n",
    "            >> factory.set_retrieve(query_encoded=True)\n",
    "            >> factory.index_scorer(query_encoded=False)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mounted-kitty",
   "metadata": {},
   "source": [
    "## Experiment on TREC 2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "retained-october",
   "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>RR(rel=2)@100</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>AP(rel=2)@100</th>\n",
       "      <th>NumRet</th>\n",
       "      <th>num_q</th>\n",
       "      <th>mrt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ColBERT E2E</td>\n",
       "      <td>0.852713</td>\n",
       "      <td>0.693407</td>\n",
       "      <td>0.602398</td>\n",
       "      <td>0.386779</td>\n",
       "      <td>309698.0</td>\n",
       "      <td>43.0</td>\n",
       "      <td>671.341082</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Approx</td>\n",
       "      <td>0.870155</td>\n",
       "      <td>0.684195</td>\n",
       "      <td>0.534308</td>\n",
       "      <td>0.349277</td>\n",
       "      <td>8600.0</td>\n",
       "      <td>43.0</td>\n",
       "      <td>180.961294</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NoMASK NoCLS NoQ</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.693194</td>\n",
       "      <td>0.602279</td>\n",
       "      <td>0.385510</td>\n",
       "      <td>187493.0</td>\n",
       "      <td>43.0</td>\n",
       "      <td>414.449386</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>QEP 5 embs</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.695987</td>\n",
       "      <td>0.606343</td>\n",
       "      <td>0.389748</td>\n",
       "      <td>140232.0</td>\n",
       "      <td>43.0</td>\n",
       "      <td>342.599890</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>QEP 9 embs</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.693194</td>\n",
       "      <td>0.602088</td>\n",
       "      <td>0.385421</td>\n",
       "      <td>209672.0</td>\n",
       "      <td>43.0</td>\n",
       "      <td>432.380500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               name  RR(rel=2)@100   nDCG@10  nDCG@100  AP(rel=2)@100  \\\n",
       "0       ColBERT E2E       0.852713  0.693407  0.602398       0.386779   \n",
       "1            Approx       0.870155  0.684195  0.534308       0.349277   \n",
       "2  NoMASK NoCLS NoQ       0.853488  0.693194  0.602279       0.385510   \n",
       "3        QEP 5 embs       0.853488  0.695987  0.606343       0.389748   \n",
       "4        QEP 9 embs       0.853488  0.693194  0.602088       0.385421   \n",
       "\n",
       "     NumRet  num_q         mrt  \n",
       "0  309698.0   43.0  671.341082  \n",
       "1    8600.0   43.0  180.961294  \n",
       "2  187493.0   43.0  414.449386  \n",
       "3  140232.0   43.0  342.599890  \n",
       "4  209672.0   43.0  432.380500  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyterrier.measures import *\n",
    "pt.Experiment(\n",
    "    [\n",
    "        e2e,\n",
    "        ann_pipe,\n",
    "        nocls_nomask_noQ,\n",
    "        qep_pipe5,\n",
    "        qep_pipe9\n",
    "    ],\n",
    "    *pt.get_dataset(\"msmarco_passage\").get_topicsqrels(\"test-2019\"),\n",
    "    eval_metrics=[RR(rel=2)@100, nDCG@10, nDCG@100, AP(rel=2)@100, NumRet, \"mrt\", \"num_q\"],\n",
    "    names=[\"ColBERT E2E\", \"Approx\", \"NoMASK NoCLS NoQ\", \"QEP 5 embs\", \"QEP 9 embs\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "elect-player",
   "metadata": {},
   "source": [
    "Observations:\n",
    " - All four approaches result in good effectiveness (e.g. MRR, nDCG@10) while reducing then number of retrieved documents\n",
    " - In particular, Approx only retrieved 2% of the documents that E2E does, while enhancing MRR, and very small reduction in nDCG@10 (0.69 -> 0.68).\n",
    " - By applying QEP to reduce the 32 query embeddings to just 5 results in no real difference in MRR, NDCG@10, NDCG@10 and even MAP, while reducing by 50% the number of retrieved documents.\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "soviet-dominant",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Both papers propose methods to adapt ColBERT's dense retrieval pipeline to be more efficient without markedly reducing effectiveness. Further results and significance tests are provided in the respective papers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "elect-groove",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cikm2021",
   "language": "python",
   "name": "cikm2021"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}