notebooks/idf-pruning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "later-arnold",
   "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": "satisfactory-birth",
   "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": "annoying-picking",
   "metadata": {},
   "source": [
    "## load the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "short-perfume",
   "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 11, 14:36:21] #> Loading model checkpoint.\n",
      "[Feb 11, 14:36:21] #> 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 11, 14:36:31] #> checkpoint['epoch'] = 0\n",
      "[Feb 11, 14:36:31] #> 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=(\"/nfs/indices/colbert_passage\",\"index_name3\")\n",
    "factory = pyterrier_colbert.ranking.ColBERTFactory(checkpoint, *index, )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dominican-brisbane",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 11, 14:36:32] #> Loading the FAISS index from /nfs/indices/colbert_passage/index_name3/ivfpq.faiss ..\n",
      "[Feb 11, 14:37:04] #> Building the emb2pid mapping..\n",
      "[Feb 11, 14:37:27] len(self.emb2pid) = 687989391\n",
      "[Feb 11, 14:37:28] #> Building the emb2tid mapping..\n",
      "687989391\n",
      "Loading doclens\n"
     ]
    }
   ],
   "source": [
    "factory.faiss_index_on_gpu = False\n",
    "nn_term = factory.nn_term(df=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "informed-universe",
   "metadata": {},
   "source": [
    "The following cell forces the ColBERT embeddings index to be loaded into memory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bigger-bulgarian",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading reranking index, memtype=mem\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f157e0b0a444437895a46f031ea5928",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading index shards to memory:   0%|          | 0/24 [00:00<?, ?shard/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[28603376,\n",
       " 28935602,\n",
       " 28997615,\n",
       " 29001911,\n",
       " 29033097,\n",
       " 29094152,\n",
       " 28649536,\n",
       " 28723825,\n",
       " 28856902,\n",
       " 28970728,\n",
       " 29018059,\n",
       " 29025293,\n",
       " 29029230,\n",
       " 29065641,\n",
       " 29474976,\n",
       " 30061199,\n",
       " 33096705,\n",
       " 30056847,\n",
       " 29293124,\n",
       " 27428802,\n",
       " 28026855,\n",
       " 27215683,\n",
       " 28375647,\n",
       " 19954586]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[x.mmap.shape[0] for x in factory._rrm().part_mmap]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "outside-bracket",
   "metadata": {},
   "source": [
    "## index stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bacterial-sample",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(687989391)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total = torch.sum(nn_term.lookup)\n",
    "total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "interested-queensland",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8841823"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(factory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "incredible-collins",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(77.8108)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total/len(factory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-content",
   "metadata": {},
   "source": [
    "## eval setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "promotional-grace",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/terrierteam/ir_measures.git@runtime\n",
      "  Cloning https://github.com/terrierteam/ir_measures.git (to revision runtime) to /tmp/pip-req-build-3gv2mv1p\n",
      "  Running command git clone -q https://github.com/terrierteam/ir_measures.git /tmp/pip-req-build-3gv2mv1p\n",
      "  Running command git checkout -b runtime --track origin/runtime\n",
      "  Switched to a new branch 'runtime'\n",
      "  Branch 'runtime' set up to track remote branch 'runtime' from 'origin'.\n",
      "Building wheels for collected packages: ir-measures\n",
      "  Building wheel for ir-measures (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for ir-measures: filename=ir_measures-0.2.3-py3-none-any.whl size=54379 sha256=ed9cce06e033cd5ef86f3c6670f99712e036453662e04cecdc0c997ea793ce54\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-atfw1vrg/wheels/77/f7/6a/b55aee2281a6d74fae948bdba4c5b80e95cd8b6b2001fbd3e3\n",
      "Successfully built ir-measures\n",
      "Installing collected packages: ir-measures\n",
      "  Attempting uninstall: ir-measures\n",
      "    Found existing installation: ir-measures 0.2.3\n",
      "    Uninstalling ir-measures-0.2.3:\n",
      "      Successfully uninstalled ir-measures-0.2.3\n",
      "Successfully installed ir-measures-0.2.3\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --force-reinstall --no-deps git+https://github.com/terrierteam/ir_measures.git@runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "prostate-writer",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _AvgDocLen(qrels, run):\n",
    "    from warnings import warn\n",
    "    if 'doc_toks' not in run.columns:\n",
    "        if len(run) > 0:\n",
    "            warn(\"run of %d rows did not have doc_toks column; available columns: %s\" % (len(run), str(run.columns)))\n",
    "        else:\n",
    "            warn(\"empty run did not have doc_toks column; available columns: %s\" % str(run.columns))\n",
    "        return 0\n",
    "    return run['doc_toks'].apply(lambda row: row.shape[0]).mean()\n",
    "    \n",
    "import ir_measures\n",
    "AvgDocLen = ir_measures.define_byquery(\n",
    "    _AvgDocLen, \n",
    "    name=\"AvgDoclen\")\n",
    "\n",
    "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\", AvgDocLen@1, AvgDocLen@10, AvgDocLen@100, AvgDocLen]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "connected-content",
   "metadata": {},
   "source": [
    "## \"Original\" baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "three-cabinet",
   "metadata": {},
   "outputs": [],
   "source": [
    "baselineApprox = (\n",
    "    (factory.ann_retrieve_score() % 1000) \n",
    "    >> factory.fetch_index_encodings(ids=True, verbose=False) \n",
    "    >> factory.scorer()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "competent-curtis",
   "metadata": {},
   "source": [
    "The following function does some of the grunt work for performing the experimental evaluations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "significant-presentation",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "pleasant-contest",
   "metadata": {},
   "source": [
    "This loop provides the results files and eval results for the '2019', '2020', 'dev' topics sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "front-people",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f65ba86fb9144b22a247dfec74573f62",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=3.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "580e3ea35d1e422fae248f169e0393ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=3.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82012424b51041ecb057d79d65289203",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=349.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</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.405597</td>\n",
       "      <td>0.692943</td>\n",
       "      <td>0.649851</td>\n",
       "      <td>0.584233</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853659</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>67.883721</td>\n",
       "      <td>76.690698</td>\n",
       "      <td>79.226279</td>\n",
       "      <td>77.001302</td>\n",
       "      <td>645.998443</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.455459</td>\n",
       "      <td>0.689205</td>\n",
       "      <td>0.657368</td>\n",
       "      <td>0.611870</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852646</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>68.981481</td>\n",
       "      <td>74.324074</td>\n",
       "      <td>78.199815</td>\n",
       "      <td>76.993000</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>approx_baseline</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.401613</td>\n",
       "      <td>0.426102</td>\n",
       "      <td>0.454005</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.341710</td>\n",
       "      <td>0.352399</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>69.069054</td>\n",
       "      <td>74.954656</td>\n",
       "      <td>78.020663</td>\n",
       "      <td>76.502952</td>\n",
       "      <td>NaN</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.405597  0.692943  0.649851  0.584233   \n",
       "0  approx_baseline        0.455459  0.689205  0.657368  0.611870   \n",
       "0  approx_baseline        0.000000  0.401613  0.426102  0.454005   \n",
       "\n",
       "   RR(rel=2)@10  RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  \\\n",
       "0      0.853488   0.853659  0.939922  0.939922       0.695360    67.883721   \n",
       "0      0.850000   0.852646  0.922840  0.925485       0.750889    68.981481   \n",
       "0      0.000000   0.000000  0.341710  0.352399       0.000000    69.069054   \n",
       "\n",
       "   AvgDoclen@10  AvgDoclen@100  AvgDoclen         mrt  \n",
       "0     76.690698      79.226279  77.001302  645.998443  \n",
       "0     74.324074      78.199815  76.993000         NaN  \n",
       "0     74.954656      78.020663  76.502952         NaN  "
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_res = [\n",
    "    do_exp(\n",
    "        [baselineApprox], \n",
    "        names=['approx_baseline'], \n",
    "        topics=t,\n",
    "        save_dir='/nfs/staticPruning/results/repro/baseline_%s/'\n",
    "    ) for t in ['2019', '2020', 'dev']\n",
    "]\n",
    "pd.concat(all_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adopted-cemetery",
   "metadata": {},
   "source": [
    "## uniform IDF pruning\n",
    "\n",
    "We now design the implementations of the uniform IDF static pruning approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "departmental-hollywood",
   "metadata": {},
   "outputs": [],
   "source": [
    "def uniform_blacklisted_tokens_transformer(blacklist, verbose=False, zero=False, gpu=torch.cuda.is_available()) -> TransformerBase:\n",
    "    \"\"\"\n",
    "    Remove tokens and their embeddings from the document dataframe\n",
    "    input: qid, query_embs, docno, doc_embs, doc_toks\n",
    "    output: qid, query_embs, docno, doc_embs, doc_toks\n",
    "    \n",
    "    The blacklist parameters must contain a list of tokenids that should be removed\n",
    "    \"\"\"\n",
    "    import pyterrier as pt\n",
    "    import torch\n",
    "    import numpy as np\n",
    "    \n",
    "    assert pt.started(), 'PyTerrier must be started'\n",
    "    \n",
    "    if gpu: \n",
    "        blacklist = torch.Tensor(blacklist).cuda()\n",
    "    else:\n",
    "        blacklist = torch.Tensor(blacklist)\n",
    "        \n",
    "    pt.tqdm.pandas()\n",
    "    \n",
    "    if verbose: print(f'Blacklist composed of {len(blacklist)} elements.')\n",
    "        \n",
    "    def _prune_remove_gpu(row):\n",
    "        tokens = row['doc_toks'].cuda()\n",
    "        embeddings = row['doc_embs'].cuda()\n",
    "\n",
    "        # create the 1-D mask\n",
    "        final_mask = torch.all(tokens[None, : ] != blacklist[:, None], axis=0)\n",
    "\n",
    "        # apply the mask to the tokens\n",
    "        row['doc_toks'] = tokens[final_mask]\n",
    "\n",
    "        # pad the 1-d mask\n",
    "        row_embs_size = embeddings.size()\n",
    "        mask_1d = torch.cat((final_mask, torch.ones(row_embs_size[0] - final_mask.size()[0], dtype=torch.bool).cuda()))\n",
    "\n",
    "        # apply the padded 1-d mask\n",
    "        row['doc_embs'] = embeddings[mask_1d, : ]\n",
    "        return row\n",
    "\n",
    "    def _prune_remove(row):\n",
    "        tokens = row['doc_toks']\n",
    "        embeddings = row['doc_embs']\n",
    "\n",
    "        # create the 1-D mask\n",
    "        final_mask = torch.all(tokens[None, : ] != blacklist[:, None], axis=0)\n",
    "\n",
    "        # apply the mask to the tokens\n",
    "        row['doc_toks'] = tokens[final_mask]\n",
    "\n",
    "        # pad the 1-d mask\n",
    "        row_embs_size = embeddings.size()\n",
    "        mask_1d = torch.cat((final_mask, torch.ones(row_embs_size[0] - final_mask.size()[0], dtype=torch.bool)))\n",
    "\n",
    "        # apply the padded 1-d mask\n",
    "        row['doc_embs'] = embeddings[mask_1d, : ]\n",
    "        return row\n",
    "    \n",
    "    def _prune_gpu_zero(row):\n",
    "        tokens = row['doc_toks'].cuda()\n",
    "        embeddings = row['doc_embs'].cuda()\n",
    "        row_embs_size = embeddings.size()\n",
    "        tokens_size = tokens.size()[0]\n",
    "        \n",
    "        # create the 1-D mask\n",
    "        final_mask = torch.zeros(row_embs_size[0], dtype=torch.bool)\n",
    "        final_mask[0:tokens_size] = torch.any(tokens[None, :] == blacklist[:, None], axis=0)\n",
    "        \n",
    "        # apply the mask\n",
    "        row['doc_toks'][final_mask[0:tokens_size]] = 0\n",
    "        row['doc_embs'][final_mask, :] = 0 \n",
    "        return row\n",
    "        \n",
    "    def _prune_zero(row):\n",
    "        tokens = row['doc_toks']\n",
    "        embeddings = row['doc_embs']\n",
    "        row_embs_size = embeddings.size()\n",
    "        tokens_size = tokens.size()[0]\n",
    "                                        \n",
    "        # create the 1-D mask\n",
    "        final_mask = torch.zeros(row_embs_size[0], dtype=torch.bool)\n",
    "        final_mask[0:tokens_size] = torch.any(tokens[None, :] == blacklist[:, None], axis=0)\n",
    "            \n",
    "        # apply the mask\n",
    "        row['doc_toks'][final_mask[0:tokens_size]] = 0\n",
    "        row['doc_embs'][final_mask, :] = 0\n",
    "        return row\n",
    "    \n",
    "    if zero:\n",
    "        prune_function = _prune_gpu_zero if gpu else _prune_zero\n",
    "    else:\n",
    "        prune_function = _prune_remove_gpu if gpu else _prune_remove\n",
    "\n",
    "    def _apply(df):\n",
    "        if verbose:\n",
    "            df = df.progress_apply(prune_function, axis=1)\n",
    "        else:\n",
    "            df = df.apply(prune_function, axis=1)\n",
    "        return df\n",
    "\n",
    "    return pt.apply.generic(_apply)\n",
    "\n",
    "def get_uniform_pipeline(factory, blacklist, k1=1000):\n",
    "    pipeline = (\n",
    "        factory.query_encoder()\n",
    "        >> (factory.ann_retrieve_score(query_encoded=True) % k1)\n",
    "        >> factory.fetch_index_encodings(ids=True, verbose=False)\n",
    "        >> uniform_blacklisted_tokens_transformer(blacklist, verbose=False, zero=False)\n",
    "        >> factory.scorer(verbose=False)\n",
    "        #>> pt.apply.doc_embs(drop=True)\n",
    "        #>> pt.apply.query_embs(drop=True)\n",
    "    )\n",
    "    return pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "expired-research",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyterrier_colbert.faiss_term_index import FaissNNTerm\n",
    "def get_tid_ordered_by_idf(faiss_nn_term : FaissNNTerm, verbose=False):\n",
    "    import math\n",
    "    assert hasattr(faiss_nn_term, 'dfs'), 'FaissNNTerm class must be initialized with dfs=True'\n",
    "\n",
    "    vocabulary = faiss_nn_term.tok.get_vocab()\n",
    "    n_docs = faiss_nn_term.num_docs\n",
    "\n",
    "    if verbose:\n",
    "        print(f'Number of docs: {n_docs}')\n",
    "        print(f'Vocabulary Length: {len(vocabulary)}')\n",
    "\n",
    "    tokens = []\n",
    "    for token in vocabulary:\n",
    "        tid = vocabulary[token]\n",
    "        df = faiss_nn_term.getDF_by_id(tid)\n",
    "        if(df != 0):\n",
    "            idf = math.log(n_docs/(df + 1), 10)\n",
    "            tokens.append((tid, idf))\n",
    "\n",
    "    # Remove items with 0 document frequency\n",
    "    if verbose: print(\"Token length (without 0 df elements):\", len(tokens))\n",
    "    # order by inverse document frequency\n",
    "    ordered_tokens = sorted(tokens, key= lambda pair: pair[1])\n",
    "    final_token_list = [_id for _id, _ in ordered_tokens]\n",
    "    return final_token_list\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "psychological-blast",
   "metadata": {},
   "source": [
    "DROPS corresponds to the $\\tau_u$ in the paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "informal-colombia",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1012, 1996, 1010, 1997, 1998, 1037, 2000, 1999, 2003, 2005]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DROPS=[1, 10, 25, 50, 100, 1000, 2000, 3000]\n",
    "\n",
    "\n",
    "tid_ordered_by_idf = get_tid_ordered_by_idf(nn_term)\n",
    "\n",
    "#we always start from 3 to ignore the CLS, D tokens etc.\n",
    "tid_ordered_by_idf[3:13]\n"
   ]
  },
  {
   "cell_type": "raw",
   "id": "thrown-mention",
   "metadata": {},
   "source": [
    "The following cells execute the uniform IDF approach on 2019, 2020 and dev topics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "arctic-dictionary",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ab3df1c411645bda537c2236515ef99",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=24.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.405961</td>\n",
       "      <td>0.694331</td>\n",
       "      <td>0.651018</td>\n",
       "      <td>0.584537</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853662</td>\n",
       "      <td>0.941860</td>\n",
       "      <td>0.941860</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>64.558140</td>\n",
       "      <td>72.981395</td>\n",
       "      <td>75.453256</td>\n",
       "      <td>73.233512</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.404652</td>\n",
       "      <td>0.695357</td>\n",
       "      <td>0.652229</td>\n",
       "      <td>0.582998</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865290</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>50.093023</td>\n",
       "      <td>57.802326</td>\n",
       "      <td>59.156512</td>\n",
       "      <td>57.486000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.404306</td>\n",
       "      <td>0.686533</td>\n",
       "      <td>0.648563</td>\n",
       "      <td>0.581893</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865295</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>44.279070</td>\n",
       "      <td>51.158140</td>\n",
       "      <td>52.490000</td>\n",
       "      <td>50.999023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.403207</td>\n",
       "      <td>0.686432</td>\n",
       "      <td>0.649681</td>\n",
       "      <td>0.580809</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853674</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>40.906977</td>\n",
       "      <td>46.897674</td>\n",
       "      <td>48.471395</td>\n",
       "      <td>47.087442</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.402592</td>\n",
       "      <td>0.688892</td>\n",
       "      <td>0.649656</td>\n",
       "      <td>0.579712</td>\n",
       "      <td>0.845736</td>\n",
       "      <td>0.845922</td>\n",
       "      <td>0.932171</td>\n",
       "      <td>0.932171</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>37.651163</td>\n",
       "      <td>43.348837</td>\n",
       "      <td>44.694186</td>\n",
       "      <td>43.435581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.378553</td>\n",
       "      <td>0.647829</td>\n",
       "      <td>0.611177</td>\n",
       "      <td>0.540572</td>\n",
       "      <td>0.843023</td>\n",
       "      <td>0.844743</td>\n",
       "      <td>0.947674</td>\n",
       "      <td>0.949042</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>25.046512</td>\n",
       "      <td>28.367442</td>\n",
       "      <td>27.860930</td>\n",
       "      <td>27.453767</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.333716</td>\n",
       "      <td>0.613761</td>\n",
       "      <td>0.565548</td>\n",
       "      <td>0.504106</td>\n",
       "      <td>0.843798</td>\n",
       "      <td>0.846744</td>\n",
       "      <td>0.937984</td>\n",
       "      <td>0.939535</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>19.418605</td>\n",
       "      <td>20.904651</td>\n",
       "      <td>21.237209</td>\n",
       "      <td>20.870930</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.315022</td>\n",
       "      <td>0.568992</td>\n",
       "      <td>0.520151</td>\n",
       "      <td>0.471269</td>\n",
       "      <td>0.814175</td>\n",
       "      <td>0.817290</td>\n",
       "      <td>0.879291</td>\n",
       "      <td>0.883096</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>15.279070</td>\n",
       "      <td>16.951163</td>\n",
       "      <td>17.212093</td>\n",
       "      <td>16.873256</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1        0.405961  0.694331  0.651018  0.584537      0.853488   \n",
       "1    uniform_10        0.404652  0.695357  0.652229  0.582998      0.865116   \n",
       "2    uniform_25        0.404306  0.686533  0.648563  0.581893      0.865116   \n",
       "3    uniform_50        0.403207  0.686432  0.649681  0.580809      0.853488   \n",
       "4   uniform_100        0.402592  0.688892  0.649656  0.579712      0.845736   \n",
       "5  uniform_1000        0.378553  0.647829  0.611177  0.540572      0.843023   \n",
       "6  uniform_2000        0.333716  0.613761  0.565548  0.504106      0.843798   \n",
       "7  uniform_3000        0.315022  0.568992  0.520151  0.471269      0.814175   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.853662  0.941860  0.941860        0.69536    64.558140     72.981395   \n",
       "1   0.865290  0.939922  0.939922        0.69536    50.093023     57.802326   \n",
       "2   0.865295  0.939922  0.939922        0.69536    44.279070     51.158140   \n",
       "3   0.853674  0.939922  0.939922        0.69536    40.906977     46.897674   \n",
       "4   0.845922  0.932171  0.932171        0.69536    37.651163     43.348837   \n",
       "5   0.844743  0.947674  0.949042        0.69536    25.046512     28.367442   \n",
       "6   0.846744  0.937984  0.939535        0.69536    19.418605     20.904651   \n",
       "7   0.817290  0.879291  0.883096        0.69536    15.279070     16.951163   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      75.453256  73.233512  \n",
       "1      59.156512  57.486000  \n",
       "2      52.490000  50.999023  \n",
       "3      48.471395  47.087442  \n",
       "4      44.694186  43.435581  \n",
       "5      27.860930  27.453767  \n",
       "6      21.237209  20.870930  \n",
       "7      17.212093  16.873256  "
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idf_2019 = do_exp(\n",
    "    [get_uniform_pipeline(factory, tid_ordered_by_idf[3:x+3]) for x in DROPS ], \n",
    "    topics='2019',\n",
    "    names=[\"uniform_%d\" % x for x in DROPS ],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/uniform_%s/\"\n",
    ")\n",
    "idf_2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "clear-alexandria",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1cf159e6c3154105a7f4266ca9eed54b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=24.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.455497</td>\n",
       "      <td>0.688954</td>\n",
       "      <td>0.657722</td>\n",
       "      <td>0.611513</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925397</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>65.759259</td>\n",
       "      <td>70.692593</td>\n",
       "      <td>74.533333</td>\n",
       "      <td>73.139389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.457219</td>\n",
       "      <td>0.690652</td>\n",
       "      <td>0.656046</td>\n",
       "      <td>0.611886</td>\n",
       "      <td>0.857716</td>\n",
       "      <td>0.860565</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925689</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>53.055556</td>\n",
       "      <td>55.650000</td>\n",
       "      <td>58.687407</td>\n",
       "      <td>57.559870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.460114</td>\n",
       "      <td>0.690334</td>\n",
       "      <td>0.659018</td>\n",
       "      <td>0.610760</td>\n",
       "      <td>0.858951</td>\n",
       "      <td>0.860634</td>\n",
       "      <td>0.924691</td>\n",
       "      <td>0.926375</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>46.166667</td>\n",
       "      <td>49.550000</td>\n",
       "      <td>52.143333</td>\n",
       "      <td>51.190611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.457625</td>\n",
       "      <td>0.691289</td>\n",
       "      <td>0.658722</td>\n",
       "      <td>0.612326</td>\n",
       "      <td>0.850485</td>\n",
       "      <td>0.852169</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.927169</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>43.129630</td>\n",
       "      <td>45.505556</td>\n",
       "      <td>47.872037</td>\n",
       "      <td>46.984296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.457333</td>\n",
       "      <td>0.688052</td>\n",
       "      <td>0.657896</td>\n",
       "      <td>0.612444</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>39.462963</td>\n",
       "      <td>41.918519</td>\n",
       "      <td>43.905185</td>\n",
       "      <td>43.076444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.421563</td>\n",
       "      <td>0.663536</td>\n",
       "      <td>0.605718</td>\n",
       "      <td>0.574324</td>\n",
       "      <td>0.815608</td>\n",
       "      <td>0.815608</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>24.111111</td>\n",
       "      <td>27.203704</td>\n",
       "      <td>28.050556</td>\n",
       "      <td>27.133370</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.383755</td>\n",
       "      <td>0.593645</td>\n",
       "      <td>0.556330</td>\n",
       "      <td>0.529415</td>\n",
       "      <td>0.754652</td>\n",
       "      <td>0.758675</td>\n",
       "      <td>0.905710</td>\n",
       "      <td>0.906867</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>18.759259</td>\n",
       "      <td>20.633333</td>\n",
       "      <td>21.314815</td>\n",
       "      <td>20.536315</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.368490</td>\n",
       "      <td>0.570477</td>\n",
       "      <td>0.544948</td>\n",
       "      <td>0.513365</td>\n",
       "      <td>0.753131</td>\n",
       "      <td>0.754365</td>\n",
       "      <td>0.889065</td>\n",
       "      <td>0.889065</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>16.055556</td>\n",
       "      <td>17.098148</td>\n",
       "      <td>17.577963</td>\n",
       "      <td>16.850463</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1        0.455497  0.688954  0.657722  0.611513      0.850000   \n",
       "1    uniform_10        0.457219  0.690652  0.656046  0.611886      0.857716   \n",
       "2    uniform_25        0.460114  0.690334  0.659018  0.610760      0.858951   \n",
       "3    uniform_50        0.457625  0.691289  0.658722  0.612326      0.850485   \n",
       "4   uniform_100        0.457333  0.688052  0.657896  0.612444      0.850044   \n",
       "5  uniform_1000        0.421563  0.663536  0.605718  0.574324      0.815608   \n",
       "6  uniform_2000        0.383755  0.593645  0.556330  0.529415      0.754652   \n",
       "7  uniform_3000        0.368490  0.570477  0.544948  0.513365      0.753131   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.852557  0.922840  0.925397       0.750889    65.759259     70.692593   \n",
       "1   0.860565  0.922840  0.925689       0.750889    53.055556     55.650000   \n",
       "2   0.860634  0.924691  0.926375       0.750889    46.166667     49.550000   \n",
       "3   0.852169  0.925485  0.927169       0.750889    43.129630     45.505556   \n",
       "4   0.851728  0.925485  0.927169       0.750889    39.462963     41.918519   \n",
       "5   0.815608  0.950617  0.950617       0.750889    24.111111     27.203704   \n",
       "6   0.758675  0.905710  0.906867       0.750889    18.759259     20.633333   \n",
       "7   0.754365  0.889065  0.889065       0.750889    16.055556     17.098148   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      74.533333  73.139389  \n",
       "1      58.687407  57.559870  \n",
       "2      52.143333  51.190611  \n",
       "3      47.872037  46.984296  \n",
       "4      43.905185  43.076444  \n",
       "5      28.050556  27.133370  \n",
       "6      21.314815  20.536315  \n",
       "7      17.577963  16.850463  "
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idf_2020 = do_exp(\n",
    "    [get_uniform_pipeline(factory, tid_ordered_by_idf[3:x+3]) for x in DROPS ],\n",
    "    topics='2020',\n",
    "    names=[\"uniform_%d\" % x for x in DROPS ],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/uniform_%s/\"\n",
    ")\n",
    "idf_2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "loved-pressure",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd5524280fc64573a654f797e9a6b77a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=2792.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.400824</td>\n",
       "      <td>0.425659</td>\n",
       "      <td>0.453466</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.340971</td>\n",
       "      <td>0.351744</td>\n",
       "      <td>0.0</td>\n",
       "      <td>65.435817</td>\n",
       "      <td>71.118438</td>\n",
       "      <td>74.046887</td>\n",
       "      <td>72.612341</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.399779</td>\n",
       "      <td>0.425007</td>\n",
       "      <td>0.452798</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.340252</td>\n",
       "      <td>0.351191</td>\n",
       "      <td>0.0</td>\n",
       "      <td>52.039971</td>\n",
       "      <td>56.304799</td>\n",
       "      <td>58.404444</td>\n",
       "      <td>57.268624</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.398953</td>\n",
       "      <td>0.424126</td>\n",
       "      <td>0.452151</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.339520</td>\n",
       "      <td>0.350502</td>\n",
       "      <td>0.0</td>\n",
       "      <td>46.336676</td>\n",
       "      <td>49.993123</td>\n",
       "      <td>51.876050</td>\n",
       "      <td>50.917969</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.398091</td>\n",
       "      <td>0.423260</td>\n",
       "      <td>0.451413</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.338807</td>\n",
       "      <td>0.349741</td>\n",
       "      <td>0.0</td>\n",
       "      <td>42.690115</td>\n",
       "      <td>45.927808</td>\n",
       "      <td>47.573274</td>\n",
       "      <td>46.675498</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.394788</td>\n",
       "      <td>0.420355</td>\n",
       "      <td>0.448464</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.335730</td>\n",
       "      <td>0.346737</td>\n",
       "      <td>0.0</td>\n",
       "      <td>39.240258</td>\n",
       "      <td>42.217837</td>\n",
       "      <td>43.711490</td>\n",
       "      <td>42.844141</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.359918</td>\n",
       "      <td>0.383584</td>\n",
       "      <td>0.415319</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.303860</td>\n",
       "      <td>0.314943</td>\n",
       "      <td>0.0</td>\n",
       "      <td>25.271347</td>\n",
       "      <td>26.715774</td>\n",
       "      <td>27.406635</td>\n",
       "      <td>26.782411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.333851</td>\n",
       "      <td>0.357203</td>\n",
       "      <td>0.390378</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.280669</td>\n",
       "      <td>0.291872</td>\n",
       "      <td>0.0</td>\n",
       "      <td>19.515616</td>\n",
       "      <td>20.455186</td>\n",
       "      <td>20.911254</td>\n",
       "      <td>20.361803</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.313844</td>\n",
       "      <td>0.336802</td>\n",
       "      <td>0.370560</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.263599</td>\n",
       "      <td>0.274841</td>\n",
       "      <td>0.0</td>\n",
       "      <td>16.043696</td>\n",
       "      <td>16.724212</td>\n",
       "      <td>17.021218</td>\n",
       "      <td>16.539694</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1             0.0  0.400824  0.425659  0.453466           0.0   \n",
       "1    uniform_10             0.0  0.399779  0.425007  0.452798           0.0   \n",
       "2    uniform_25             0.0  0.398953  0.424126  0.452151           0.0   \n",
       "3    uniform_50             0.0  0.398091  0.423260  0.451413           0.0   \n",
       "4   uniform_100             0.0  0.394788  0.420355  0.448464           0.0   \n",
       "5  uniform_1000             0.0  0.359918  0.383584  0.415319           0.0   \n",
       "6  uniform_2000             0.0  0.333851  0.357203  0.390378           0.0   \n",
       "7  uniform_3000             0.0  0.313844  0.336802  0.370560           0.0   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0        0.0  0.340971  0.351744            0.0    65.435817     71.118438   \n",
       "1        0.0  0.340252  0.351191            0.0    52.039971     56.304799   \n",
       "2        0.0  0.339520  0.350502            0.0    46.336676     49.993123   \n",
       "3        0.0  0.338807  0.349741            0.0    42.690115     45.927808   \n",
       "4        0.0  0.335730  0.346737            0.0    39.240258     42.217837   \n",
       "5        0.0  0.303860  0.314943            0.0    25.271347     26.715774   \n",
       "6        0.0  0.280669  0.291872            0.0    19.515616     20.455186   \n",
       "7        0.0  0.263599  0.274841            0.0    16.043696     16.724212   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      74.046887  72.612341  \n",
       "1      58.404444  57.268624  \n",
       "2      51.876050  50.917969  \n",
       "3      47.573274  46.675498  \n",
       "4      43.711490  42.844141  \n",
       "5      27.406635  26.782411  \n",
       "6      20.911254  20.361803  \n",
       "7      17.021218  16.539694  "
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idf_dev = do_exp(\n",
    "    [get_uniform_pipeline(factory, tid_ordered_by_idf[3:x+3]) for x in DROPS ],\n",
    "    topics='dev',\n",
    "    names=[\"uniform_%d\" % x for x in DROPS ],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/uniform_%s/\"\n",
    ")\n",
    "idf_dev"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banned-tissue",
   "metadata": {},
   "source": [
    "The next 2 cells ran some additional datapoints for the uniform pruning, and copied them into the files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "hired-shield",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "716d9840b5914e2fb1f3067958527620",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/9 [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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.388886</td>\n",
       "      <td>0.657279</td>\n",
       "      <td>0.624054</td>\n",
       "      <td>0.566356</td>\n",
       "      <td>0.807586</td>\n",
       "      <td>0.809249</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.937271</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>32.209302</td>\n",
       "      <td>38.093023</td>\n",
       "      <td>38.870233</td>\n",
       "      <td>38.075884</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.375526</td>\n",
       "      <td>0.643821</td>\n",
       "      <td>0.607205</td>\n",
       "      <td>0.541373</td>\n",
       "      <td>0.802326</td>\n",
       "      <td>0.804457</td>\n",
       "      <td>0.924419</td>\n",
       "      <td>0.926208</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>30.209302</td>\n",
       "      <td>33.630233</td>\n",
       "      <td>33.715814</td>\n",
       "      <td>33.177233</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.384139</td>\n",
       "      <td>0.648789</td>\n",
       "      <td>0.611416</td>\n",
       "      <td>0.544533</td>\n",
       "      <td>0.815891</td>\n",
       "      <td>0.817790</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.937597</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>27.720930</td>\n",
       "      <td>30.618605</td>\n",
       "      <td>30.358140</td>\n",
       "      <td>29.916605</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0  uniform_250        0.388886  0.657279  0.624054  0.566356      0.807586   \n",
       "1  uniform_500        0.375526  0.643821  0.607205  0.541373      0.802326   \n",
       "2  uniform_750        0.384139  0.648789  0.611416  0.544533      0.815891   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.809249  0.936047  0.937271       0.695251    32.209302     38.093023   \n",
       "1   0.804457  0.924419  0.926208       0.695251    30.209302     33.630233   \n",
       "2   0.817790  0.936047  0.937597       0.695251    27.720930     30.618605   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      38.870233  38.075884  \n",
       "1      33.715814  33.177233  \n",
       "2      30.358140  29.916605  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d7aacc3c9b746cda27cf666094fceaa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/9 [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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.454676</td>\n",
       "      <td>0.690524</td>\n",
       "      <td>0.652513</td>\n",
       "      <td>0.607058</td>\n",
       "      <td>0.860516</td>\n",
       "      <td>0.860516</td>\n",
       "      <td>0.943673</td>\n",
       "      <td>0.943673</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>33.185185</td>\n",
       "      <td>36.896296</td>\n",
       "      <td>38.307407</td>\n",
       "      <td>37.609278</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.444291</td>\n",
       "      <td>0.681340</td>\n",
       "      <td>0.631334</td>\n",
       "      <td>0.596277</td>\n",
       "      <td>0.844753</td>\n",
       "      <td>0.844753</td>\n",
       "      <td>0.958333</td>\n",
       "      <td>0.958333</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>29.462963</td>\n",
       "      <td>32.405556</td>\n",
       "      <td>33.394074</td>\n",
       "      <td>32.691685</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.436530</td>\n",
       "      <td>0.679331</td>\n",
       "      <td>0.625340</td>\n",
       "      <td>0.587080</td>\n",
       "      <td>0.846296</td>\n",
       "      <td>0.846296</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>26.537037</td>\n",
       "      <td>29.579630</td>\n",
       "      <td>30.361481</td>\n",
       "      <td>29.504944</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0  uniform_250        0.454676  0.690524  0.652513  0.607058      0.860516   \n",
       "1  uniform_500        0.444291  0.681340  0.631334  0.596277      0.844753   \n",
       "2  uniform_750        0.436530  0.679331  0.625340  0.587080      0.846296   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.860516  0.943673  0.943673       0.754593    33.185185     36.896296   \n",
       "1   0.844753  0.958333  0.958333       0.754593    29.462963     32.405556   \n",
       "2   0.846296  0.950617  0.950617       0.754593    26.537037     29.579630   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      38.307407  37.609278  \n",
       "1      33.394074  32.691685  \n",
       "2      30.361481  29.504944  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df86f3b080ab4e37ab63c385cfa82460",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/1047 [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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.387861</td>\n",
       "      <td>0.412696</td>\n",
       "      <td>0.441257</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.328845</td>\n",
       "      <td>0.339812</td>\n",
       "      <td>0.0</td>\n",
       "      <td>34.629513</td>\n",
       "      <td>37.015287</td>\n",
       "      <td>38.195284</td>\n",
       "      <td>37.384855</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.373959</td>\n",
       "      <td>0.398772</td>\n",
       "      <td>0.428734</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.316276</td>\n",
       "      <td>0.327443</td>\n",
       "      <td>0.0</td>\n",
       "      <td>30.188682</td>\n",
       "      <td>32.125129</td>\n",
       "      <td>33.118317</td>\n",
       "      <td>32.445901</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.365723</td>\n",
       "      <td>0.389993</td>\n",
       "      <td>0.420990</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.308556</td>\n",
       "      <td>0.319766</td>\n",
       "      <td>0.0</td>\n",
       "      <td>27.420057</td>\n",
       "      <td>29.093968</td>\n",
       "      <td>29.898242</td>\n",
       "      <td>29.234465</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0  uniform_250             0.0  0.387861  0.412696  0.441257           0.0   \n",
       "1  uniform_500             0.0  0.373959  0.398772  0.428734           0.0   \n",
       "2  uniform_750             0.0  0.365723  0.389993  0.420990           0.0   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0        0.0  0.328845  0.339812            0.0    34.629513     37.015287   \n",
       "1        0.0  0.316276  0.327443            0.0    30.188682     32.125129   \n",
       "2        0.0  0.308556  0.319766            0.0    27.420057     29.093968   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      38.195284  37.384855  \n",
       "1      33.118317  32.445901  \n",
       "2      29.898242  29.234465  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display\n",
    "for t in ['2019', '2020', 'dev']:\n",
    "    idf = do_exp(\n",
    "        [get_uniform_pipeline(factory, tid_ordered_by_idf[3:x+3]) for x in [250,500,750] ], \n",
    "        topics=t,\n",
    "        names=[\"uniform_%d\" % x for x in [250,500,750] ],\n",
    "        save_dir=\"/nfs/staticPruning/results/repro/supplement_uniform_%s/\"\n",
    "    )\n",
    "    display(idf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "champion-indication",
   "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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.405961</td>\n",
       "      <td>0.694331</td>\n",
       "      <td>0.651018</td>\n",
       "      <td>0.584537</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853662</td>\n",
       "      <td>0.941860</td>\n",
       "      <td>0.941860</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>64.558140</td>\n",
       "      <td>72.981395</td>\n",
       "      <td>75.453256</td>\n",
       "      <td>73.233512</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.404652</td>\n",
       "      <td>0.695357</td>\n",
       "      <td>0.652229</td>\n",
       "      <td>0.582998</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865290</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>50.093023</td>\n",
       "      <td>57.802326</td>\n",
       "      <td>59.156512</td>\n",
       "      <td>57.486000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.404306</td>\n",
       "      <td>0.686533</td>\n",
       "      <td>0.648563</td>\n",
       "      <td>0.581893</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865295</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>44.279070</td>\n",
       "      <td>51.158140</td>\n",
       "      <td>52.490000</td>\n",
       "      <td>50.999023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.403207</td>\n",
       "      <td>0.686432</td>\n",
       "      <td>0.649681</td>\n",
       "      <td>0.580809</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853674</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>40.906977</td>\n",
       "      <td>46.897674</td>\n",
       "      <td>48.471395</td>\n",
       "      <td>47.087442</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.402592</td>\n",
       "      <td>0.688892</td>\n",
       "      <td>0.649656</td>\n",
       "      <td>0.579712</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>37.651163</td>\n",
       "      <td>43.348837</td>\n",
       "      <td>44.694186</td>\n",
       "      <td>43.435581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.388886</td>\n",
       "      <td>0.657279</td>\n",
       "      <td>0.624054</td>\n",
       "      <td>0.566356</td>\n",
       "      <td>0.807586</td>\n",
       "      <td>0.809249</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.937271</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>32.209302</td>\n",
       "      <td>38.093023</td>\n",
       "      <td>38.870233</td>\n",
       "      <td>38.075884</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.375526</td>\n",
       "      <td>0.643821</td>\n",
       "      <td>0.607205</td>\n",
       "      <td>0.541373</td>\n",
       "      <td>0.802326</td>\n",
       "      <td>0.804457</td>\n",
       "      <td>0.924419</td>\n",
       "      <td>0.926208</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>30.209302</td>\n",
       "      <td>33.630233</td>\n",
       "      <td>33.715814</td>\n",
       "      <td>33.177233</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.384139</td>\n",
       "      <td>0.648789</td>\n",
       "      <td>0.611416</td>\n",
       "      <td>0.544533</td>\n",
       "      <td>0.815891</td>\n",
       "      <td>0.817790</td>\n",
       "      <td>0.936047</td>\n",
       "      <td>0.937597</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>27.720930</td>\n",
       "      <td>30.618605</td>\n",
       "      <td>30.358140</td>\n",
       "      <td>29.916605</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.378553</td>\n",
       "      <td>0.647829</td>\n",
       "      <td>0.611177</td>\n",
       "      <td>0.540572</td>\n",
       "      <td>0.843023</td>\n",
       "      <td>0.844743</td>\n",
       "      <td>0.947674</td>\n",
       "      <td>0.949042</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>25.046512</td>\n",
       "      <td>28.367442</td>\n",
       "      <td>27.860930</td>\n",
       "      <td>27.453767</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.333716</td>\n",
       "      <td>0.613761</td>\n",
       "      <td>0.565548</td>\n",
       "      <td>0.504106</td>\n",
       "      <td>0.843798</td>\n",
       "      <td>0.846744</td>\n",
       "      <td>0.937984</td>\n",
       "      <td>0.939535</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>19.418605</td>\n",
       "      <td>20.904651</td>\n",
       "      <td>21.237209</td>\n",
       "      <td>20.870930</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.315022</td>\n",
       "      <td>0.568992</td>\n",
       "      <td>0.520151</td>\n",
       "      <td>0.471269</td>\n",
       "      <td>0.814175</td>\n",
       "      <td>0.817290</td>\n",
       "      <td>0.879291</td>\n",
       "      <td>0.883096</td>\n",
       "      <td>0.695360</td>\n",
       "      <td>15.279070</td>\n",
       "      <td>16.951163</td>\n",
       "      <td>17.212093</td>\n",
       "      <td>16.873256</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1        0.405961  0.694331  0.651018  0.584537      0.853488   \n",
       "1    uniform_10        0.404652  0.695357  0.652229  0.582998      0.865116   \n",
       "2    uniform_25        0.404306  0.686533  0.648563  0.581893      0.865116   \n",
       "3    uniform_50        0.403207  0.686432  0.649681  0.580809      0.853488   \n",
       "4   uniform_100        0.402592  0.688892  0.649656  0.579712      0.845736   \n",
       "0   uniform_250        0.388886  0.657279  0.624054  0.566356      0.807586   \n",
       "1   uniform_500        0.375526  0.643821  0.607205  0.541373      0.802326   \n",
       "2   uniform_750        0.384139  0.648789  0.611416  0.544533      0.815891   \n",
       "5  uniform_1000        0.378553  0.647829  0.611177  0.540572      0.843023   \n",
       "6  uniform_2000        0.333716  0.613761  0.565548  0.504106      0.843798   \n",
       "7  uniform_3000        0.315022  0.568992  0.520151  0.471269      0.814175   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.853662  0.941860  0.941860       0.695360    64.558140     72.981395   \n",
       "1   0.865290  0.939922  0.939922       0.695360    50.093023     57.802326   \n",
       "2   0.865295  0.939922  0.939922       0.695360    44.279070     51.158140   \n",
       "3   0.853674  0.939922  0.939922       0.695360    40.906977     46.897674   \n",
       "4   0.845922  0.932171  0.932171       0.695360    37.651163     43.348837   \n",
       "0   0.809249  0.936047  0.937271       0.695251    32.209302     38.093023   \n",
       "1   0.804457  0.924419  0.926208       0.695251    30.209302     33.630233   \n",
       "2   0.817790  0.936047  0.937597       0.695251    27.720930     30.618605   \n",
       "5   0.844743  0.947674  0.949042       0.695360    25.046512     28.367442   \n",
       "6   0.846744  0.937984  0.939535       0.695360    19.418605     20.904651   \n",
       "7   0.817290  0.879291  0.883096       0.695360    15.279070     16.951163   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      75.453256  73.233512  \n",
       "1      59.156512  57.486000  \n",
       "2      52.490000  50.999023  \n",
       "3      48.471395  47.087442  \n",
       "4      44.694186  43.435581  \n",
       "0      38.870233  38.075884  \n",
       "1      33.715814  33.177233  \n",
       "2      30.358140  29.916605  \n",
       "5      27.860930  27.453767  \n",
       "6      21.237209  20.870930  \n",
       "7      17.212093  16.873256  "
      ]
     },
     "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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.455497</td>\n",
       "      <td>0.688954</td>\n",
       "      <td>0.657722</td>\n",
       "      <td>0.611513</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925397</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>65.759259</td>\n",
       "      <td>70.692593</td>\n",
       "      <td>74.533333</td>\n",
       "      <td>73.139389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.457219</td>\n",
       "      <td>0.690652</td>\n",
       "      <td>0.656046</td>\n",
       "      <td>0.611886</td>\n",
       "      <td>0.857716</td>\n",
       "      <td>0.860565</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925689</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>53.055556</td>\n",
       "      <td>55.650000</td>\n",
       "      <td>58.687407</td>\n",
       "      <td>57.559870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.460114</td>\n",
       "      <td>0.690334</td>\n",
       "      <td>0.659018</td>\n",
       "      <td>0.610760</td>\n",
       "      <td>0.858951</td>\n",
       "      <td>0.860634</td>\n",
       "      <td>0.924691</td>\n",
       "      <td>0.926375</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>46.166667</td>\n",
       "      <td>49.550000</td>\n",
       "      <td>52.143333</td>\n",
       "      <td>51.190611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.457625</td>\n",
       "      <td>0.691289</td>\n",
       "      <td>0.658722</td>\n",
       "      <td>0.612326</td>\n",
       "      <td>0.850485</td>\n",
       "      <td>0.852169</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.927169</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>43.129630</td>\n",
       "      <td>45.505556</td>\n",
       "      <td>47.872037</td>\n",
       "      <td>46.984296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.457333</td>\n",
       "      <td>0.688052</td>\n",
       "      <td>0.657896</td>\n",
       "      <td>0.612444</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>39.462963</td>\n",
       "      <td>41.918519</td>\n",
       "      <td>43.905185</td>\n",
       "      <td>43.076444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.454676</td>\n",
       "      <td>0.690524</td>\n",
       "      <td>0.652513</td>\n",
       "      <td>0.607058</td>\n",
       "      <td>0.860516</td>\n",
       "      <td>0.860516</td>\n",
       "      <td>0.943673</td>\n",
       "      <td>0.943673</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>33.185185</td>\n",
       "      <td>36.896296</td>\n",
       "      <td>38.307407</td>\n",
       "      <td>37.609278</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.444291</td>\n",
       "      <td>0.681340</td>\n",
       "      <td>0.631334</td>\n",
       "      <td>0.596277</td>\n",
       "      <td>0.844753</td>\n",
       "      <td>0.844753</td>\n",
       "      <td>0.958333</td>\n",
       "      <td>0.958333</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>29.462963</td>\n",
       "      <td>32.405556</td>\n",
       "      <td>33.394074</td>\n",
       "      <td>32.691685</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.436530</td>\n",
       "      <td>0.679331</td>\n",
       "      <td>0.625340</td>\n",
       "      <td>0.587080</td>\n",
       "      <td>0.846296</td>\n",
       "      <td>0.846296</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>26.537037</td>\n",
       "      <td>29.579630</td>\n",
       "      <td>30.361481</td>\n",
       "      <td>29.504944</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.421563</td>\n",
       "      <td>0.663536</td>\n",
       "      <td>0.605718</td>\n",
       "      <td>0.574324</td>\n",
       "      <td>0.815608</td>\n",
       "      <td>0.815608</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>24.111111</td>\n",
       "      <td>27.203704</td>\n",
       "      <td>28.050556</td>\n",
       "      <td>27.133370</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.383755</td>\n",
       "      <td>0.593645</td>\n",
       "      <td>0.556330</td>\n",
       "      <td>0.529415</td>\n",
       "      <td>0.754652</td>\n",
       "      <td>0.758675</td>\n",
       "      <td>0.905710</td>\n",
       "      <td>0.906867</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>18.759259</td>\n",
       "      <td>20.633333</td>\n",
       "      <td>21.314815</td>\n",
       "      <td>20.536315</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.368490</td>\n",
       "      <td>0.570477</td>\n",
       "      <td>0.544948</td>\n",
       "      <td>0.513365</td>\n",
       "      <td>0.753131</td>\n",
       "      <td>0.754365</td>\n",
       "      <td>0.889065</td>\n",
       "      <td>0.889065</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>16.055556</td>\n",
       "      <td>17.098148</td>\n",
       "      <td>17.577963</td>\n",
       "      <td>16.850463</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1        0.455497  0.688954  0.657722  0.611513      0.850000   \n",
       "1    uniform_10        0.457219  0.690652  0.656046  0.611886      0.857716   \n",
       "2    uniform_25        0.460114  0.690334  0.659018  0.610760      0.858951   \n",
       "3    uniform_50        0.457625  0.691289  0.658722  0.612326      0.850485   \n",
       "4   uniform_100        0.457333  0.688052  0.657896  0.612444      0.850044   \n",
       "0   uniform_250        0.454676  0.690524  0.652513  0.607058      0.860516   \n",
       "1   uniform_500        0.444291  0.681340  0.631334  0.596277      0.844753   \n",
       "2   uniform_750        0.436530  0.679331  0.625340  0.587080      0.846296   \n",
       "5  uniform_1000        0.421563  0.663536  0.605718  0.574324      0.815608   \n",
       "6  uniform_2000        0.383755  0.593645  0.556330  0.529415      0.754652   \n",
       "7  uniform_3000        0.368490  0.570477  0.544948  0.513365      0.753131   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.852557  0.922840  0.925397       0.750889    65.759259     70.692593   \n",
       "1   0.860565  0.922840  0.925689       0.750889    53.055556     55.650000   \n",
       "2   0.860634  0.924691  0.926375       0.750889    46.166667     49.550000   \n",
       "3   0.852169  0.925485  0.927169       0.750889    43.129630     45.505556   \n",
       "4   0.851728  0.925485  0.927169       0.750889    39.462963     41.918519   \n",
       "0   0.860516  0.943673  0.943673       0.754593    33.185185     36.896296   \n",
       "1   0.844753  0.958333  0.958333       0.754593    29.462963     32.405556   \n",
       "2   0.846296  0.950617  0.950617       0.754593    26.537037     29.579630   \n",
       "5   0.815608  0.950617  0.950617       0.750889    24.111111     27.203704   \n",
       "6   0.758675  0.905710  0.906867       0.750889    18.759259     20.633333   \n",
       "7   0.754365  0.889065  0.889065       0.750889    16.055556     17.098148   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      74.533333  73.139389  \n",
       "1      58.687407  57.559870  \n",
       "2      52.143333  51.190611  \n",
       "3      47.872037  46.984296  \n",
       "4      43.905185  43.076444  \n",
       "0      38.307407  37.609278  \n",
       "1      33.394074  32.691685  \n",
       "2      30.361481  29.504944  \n",
       "5      28.050556  27.133370  \n",
       "6      21.314815  20.536315  \n",
       "7      17.577963  16.850463  "
      ]
     },
     "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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.400824</td>\n",
       "      <td>0.425659</td>\n",
       "      <td>0.453466</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.340971</td>\n",
       "      <td>0.351744</td>\n",
       "      <td>0.0</td>\n",
       "      <td>65.435817</td>\n",
       "      <td>71.118438</td>\n",
       "      <td>74.046887</td>\n",
       "      <td>72.612341</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_10</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.399779</td>\n",
       "      <td>0.425007</td>\n",
       "      <td>0.452798</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.340252</td>\n",
       "      <td>0.351191</td>\n",
       "      <td>0.0</td>\n",
       "      <td>52.039971</td>\n",
       "      <td>56.304799</td>\n",
       "      <td>58.404444</td>\n",
       "      <td>57.268624</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.398953</td>\n",
       "      <td>0.424126</td>\n",
       "      <td>0.452151</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.339520</td>\n",
       "      <td>0.350502</td>\n",
       "      <td>0.0</td>\n",
       "      <td>46.336676</td>\n",
       "      <td>49.993123</td>\n",
       "      <td>51.876050</td>\n",
       "      <td>50.917969</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>uniform_50</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.398091</td>\n",
       "      <td>0.423260</td>\n",
       "      <td>0.451413</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.338807</td>\n",
       "      <td>0.349741</td>\n",
       "      <td>0.0</td>\n",
       "      <td>42.690115</td>\n",
       "      <td>45.927808</td>\n",
       "      <td>47.573274</td>\n",
       "      <td>46.675498</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>uniform_100</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.394788</td>\n",
       "      <td>0.420355</td>\n",
       "      <td>0.448464</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.335730</td>\n",
       "      <td>0.346737</td>\n",
       "      <td>0.0</td>\n",
       "      <td>39.240258</td>\n",
       "      <td>42.217837</td>\n",
       "      <td>43.711490</td>\n",
       "      <td>42.844141</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>uniform_250</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.387861</td>\n",
       "      <td>0.412696</td>\n",
       "      <td>0.441257</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.328845</td>\n",
       "      <td>0.339812</td>\n",
       "      <td>0.0</td>\n",
       "      <td>34.629513</td>\n",
       "      <td>37.015287</td>\n",
       "      <td>38.195284</td>\n",
       "      <td>37.384855</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>uniform_500</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.373959</td>\n",
       "      <td>0.398772</td>\n",
       "      <td>0.428734</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.316276</td>\n",
       "      <td>0.327443</td>\n",
       "      <td>0.0</td>\n",
       "      <td>30.188682</td>\n",
       "      <td>32.125129</td>\n",
       "      <td>33.118317</td>\n",
       "      <td>32.445901</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>uniform_750</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.365723</td>\n",
       "      <td>0.389993</td>\n",
       "      <td>0.420990</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.308556</td>\n",
       "      <td>0.319766</td>\n",
       "      <td>0.0</td>\n",
       "      <td>27.420057</td>\n",
       "      <td>29.093968</td>\n",
       "      <td>29.898242</td>\n",
       "      <td>29.234465</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>uniform_1000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.359918</td>\n",
       "      <td>0.383584</td>\n",
       "      <td>0.415319</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.303860</td>\n",
       "      <td>0.314943</td>\n",
       "      <td>0.0</td>\n",
       "      <td>25.271347</td>\n",
       "      <td>26.715774</td>\n",
       "      <td>27.406635</td>\n",
       "      <td>26.782411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>uniform_2000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.333851</td>\n",
       "      <td>0.357203</td>\n",
       "      <td>0.390378</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.280669</td>\n",
       "      <td>0.291872</td>\n",
       "      <td>0.0</td>\n",
       "      <td>19.515616</td>\n",
       "      <td>20.455186</td>\n",
       "      <td>20.911254</td>\n",
       "      <td>20.361803</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>uniform_3000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.313844</td>\n",
       "      <td>0.336802</td>\n",
       "      <td>0.370560</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.263599</td>\n",
       "      <td>0.274841</td>\n",
       "      <td>0.0</td>\n",
       "      <td>16.043696</td>\n",
       "      <td>16.724212</td>\n",
       "      <td>17.021218</td>\n",
       "      <td>16.539694</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0     uniform_1             0.0  0.400824  0.425659  0.453466           0.0   \n",
       "1    uniform_10             0.0  0.399779  0.425007  0.452798           0.0   \n",
       "2    uniform_25             0.0  0.398953  0.424126  0.452151           0.0   \n",
       "3    uniform_50             0.0  0.398091  0.423260  0.451413           0.0   \n",
       "4   uniform_100             0.0  0.394788  0.420355  0.448464           0.0   \n",
       "0   uniform_250             0.0  0.387861  0.412696  0.441257           0.0   \n",
       "1   uniform_500             0.0  0.373959  0.398772  0.428734           0.0   \n",
       "2   uniform_750             0.0  0.365723  0.389993  0.420990           0.0   \n",
       "5  uniform_1000             0.0  0.359918  0.383584  0.415319           0.0   \n",
       "6  uniform_2000             0.0  0.333851  0.357203  0.390378           0.0   \n",
       "7  uniform_3000             0.0  0.313844  0.336802  0.370560           0.0   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0        0.0  0.340971  0.351744            0.0    65.435817     71.118438   \n",
       "1        0.0  0.340252  0.351191            0.0    52.039971     56.304799   \n",
       "2        0.0  0.339520  0.350502            0.0    46.336676     49.993123   \n",
       "3        0.0  0.338807  0.349741            0.0    42.690115     45.927808   \n",
       "4        0.0  0.335730  0.346737            0.0    39.240258     42.217837   \n",
       "0        0.0  0.328845  0.339812            0.0    34.629513     37.015287   \n",
       "1        0.0  0.316276  0.327443            0.0    30.188682     32.125129   \n",
       "2        0.0  0.308556  0.319766            0.0    27.420057     29.093968   \n",
       "5        0.0  0.303860  0.314943            0.0    25.271347     26.715774   \n",
       "6        0.0  0.280669  0.291872            0.0    19.515616     20.455186   \n",
       "7        0.0  0.263599  0.274841            0.0    16.043696     16.724212   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      74.046887  72.612341  \n",
       "1      58.404444  57.268624  \n",
       "2      51.876050  50.917969  \n",
       "3      47.573274  46.675498  \n",
       "4      43.711490  42.844141  \n",
       "0      38.195284  37.384855  \n",
       "1      33.118317  32.445901  \n",
       "2      29.898242  29.234465  \n",
       "5      27.406635  26.782411  \n",
       "6      20.911254  20.361803  \n",
       "7      17.021218  16.539694  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import shutil\n",
    "for t in ['2019', '2020', 'dev']:\n",
    "    add_dir = \"/nfs/staticPruning/results/repro/supplement_uniform_%s/\" % t\n",
    "    main_dir = \"/nfs/staticPruning/results/repro/uniform_%s/\" % t\n",
    "    newdf = pd.concat([\n",
    "        pd.read_pickle(add_dir + \"results_df.pkl\"),\n",
    "        pd.read_pickle(main_dir + \"results_df.pkl\"),\n",
    "    ]).iloc[[3,4,5,6,7,0,1,2,8,9,10]]\n",
    "    display(newdf)\n",
    "    os.rename(main_dir + \"results_df.pkl\", main_dir + \"results_df.backup.pkl\")\n",
    "    newdf.to_pickle(main_dir + \"results_df.pkl\")\n",
    "    for x in [250,500,750]:\n",
    "        shutil.copyfile(add_dir + \"uniform_%d.res.gz\" %x, main_dir+ \"uniform_%d.res.gz\" %x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "global-former",
   "metadata": {},
   "source": [
    "## IDF document-centric pruning \n",
    "\n",
    "We now turn to implement the transformers for the document-centric pruning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "demonstrated-weather",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from typing import Tuple, Sequence\n",
    "\n",
    "def prune_sorted_natural(x, *Ys):\n",
    "    assert len(x.shape)==1\n",
    "    \n",
    "    indices = torch.sort(x, descending=True).indices\n",
    "\n",
    "    # concatenate neg and pos sides\n",
    "    rtr = [\n",
    "        x[indices]\n",
    "    ]\n",
    "    \n",
    "    # apply same slicing for other arrays\n",
    "    for y in Ys:        \n",
    "        y_sorted = y[indices]\n",
    "        rtr.append(y_sorted)\n",
    "    return rtr\n",
    "\n",
    "def doc_centric_idf(tokenids : torch.Tensor, embs : torch.Tensor, num_prune) ->  Tuple[Sequence[int], torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Returns the tokenids and embs after remvoing num_prune embeddings\n",
    "    \"\"\"\n",
    "       \n",
    "    length = len(tokenids)\n",
    "    newmask = torch.BoolTensor(length)\n",
    "    newmask[:] = True\n",
    "    embs = embs[0:length, :]\n",
    "    \n",
    "    dfs = nn_term.lookup[tokenids]\n",
    "    \n",
    "    #sort the embeddings according to pruning power\n",
    "    sorted_dfs, pos = prune_sorted_natural(\n",
    "        dfs, \n",
    "        torch.LongTensor(list(range(length)))\n",
    "    )\n",
    "\n",
    "    # prune the requisite numbers\n",
    "    newmask[pos[0:num_prune]] = False\n",
    "    \n",
    "    return tokenids[newmask], embs[newmask,:]\n",
    "\n",
    "\n",
    "def prune_doccentric_transformerH0(factory, topK, active_ids=None, min_obs=0):\n",
    "    assert active_ids is None \n",
    "    #if active_ids is None:\n",
    "    #    active_ids = list(range(nn_term.lookup.shape[0]))\n",
    "        \n",
    "        \n",
    "    \n",
    "    def _prune(df):\n",
    "        #pt.tqdm.pandas(desc='prune_doccentric')\n",
    "        #progress_apply\n",
    "        df[[\"doc_toks\", \"doc_embs\"]] = df.apply(\n",
    "            lambda row: doc_centric_idf(row.doc_toks, row.doc_embs, topK),\n",
    "            #lambda row: tc.ignore_prune(row.doc_toks, row.doc_embs, topK, pruning='global', blacklist=blacklist), \n",
    "            result_type=\"expand\", \n",
    "            axis=1)\n",
    "        return df\n",
    "    return pt.apply.generic(_prune)\n",
    "\n",
    "def prune_doccentric_pipeH0(topK, approx=True):\n",
    "    first_pass = (factory.ann_retrieve_score() % 1000) if approx else factory.set_retrieve()\n",
    "    return ( first_pass\n",
    "        >> factory.fetch_index_encodings(ids=True, verbose=False) \n",
    "        >> prune_doccentric_transformerH0(factory, topK)\n",
    "        >> factory.scorer(verbose=False))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "local-raising",
   "metadata": {},
   "source": [
    "check that the document-centric pruning works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "ideal-dancing",
   "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",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>1</td>\n",
       "      <td>7605154</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1</td>\n",
       "      <td>5453527</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>121</th>\n",
       "      <td>1</td>\n",
       "      <td>1833252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>1</td>\n",
       "      <td>7575552</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52</th>\n",
       "      <td>1</td>\n",
       "      <td>1525431</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1701</th>\n",
       "      <td>1</td>\n",
       "      <td>2130934</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1889</th>\n",
       "      <td>1</td>\n",
       "      <td>5812915</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1732</th>\n",
       "      <td>1</td>\n",
       "      <td>6823853</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1995</th>\n",
       "      <td>1</td>\n",
       "      <td>976146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>312</th>\n",
       "      <td>1</td>\n",
       "      <td>5436251</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     qid    docno\n",
       "11     1  7605154\n",
       "8      1  5453527\n",
       "121    1  1833252\n",
       "19     1  7575552\n",
       "52     1  1525431\n",
       "...   ..      ...\n",
       "1701   1  2130934\n",
       "1889   1  5812915\n",
       "1732   1  6823853\n",
       "1995   1   976146\n",
       "312    1  5436251\n",
       "\n",
       "[1000 rows x 2 columns]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prune_doccentric_pipeH0(10).search(\"chemical reactions\")[['qid', 'docno']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "combined-management",
   "metadata": {},
   "source": [
    "now conduct the experiments for document-centric pruning. \n",
    "\n",
    "`DOC_CENTRIC_THRESHOLDS` corresponds to $\\tau_d$ in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "outer-temperature",
   "metadata": {},
   "outputs": [],
   "source": [
    "DOC_CENTRIC_THRESHOLDS=[1,3,5,8,10,15,20,25,30]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "plastic-assist",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97fb0b0699894adab0f0f6a4f36dca41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=27.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>R(rel=2)@1000</th>\n",
       "      <th>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>doc_centric_1</td>\n",
       "      <td>0.405220</td>\n",
       "      <td>0.692956</td>\n",
       "      <td>0.649440</td>\n",
       "      <td>0.584247</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853659</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>66.883721</td>\n",
       "      <td>75.690698</td>\n",
       "      <td>78.282326</td>\n",
       "      <td>76.001302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>doc_centric_3</td>\n",
       "      <td>0.408684</td>\n",
       "      <td>0.698561</td>\n",
       "      <td>0.652257</td>\n",
       "      <td>0.586583</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865287</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>66.581395</td>\n",
       "      <td>74.383721</td>\n",
       "      <td>76.500930</td>\n",
       "      <td>74.001302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>doc_centric_5</td>\n",
       "      <td>0.411256</td>\n",
       "      <td>0.701147</td>\n",
       "      <td>0.654258</td>\n",
       "      <td>0.586850</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865292</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>64.581395</td>\n",
       "      <td>72.330233</td>\n",
       "      <td>74.726512</td>\n",
       "      <td>72.001302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>doc_centric_8</td>\n",
       "      <td>0.400595</td>\n",
       "      <td>0.694253</td>\n",
       "      <td>0.647600</td>\n",
       "      <td>0.580938</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865294</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>61.813953</td>\n",
       "      <td>69.530233</td>\n",
       "      <td>71.956512</td>\n",
       "      <td>69.001302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>doc_centric_10</td>\n",
       "      <td>0.400391</td>\n",
       "      <td>0.690964</td>\n",
       "      <td>0.645815</td>\n",
       "      <td>0.581171</td>\n",
       "      <td>0.868992</td>\n",
       "      <td>0.869171</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>59.744186</td>\n",
       "      <td>68.290698</td>\n",
       "      <td>70.318837</td>\n",
       "      <td>67.001302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>doc_centric_15</td>\n",
       "      <td>0.401108</td>\n",
       "      <td>0.691849</td>\n",
       "      <td>0.651776</td>\n",
       "      <td>0.581463</td>\n",
       "      <td>0.868992</td>\n",
       "      <td>0.869188</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>56.488372</td>\n",
       "      <td>65.255814</td>\n",
       "      <td>66.801860</td>\n",
       "      <td>62.001395</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>doc_centric_20</td>\n",
       "      <td>0.385735</td>\n",
       "      <td>0.680996</td>\n",
       "      <td>0.636953</td>\n",
       "      <td>0.569487</td>\n",
       "      <td>0.863178</td>\n",
       "      <td>0.863398</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>54.674419</td>\n",
       "      <td>63.872093</td>\n",
       "      <td>63.850698</td>\n",
       "      <td>57.002837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>doc_centric_25</td>\n",
       "      <td>0.370902</td>\n",
       "      <td>0.658937</td>\n",
       "      <td>0.626693</td>\n",
       "      <td>0.555300</td>\n",
       "      <td>0.839922</td>\n",
       "      <td>0.840106</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>54.906977</td>\n",
       "      <td>61.834884</td>\n",
       "      <td>61.062093</td>\n",
       "      <td>52.013651</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>doc_centric_30</td>\n",
       "      <td>0.346869</td>\n",
       "      <td>0.641218</td>\n",
       "      <td>0.602858</td>\n",
       "      <td>0.534962</td>\n",
       "      <td>0.787597</td>\n",
       "      <td>0.788548</td>\n",
       "      <td>0.69536</td>\n",
       "      <td>55.883721</td>\n",
       "      <td>59.676744</td>\n",
       "      <td>58.417674</td>\n",
       "      <td>47.050721</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0   doc_centric_1        0.405220  0.692956  0.649440  0.584247      0.853488   \n",
       "1   doc_centric_3        0.408684  0.698561  0.652257  0.586583      0.865116   \n",
       "2   doc_centric_5        0.411256  0.701147  0.654258  0.586850      0.865116   \n",
       "3   doc_centric_8        0.400595  0.694253  0.647600  0.580938      0.865116   \n",
       "4  doc_centric_10        0.400391  0.690964  0.645815  0.581171      0.868992   \n",
       "5  doc_centric_15        0.401108  0.691849  0.651776  0.581463      0.868992   \n",
       "6  doc_centric_20        0.385735  0.680996  0.636953  0.569487      0.863178   \n",
       "7  doc_centric_25        0.370902  0.658937  0.626693  0.555300      0.839922   \n",
       "8  doc_centric_30        0.346869  0.641218  0.602858  0.534962      0.787597   \n",
       "\n",
       "   RR(rel=2)  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  AvgDoclen@100  \\\n",
       "0   0.853659        0.69536    66.883721     75.690698      78.282326   \n",
       "1   0.865287        0.69536    66.581395     74.383721      76.500930   \n",
       "2   0.865292        0.69536    64.581395     72.330233      74.726512   \n",
       "3   0.865294        0.69536    61.813953     69.530233      71.956512   \n",
       "4   0.869171        0.69536    59.744186     68.290698      70.318837   \n",
       "5   0.869188        0.69536    56.488372     65.255814      66.801860   \n",
       "6   0.863398        0.69536    54.674419     63.872093      63.850698   \n",
       "7   0.840106        0.69536    54.906977     61.834884      61.062093   \n",
       "8   0.788548        0.69536    55.883721     59.676744      58.417674   \n",
       "\n",
       "   AvgDoclen  \n",
       "0  76.001302  \n",
       "1  74.001302  \n",
       "2  72.001302  \n",
       "3  69.001302  \n",
       "4  67.001302  \n",
       "5  62.001395  \n",
       "6  57.002837  \n",
       "7  52.013651  \n",
       "8  47.050721  "
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doccentric_2019 = do_exp(\n",
    "    [prune_doccentric_pipeH0(x) for x in DOC_CENTRIC_THRESHOLDS], \n",
    "    topics='2019',\n",
    "    names=[\"doc_centric_%d\" % x for x in DOC_CENTRIC_THRESHOLDS],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/doccentric_%s/\"\n",
    ")\n",
    "doccentric_2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "resident-incident",
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d60663fa24cb440ca699c6e7fdbc0deb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=27.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>R(rel=2)@1000</th>\n",
       "      <th>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>doc_centric_1</td>\n",
       "      <td>0.455460</td>\n",
       "      <td>0.689148</td>\n",
       "      <td>0.657294</td>\n",
       "      <td>0.611809</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>67.981481</td>\n",
       "      <td>73.320370</td>\n",
       "      <td>77.244815</td>\n",
       "      <td>75.993000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>doc_centric_3</td>\n",
       "      <td>0.456174</td>\n",
       "      <td>0.689630</td>\n",
       "      <td>0.658862</td>\n",
       "      <td>0.611955</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>66.129630</td>\n",
       "      <td>71.200000</td>\n",
       "      <td>75.377778</td>\n",
       "      <td>73.993000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>doc_centric_5</td>\n",
       "      <td>0.456169</td>\n",
       "      <td>0.689525</td>\n",
       "      <td>0.660041</td>\n",
       "      <td>0.611555</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>64.129630</td>\n",
       "      <td>69.735185</td>\n",
       "      <td>73.578704</td>\n",
       "      <td>71.993000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>doc_centric_8</td>\n",
       "      <td>0.454840</td>\n",
       "      <td>0.691973</td>\n",
       "      <td>0.658627</td>\n",
       "      <td>0.611849</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852557</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>61.500000</td>\n",
       "      <td>67.179630</td>\n",
       "      <td>71.000556</td>\n",
       "      <td>68.993000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>doc_centric_10</td>\n",
       "      <td>0.457086</td>\n",
       "      <td>0.694662</td>\n",
       "      <td>0.659177</td>\n",
       "      <td>0.612022</td>\n",
       "      <td>0.850000</td>\n",
       "      <td>0.852412</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>59.500000</td>\n",
       "      <td>65.453704</td>\n",
       "      <td>69.422778</td>\n",
       "      <td>66.993000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>doc_centric_15</td>\n",
       "      <td>0.451205</td>\n",
       "      <td>0.692655</td>\n",
       "      <td>0.652491</td>\n",
       "      <td>0.608864</td>\n",
       "      <td>0.838889</td>\n",
       "      <td>0.840046</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>59.888889</td>\n",
       "      <td>63.144444</td>\n",
       "      <td>66.421481</td>\n",
       "      <td>61.993185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>doc_centric_20</td>\n",
       "      <td>0.450207</td>\n",
       "      <td>0.689134</td>\n",
       "      <td>0.652400</td>\n",
       "      <td>0.605126</td>\n",
       "      <td>0.847531</td>\n",
       "      <td>0.849074</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>58.185185</td>\n",
       "      <td>61.144444</td>\n",
       "      <td>63.957222</td>\n",
       "      <td>56.997167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>doc_centric_25</td>\n",
       "      <td>0.438444</td>\n",
       "      <td>0.675905</td>\n",
       "      <td>0.642909</td>\n",
       "      <td>0.594756</td>\n",
       "      <td>0.833025</td>\n",
       "      <td>0.833025</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>56.351852</td>\n",
       "      <td>58.292593</td>\n",
       "      <td>61.602593</td>\n",
       "      <td>52.018222</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>doc_centric_30</td>\n",
       "      <td>0.423784</td>\n",
       "      <td>0.666966</td>\n",
       "      <td>0.632528</td>\n",
       "      <td>0.584582</td>\n",
       "      <td>0.835391</td>\n",
       "      <td>0.836366</td>\n",
       "      <td>0.750889</td>\n",
       "      <td>51.962963</td>\n",
       "      <td>56.733333</td>\n",
       "      <td>58.869630</td>\n",
       "      <td>47.066019</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0   doc_centric_1        0.455460  0.689148  0.657294  0.611809      0.850000   \n",
       "1   doc_centric_3        0.456174  0.689630  0.658862  0.611955      0.850000   \n",
       "2   doc_centric_5        0.456169  0.689525  0.660041  0.611555      0.850000   \n",
       "3   doc_centric_8        0.454840  0.691973  0.658627  0.611849      0.850000   \n",
       "4  doc_centric_10        0.457086  0.694662  0.659177  0.612022      0.850000   \n",
       "5  doc_centric_15        0.451205  0.692655  0.652491  0.608864      0.838889   \n",
       "6  doc_centric_20        0.450207  0.689134  0.652400  0.605126      0.847531   \n",
       "7  doc_centric_25        0.438444  0.675905  0.642909  0.594756      0.833025   \n",
       "8  doc_centric_30        0.423784  0.666966  0.632528  0.584582      0.835391   \n",
       "\n",
       "   RR(rel=2)  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  AvgDoclen@100  \\\n",
       "0   0.852557       0.750889    67.981481     73.320370      77.244815   \n",
       "1   0.852557       0.750889    66.129630     71.200000      75.377778   \n",
       "2   0.852557       0.750889    64.129630     69.735185      73.578704   \n",
       "3   0.852557       0.750889    61.500000     67.179630      71.000556   \n",
       "4   0.852412       0.750889    59.500000     65.453704      69.422778   \n",
       "5   0.840046       0.750889    59.888889     63.144444      66.421481   \n",
       "6   0.849074       0.750889    58.185185     61.144444      63.957222   \n",
       "7   0.833025       0.750889    56.351852     58.292593      61.602593   \n",
       "8   0.836366       0.750889    51.962963     56.733333      58.869630   \n",
       "\n",
       "   AvgDoclen  \n",
       "0  75.993000  \n",
       "1  73.993000  \n",
       "2  71.993000  \n",
       "3  68.993000  \n",
       "4  66.993000  \n",
       "5  61.993185  \n",
       "6  56.997167  \n",
       "7  52.018222  \n",
       "8  47.066019  "
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doccentric_2020 = do_exp(\n",
    "    [prune_doccentric_pipeH0(x) for x in DOC_CENTRIC_THRESHOLDS], \n",
    "    topics='2020',\n",
    "    names=[\"doc_centric_%d\" % x for x in DOC_CENTRIC_THRESHOLDS],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/doccentric_%s/\"\n",
    ")\n",
    "doccentric_2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "worth-sessions",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "35a569a215ed42b6ad7d6afb604ffef6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='pt.Experiment'), FloatProgress(value=0.0, max=3141.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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>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>AvgDoclen@1</th>\n",
       "      <th>AvgDoclen@10</th>\n",
       "      <th>AvgDoclen@100</th>\n",
       "      <th>AvgDoclen</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>doc_centric_1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.401628</td>\n",
       "      <td>0.426198</td>\n",
       "      <td>0.454013</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.341719</td>\n",
       "      <td>0.352404</td>\n",
       "      <td>0.0</td>\n",
       "      <td>68.102006</td>\n",
       "      <td>73.984928</td>\n",
       "      <td>77.043357</td>\n",
       "      <td>75.502952</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>doc_centric_3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.401501</td>\n",
       "      <td>0.425715</td>\n",
       "      <td>0.453846</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.341451</td>\n",
       "      <td>0.352113</td>\n",
       "      <td>0.0</td>\n",
       "      <td>66.298567</td>\n",
       "      <td>72.158453</td>\n",
       "      <td>75.169029</td>\n",
       "      <td>73.502952</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>doc_centric_5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.401206</td>\n",
       "      <td>0.425598</td>\n",
       "      <td>0.453503</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.341014</td>\n",
       "      <td>0.351669</td>\n",
       "      <td>0.0</td>\n",
       "      <td>64.690115</td>\n",
       "      <td>70.373238</td>\n",
       "      <td>73.296202</td>\n",
       "      <td>71.502952</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>doc_centric_8</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.401186</td>\n",
       "      <td>0.425374</td>\n",
       "      <td>0.453416</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.340853</td>\n",
       "      <td>0.351475</td>\n",
       "      <td>0.0</td>\n",
       "      <td>62.813467</td>\n",
       "      <td>67.863926</td>\n",
       "      <td>70.646749</td>\n",
       "      <td>68.502958</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>doc_centric_10</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.399992</td>\n",
       "      <td>0.424304</td>\n",
       "      <td>0.452097</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.339449</td>\n",
       "      <td>0.350034</td>\n",
       "      <td>0.0</td>\n",
       "      <td>61.748138</td>\n",
       "      <td>66.401862</td>\n",
       "      <td>69.049629</td>\n",
       "      <td>66.502966</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>doc_centric_15</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.392814</td>\n",
       "      <td>0.417930</td>\n",
       "      <td>0.446841</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.333808</td>\n",
       "      <td>0.344826</td>\n",
       "      <td>0.0</td>\n",
       "      <td>60.719628</td>\n",
       "      <td>64.008983</td>\n",
       "      <td>65.844202</td>\n",
       "      <td>61.503098</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>doc_centric_20</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.384567</td>\n",
       "      <td>0.410787</td>\n",
       "      <td>0.439235</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.326755</td>\n",
       "      <td>0.338073</td>\n",
       "      <td>0.0</td>\n",
       "      <td>60.881948</td>\n",
       "      <td>62.390616</td>\n",
       "      <td>63.108341</td>\n",
       "      <td>56.505853</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>doc_centric_25</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.370700</td>\n",
       "      <td>0.395180</td>\n",
       "      <td>0.425006</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.313317</td>\n",
       "      <td>0.324463</td>\n",
       "      <td>0.0</td>\n",
       "      <td>60.200430</td>\n",
       "      <td>60.563883</td>\n",
       "      <td>60.232650</td>\n",
       "      <td>51.521634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>doc_centric_30</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.354905</td>\n",
       "      <td>0.379065</td>\n",
       "      <td>0.408670</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.301130</td>\n",
       "      <td>0.312367</td>\n",
       "      <td>0.0</td>\n",
       "      <td>59.107450</td>\n",
       "      <td>58.461304</td>\n",
       "      <td>57.304550</td>\n",
       "      <td>46.563818</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             name  AP(rel=2)@1000   nDCG@10   nDCG@20  nDCG@100  RR(rel=2)@10  \\\n",
       "0   doc_centric_1             0.0  0.401628  0.426198  0.454013           0.0   \n",
       "1   doc_centric_3             0.0  0.401501  0.425715  0.453846           0.0   \n",
       "2   doc_centric_5             0.0  0.401206  0.425598  0.453503           0.0   \n",
       "3   doc_centric_8             0.0  0.401186  0.425374  0.453416           0.0   \n",
       "4  doc_centric_10             0.0  0.399992  0.424304  0.452097           0.0   \n",
       "5  doc_centric_15             0.0  0.392814  0.417930  0.446841           0.0   \n",
       "6  doc_centric_20             0.0  0.384567  0.410787  0.439235           0.0   \n",
       "7  doc_centric_25             0.0  0.370700  0.395180  0.425006           0.0   \n",
       "8  doc_centric_30             0.0  0.354905  0.379065  0.408670           0.0   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0        0.0  0.341719  0.352404            0.0    68.102006     73.984928   \n",
       "1        0.0  0.341451  0.352113            0.0    66.298567     72.158453   \n",
       "2        0.0  0.341014  0.351669            0.0    64.690115     70.373238   \n",
       "3        0.0  0.340853  0.351475            0.0    62.813467     67.863926   \n",
       "4        0.0  0.339449  0.350034            0.0    61.748138     66.401862   \n",
       "5        0.0  0.333808  0.344826            0.0    60.719628     64.008983   \n",
       "6        0.0  0.326755  0.338073            0.0    60.881948     62.390616   \n",
       "7        0.0  0.313317  0.324463            0.0    60.200430     60.563883   \n",
       "8        0.0  0.301130  0.312367            0.0    59.107450     58.461304   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      77.043357  75.502952  \n",
       "1      75.169029  73.502952  \n",
       "2      73.296202  71.502952  \n",
       "3      70.646749  68.502958  \n",
       "4      69.049629  66.502966  \n",
       "5      65.844202  61.503098  \n",
       "6      63.108341  56.505853  \n",
       "7      60.232650  51.521634  \n",
       "8      57.304550  46.563818  "
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doccentric_dev = do_exp(\n",
    "    [prune_doccentric_pipeH0(x) for x in DOC_CENTRIC_THRESHOLDS], \n",
    "    topics='dev',\n",
    "    names=[\"doc_centric_%d\" % x for x in DOC_CENTRIC_THRESHOLDS],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/doccentric_%s/\"\n",
    ")\n",
    "doccentric_dev"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "union-parish",
   "metadata": {},
   "source": [
    "# uniform idf index generation\n",
    "\n",
    "We now turn to creating real indices. After index generation has completed, the faiss_index.py script within ColBERT must be executed, using the `--sample 0.05` setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "elementary-divorce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prune_index_to_disk_by_idf(factory, new_index_path, num_tokens, flush):\n",
    "    import torch\n",
    "    \n",
    "    removed_tids = get_tid_ordered_by_idf(factory.nn_term(), verbose=False)[3:num_tokens+3]\n",
    "    #removed_tids = torch.Tensor(removed_tids)\n",
    "    removed_tids = set(removed_tids)\n",
    "    \n",
    "    import os, torch\n",
    "    embs_to_save = []\n",
    "    ids_to_save = []\n",
    "    shard=0\n",
    "    shard_size=0\n",
    "    \n",
    "    def _save_shard(shard, all_embs, all_ids):\n",
    "        import ujson\n",
    "        torch.save(torch.cat(all_embs), os.path.join(new_index_path, \"%d.pt\" % shard))\n",
    "        torch.save(torch.cat(all_ids), os.path.join(new_index_path, \"%d.tokenids\" % shard))\n",
    "        doclens_path = os.path.join(new_index_path, \"doclens.%d.json\" % shard)\n",
    "        doclens = [d.size(0) for d in all_embs]\n",
    "        with open(doclens_path, 'w') as output_doclens:\n",
    "            ujson.dump(doclens, output_doclens)\n",
    "        \n",
    "    \n",
    "    mask = torch.BoolTensor(180)\n",
    "    mask[:] = False\n",
    "    numdocs = len(factory)\n",
    "\n",
    "    for pid in pt.tqdm(range(numdocs), total=numdocs, unit='d'):        \n",
    "        ids = factory.nn_term().get_tokens_for_doc(pid)\n",
    "        num_toks = ids.size(0)\n",
    "        embs = factory._rrm().get_embedding(pid).type(torch.float16)\n",
    "        embs = embs[0:num_toks,:]\n",
    "        \n",
    "        mask[0:num_toks] = torch.BoolTensor([x.item() in removed_tids for x in ids]) # torch.isin(ids, removed_tids)\n",
    "        #This is 5x faster says Antonio\n",
    "        #mask[0:num_tokens] = torch.any(ids[None, : ] == removed_tids[:, None], axis=0)\n",
    "\n",
    "        newids = ids[~mask[:num_toks]]\n",
    "        newembs = embs[~mask[:num_toks], : ]\n",
    "        \n",
    "        #assert newids.shape[0] > 0, newids.shape\n",
    "        mask[:]  = False\n",
    "        \n",
    "        embs_to_save.append(newembs)\n",
    "        ids_to_save.append(newids)\n",
    "        shard_size += 1\n",
    "        \n",
    "        if shard_size == flush:\n",
    "            print(\"Flushing shard %d\" % shard)\n",
    "            _save_shard(shard, embs_to_save, ids_to_save)\n",
    "            shard += 1\n",
    "            shard_size = 0\n",
    "            embs_to_save=[]\n",
    "            ids_to_save=[]\n",
    "            print(\"Done\")\n",
    "            \n",
    "    if shard_size > 0:\n",
    "        _save_shard(shard, embs_to_save, ids_to_save)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "orange-merchant",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26f12a8cb9784493a32c7bd39e26d04d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8841823 [00:00<?, ?d/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flushing shard 0\n",
      "Done\n",
      "Flushing shard 1\n",
      "Done\n",
      "Flushing shard 2\n",
      "Done\n",
      "Flushing shard 3\n",
      "Done\n",
      "Flushing shard 4\n",
      "Done\n",
      "Flushing shard 5\n",
      "Done\n",
      "Flushing shard 6\n",
      "Done\n",
      "Flushing shard 7\n",
      "Done\n",
      "Flushing shard 8\n",
      "Done\n",
      "Flushing shard 9\n",
      "Done\n",
      "Flushing shard 10\n",
      "Done\n",
      "Flushing shard 11\n",
      "Done\n",
      "Flushing shard 12\n",
      "Done\n",
      "Flushing shard 13\n",
      "Done\n",
      "Flushing shard 14\n",
      "Done\n",
      "Flushing shard 15\n",
      "Done\n",
      "Flushing shard 16\n",
      "Done\n",
      "Flushing shard 17\n",
      "Done\n",
      "Flushing shard 18\n",
      "Done\n",
      "Flushing shard 19\n",
      "Done\n",
      "Flushing shard 20\n",
      "Done\n",
      "Flushing shard 21\n",
      "Done\n",
      "Flushing shard 22\n",
      "Done\n",
      "Flushing shard 23\n",
      "Done\n",
      "Flushing shard 24\n",
      "Done\n",
      "Flushing shard 25\n",
      "Done\n",
      "Flushing shard 26\n",
      "Done\n",
      "Flushing shard 27\n",
      "Done\n",
      "Flushing shard 28\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.makedirs(\"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/idf_uniform_100\", exist_ok=True)\n",
    "prune_index_to_disk_by_idf(\n",
    "    factory, \n",
    "    \"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/idf_uniform_100\", \n",
    "    100, \n",
    "    300_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "iraqi-netscape",
   "metadata": {},
   "source": [
    "# stopwords removed index generation\n",
    "\n",
    "The file `stopword-list.txt` corresponds to Terrier's stopwords file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "found-duplicate",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"stopword-list.txt\"):\n",
    "    !wget \"https://raw.githubusercontent.com/terrier-org/terrier-core/5.x/modules/core/src/main/resources/stopword-list.txt\"\n",
    "\n",
    "stops=[]\n",
    "with open(\"stopword-list.txt\") as f:\n",
    "    for l in f:\n",
    "        stops.append(l.strip())\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "stop_ids = [x for x in tokenizer.convert_tokens_to_ids(stops) if x != 100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "complimentary-democrat",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['.', 'the', ',', 'of', 'and', 'a', 'to', 'in', 'is', 'for']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.convert_ids_to_tokens([1012, 1996, 1010, 1997, 1998, 1037, 2000, 1999, 2003, 2005])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "advance-lindsay",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prune_index_to_disk_by_stops(factory, new_index_path, flush):\n",
    "    import torch\n",
    "    \n",
    "    removed_tids = stop_ids\n",
    "    #removed_tids = torch.Tensor(removed_tids)\n",
    "    removed_tids = set(removed_tids)\n",
    "    \n",
    "    import os, torch\n",
    "    embs_to_save = []\n",
    "    ids_to_save = []\n",
    "    shard=0\n",
    "    shard_size=0\n",
    "    \n",
    "    def _save_shard(shard, all_embs, all_ids):\n",
    "        import ujson\n",
    "        torch.save(torch.cat(all_embs), os.path.join(new_index_path, \"%d.pt\" % shard))\n",
    "        torch.save(torch.cat(all_ids), os.path.join(new_index_path, \"%d.tokenids\" % shard))\n",
    "        doclens_path = os.path.join(new_index_path, \"doclens.%d.json\" % shard)\n",
    "        doclens = [d.size(0) for d in all_embs]\n",
    "        with open(doclens_path, 'w') as output_doclens:\n",
    "            ujson.dump(doclens, output_doclens)\n",
    "        \n",
    "    \n",
    "    mask = torch.BoolTensor(180)\n",
    "    mask[:] = False\n",
    "    numdocs = len(factory)\n",
    "\n",
    "    for pid in pt.tqdm(range(numdocs), total=numdocs, unit='d'):        \n",
    "        ids = factory.nn_term().get_tokens_for_doc(pid)\n",
    "        num_toks = ids.size(0)\n",
    "        embs = factory._rrm().get_embedding(pid).type(torch.float16)\n",
    "        embs = embs[0:num_toks,:]\n",
    "        \n",
    "        mask[0:num_toks] = torch.BoolTensor([x.item() in removed_tids for x in ids]) # torch.isin(ids, removed_tids)\n",
    "        #This is 5x faster says Antonio\n",
    "        #mask[0:num_tokens] = torch.any(ids[None, : ] == removed_tids[:, None], axis=0)\n",
    "\n",
    "        newids = ids[~mask[:num_toks]]\n",
    "        newembs = embs[~mask[:num_toks], : ]\n",
    "        \n",
    "        #assert newids.shape[0] > 0, newids.shape\n",
    "        mask[:]  = False\n",
    "        \n",
    "        embs_to_save.append(newembs)\n",
    "        ids_to_save.append(newids)\n",
    "        shard_size += 1\n",
    "        \n",
    "        if shard_size == flush:\n",
    "            print(\"Flushing shard %d\" % shard)\n",
    "            _save_shard(shard, embs_to_save, ids_to_save)\n",
    "            shard += 1\n",
    "            shard_size = 0\n",
    "            embs_to_save=[]\n",
    "            ids_to_save=[]\n",
    "            print(\"Done\")\n",
    "            \n",
    "    if shard_size > 0:\n",
    "        _save_shard(shard, embs_to_save, ids_to_save)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "educational-messaging",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d785e6dae9c4bf88ae15dd2cde497db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8841823 [00:00<?, ?d/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flushing shard 0\n",
      "Done\n",
      "Flushing shard 1\n",
      "Done\n",
      "Flushing shard 2\n",
      "Done\n",
      "Flushing shard 3\n",
      "Done\n",
      "Flushing shard 4\n",
      "Done\n",
      "Flushing shard 5\n",
      "Done\n",
      "Flushing shard 6\n",
      "Done\n",
      "Flushing shard 7\n",
      "Done\n",
      "Flushing shard 8\n",
      "Done\n",
      "Flushing shard 14\n",
      "Done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.makedirs(\"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/stops\", exist_ok=True)\n",
    "prune_index_to_disk_by_stops(\n",
    "    factory, \n",
    "    \"/nfsindices/staticPrunedIndices/repro_paper_SIGIR2022/stops\",\n",
    "    300_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "welcome-candidate",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}