notebooks/baseline-pruning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "streaming-bench",
   "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": "simple-driver",
   "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": "popular-republican",
   "metadata": {},
   "source": [
    "## load the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "insured-metabolism",
   "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 10, 19:01:34] #> Loading model checkpoint.\n",
      "[Feb 10, 19:01:34] #> 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 10, 19:01:43] #> checkpoint['epoch'] = 0\n",
      "[Feb 10, 19:01:43] #> 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",
    "# from pyterrier_colbert.ranking import ColBERTFactory\n",
    "# factory = ColBERTFactory.from_dataset('vaswani', 'colbert_uog44k')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "future-interest",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Feb 10, 19:01:44] #> Loading the FAISS index from /nfs/indices/colbert_passage/index_name3/ivfpq.faiss ..\n",
      "[Feb 10, 19:01:55] #> Building the emb2pid mapping..\n",
      "[Feb 10, 19:02:17] len(self.emb2pid) = 687989391\n",
      "[Feb 10, 19:02:23] #> 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": "code",
   "execution_count": 5,
   "id": "athletic-domestic",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading reranking index, memtype=mem\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c4104640868d49f6899db713b02aeaf0",
       "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": "guilty-liquid",
   "metadata": {},
   "source": [
    "## index stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "still-letters",
   "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": "standing-communist",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8841823"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(factory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "amino-russell",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(77.8108)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total/len(factory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stylish-anderson",
   "metadata": {},
   "source": [
    "## eval setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "raising-palestinian",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --force-reinstall --no-deps git+https://github.com/terrierteam/ir_measures.git@runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "portable-carol",
   "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": "code",
   "execution_count": 10,
   "id": "young-literature",
   "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": "handed-synthesis",
   "metadata": {},
   "source": [
    "## stops uniform pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "arbitrary-battle",
   "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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "differential-jersey",
   "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_stop_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": 13,
   "id": "gross-scope",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": 14,
   "id": "subtle-florence",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading msmarco_passage topics to /home/me/.pyterrier/corpora/msmarco_passage/msmarco-test2019-queries.tsv.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b34ef89bb5f44826bfca679154f22dfd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "msmarco-test2019-queries.tsv.gz:   0%|          | 0.00/4.18k [00:00<?, ?iB/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading msmarco_passage qrels to /home/me/.pyterrier/corpora/msmarco_passage/2019qrels-docs.txt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e617d7d10704c7fb6db1f62a12edd41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "2019qrels-docs.txt:   0%|          | 0.00/183k [00:00<?, ?iB/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f7d6df15ce0a43258aecd718afba07ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/3 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\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>stops_uniform</td>\n",
       "      <td>0.398659</td>\n",
       "      <td>0.687096</td>\n",
       "      <td>0.638232</td>\n",
       "      <td>0.572497</td>\n",
       "      <td>0.830233</td>\n",
       "      <td>0.83041</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.928295</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>45.581395</td>\n",
       "      <td>52.155814</td>\n",
       "      <td>53.665116</td>\n",
       "      <td>52.245628</td>\n",
       "      <td>657.402888</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  stops_uniform        0.398659  0.687096  0.638232  0.572497      0.830233   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0    0.83041  0.928295  0.928295       0.695251    45.581395     52.155814   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen         mrt  \n",
       "0      53.665116  52.245628  657.402888  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stops_2019 = do_exp(\n",
    "    [get_uniform_stop_pipeline(factory, stop_ids)], \n",
    "    topics='2019',\n",
    "    names=[\"stops_uniform\" ],\n",
    "    save_dir=\"/nfs/craigm/staticPruning/results/repro/Cstops_uniform_%s/\"\n",
    ")\n",
    "stops_2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "equal-lithuania",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading msmarco_passage topics to /home/me/.pyterrier/corpora/msmarco_passage/msmarco-test2020-queries.tsv.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1368e2d220354881834edb2d0cb7add1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "msmarco-test2020-queries.tsv.gz:   0%|          | 0.00/4.03k [00:00<?, ?iB/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading msmarco_passage qrels to /home/me/.pyterrier/corpora/msmarco_passage/2020qrels-docs.txt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d8cafb5c034487abb30b0cf70629bb1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "2020qrels-docs.txt:   0%|          | 0.00/213k [00:00<?, ?iB/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16e895eb75684949bdf8c997bcb2943d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/3 [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>stops_uniform</td>\n",
       "      <td>0.454838</td>\n",
       "      <td>0.687692</td>\n",
       "      <td>0.653287</td>\n",
       "      <td>0.606888</td>\n",
       "      <td>0.86034</td>\n",
       "      <td>0.861764</td>\n",
       "      <td>0.916667</td>\n",
       "      <td>0.918091</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>45.407407</td>\n",
       "      <td>49.985185</td>\n",
       "      <td>52.892407</td>\n",
       "      <td>51.988426</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  stops_uniform        0.454838  0.687692  0.653287  0.606888       0.86034   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.861764  0.916667  0.918091       0.754593    45.407407     49.985185   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      52.892407  51.988426  "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stops_2019 = do_exp(\n",
    "    [get_uniform_stop_pipeline(factory, stop_ids)], \n",
    "    topics='2020',\n",
    "    names=[\"stops_uniform\" ],\n",
    "    save_dir=\"/nfs/craigm/staticPruning/results/repro/Cstops_uniform_%s/\"\n",
    ")\n",
    "stops_2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "junior-basketball",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3ab9827be184b3ab140a4dbc48d7000",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/349 [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>stops_uniform</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.391132</td>\n",
       "      <td>0.416073</td>\n",
       "      <td>0.444741</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.331263</td>\n",
       "      <td>0.342257</td>\n",
       "      <td>0.0</td>\n",
       "      <td>47.682235</td>\n",
       "      <td>51.46616</td>\n",
       "      <td>53.279136</td>\n",
       "      <td>51.901785</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  stops_uniform             0.0  0.391132  0.416073  0.444741           0.0   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0        0.0  0.331263  0.342257            0.0    47.682235      51.46616   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      53.279136  51.901785  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stops_dev = do_exp(\n",
    "    [get_uniform_stop_pipeline(factory, stop_ids)],\n",
    "    topics='dev',\n",
    "    names=[\"stops_uniform\" ],\n",
    "    save_dir=\"/nfs/staticPruning/results/repro/Cstops_uniform_%s/\"\n",
    ")\n",
    "stops_dev"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "swedish-devices",
   "metadata": {},
   "source": [
    "## random document-centrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "known-graphic",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from typing import Tuple, Sequence\n",
    "\n",
    "\n",
    "def doc_centric_random(tokenids : torch.Tensor, embs : torch.Tensor, num_prune, rng) ->  Tuple[Sequence[int], torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Returns the tokenids and embs after remvoing num_prune embeddings\n",
    "    \"\"\"\n",
    "    \n",
    "    length = len(tokenids)\n",
    "    \n",
    "    newmask = torch.BoolTensor(length)\n",
    "    newmask[:] = True\n",
    "    embs = embs[0:length, :]   \n",
    "    \n",
    "    for i in rng.sample(range(length), min(length, num_prune)):\n",
    "        newmask[i] = False\n",
    "        \n",
    "    return tokenids[newmask], embs[newmask,:]\n",
    "\n",
    "\n",
    "def prune_doccentric_transformerH0(factory, topK, active_ids=None, min_obs=0, seed=42):\n",
    "    import random\n",
    "    \n",
    "    rng = random.Random(seed)\n",
    "    assert active_ids is None \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_random(row.doc_toks, row.doc_embs, topK, rng),\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_random_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": "code",
   "execution_count": 22,
   "id": "finished-dominant",
   "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>1537</th>\n",
       "      <td>1</td>\n",
       "      <td>5453527</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1999</th>\n",
       "      <td>1</td>\n",
       "      <td>7605154</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>665</th>\n",
       "      <td>1</td>\n",
       "      <td>1833253</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>664</th>\n",
       "      <td>1</td>\n",
       "      <td>1833252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>586</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>1533</th>\n",
       "      <td>1</td>\n",
       "      <td>5436251</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2164</th>\n",
       "      <td>1</td>\n",
       "      <td>8311152</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2053</th>\n",
       "      <td>1</td>\n",
       "      <td>7875435</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1023</th>\n",
       "      <td>1</td>\n",
       "      <td>3444362</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1067</th>\n",
       "      <td>1</td>\n",
       "      <td>3605270</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     qid    docno\n",
       "1537   1  5453527\n",
       "1999   1  7605154\n",
       "665    1  1833253\n",
       "664    1  1833252\n",
       "586    1  1525431\n",
       "...   ..      ...\n",
       "1533   1  5436251\n",
       "2164   1  8311152\n",
       "2053   1  7875435\n",
       "1023   1  3444362\n",
       "1067   1  3605270\n",
       "\n",
       "[1000 rows x 2 columns]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prune_doccentric_random_pipeH0(10).search(\"chemical reactions\")[['qid', 'docno']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "automatic-impression",
   "metadata": {},
   "outputs": [],
   "source": [
    "DOC_CENTRIC_THRESHOLDS=[1,3,5,8,10,15,20,25,30]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bright-wilson",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9447514ae6f44e6fa7be6148875d4ffe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/27 [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>doc_centric_1</td>\n",
       "      <td>0.406474</td>\n",
       "      <td>0.694532</td>\n",
       "      <td>0.650674</td>\n",
       "      <td>0.584214</td>\n",
       "      <td>0.853488</td>\n",
       "      <td>0.853659</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>67.255814</td>\n",
       "      <td>75.909302</td>\n",
       "      <td>78.472326</td>\n",
       "      <td>75.997488</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>doc_centric_3</td>\n",
       "      <td>0.403502</td>\n",
       "      <td>0.695778</td>\n",
       "      <td>0.654692</td>\n",
       "      <td>0.584474</td>\n",
       "      <td>0.865116</td>\n",
       "      <td>0.865290</td>\n",
       "      <td>0.951550</td>\n",
       "      <td>0.951550</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>64.976744</td>\n",
       "      <td>75.209302</td>\n",
       "      <td>76.796047</td>\n",
       "      <td>73.997488</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>doc_centric_5</td>\n",
       "      <td>0.395113</td>\n",
       "      <td>0.678131</td>\n",
       "      <td>0.646355</td>\n",
       "      <td>0.577801</td>\n",
       "      <td>0.841860</td>\n",
       "      <td>0.842045</td>\n",
       "      <td>0.943798</td>\n",
       "      <td>0.943798</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>68.372093</td>\n",
       "      <td>73.232558</td>\n",
       "      <td>75.153256</td>\n",
       "      <td>71.997488</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>doc_centric_8</td>\n",
       "      <td>0.390897</td>\n",
       "      <td>0.676074</td>\n",
       "      <td>0.644773</td>\n",
       "      <td>0.578178</td>\n",
       "      <td>0.818605</td>\n",
       "      <td>0.818789</td>\n",
       "      <td>0.934109</td>\n",
       "      <td>0.934109</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>62.162791</td>\n",
       "      <td>70.541860</td>\n",
       "      <td>72.909767</td>\n",
       "      <td>68.997488</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>doc_centric_10</td>\n",
       "      <td>0.386604</td>\n",
       "      <td>0.670099</td>\n",
       "      <td>0.634522</td>\n",
       "      <td>0.571584</td>\n",
       "      <td>0.860465</td>\n",
       "      <td>0.860656</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.939922</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>65.441860</td>\n",
       "      <td>69.848837</td>\n",
       "      <td>71.484186</td>\n",
       "      <td>66.997488</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>doc_centric_15</td>\n",
       "      <td>0.376311</td>\n",
       "      <td>0.668657</td>\n",
       "      <td>0.632718</td>\n",
       "      <td>0.564782</td>\n",
       "      <td>0.818798</td>\n",
       "      <td>0.818999</td>\n",
       "      <td>0.959302</td>\n",
       "      <td>0.959302</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>58.976744</td>\n",
       "      <td>66.506977</td>\n",
       "      <td>67.436279</td>\n",
       "      <td>61.997581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>doc_centric_20</td>\n",
       "      <td>0.369363</td>\n",
       "      <td>0.651749</td>\n",
       "      <td>0.611360</td>\n",
       "      <td>0.547354</td>\n",
       "      <td>0.846899</td>\n",
       "      <td>0.847103</td>\n",
       "      <td>0.937984</td>\n",
       "      <td>0.937984</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>57.813953</td>\n",
       "      <td>63.383721</td>\n",
       "      <td>64.293256</td>\n",
       "      <td>56.999023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>doc_centric_25</td>\n",
       "      <td>0.347077</td>\n",
       "      <td>0.635784</td>\n",
       "      <td>0.601294</td>\n",
       "      <td>0.538971</td>\n",
       "      <td>0.854651</td>\n",
       "      <td>0.854831</td>\n",
       "      <td>0.934109</td>\n",
       "      <td>0.934109</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>57.930233</td>\n",
       "      <td>61.118605</td>\n",
       "      <td>61.310930</td>\n",
       "      <td>52.009907</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>doc_centric_30</td>\n",
       "      <td>0.340047</td>\n",
       "      <td>0.623785</td>\n",
       "      <td>0.586891</td>\n",
       "      <td>0.521498</td>\n",
       "      <td>0.825581</td>\n",
       "      <td>0.825834</td>\n",
       "      <td>0.914729</td>\n",
       "      <td>0.914729</td>\n",
       "      <td>0.695251</td>\n",
       "      <td>53.906977</td>\n",
       "      <td>59.611628</td>\n",
       "      <td>58.687674</td>\n",
       "      <td>47.047233</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.406474  0.694532  0.650674  0.584214      0.853488   \n",
       "1   doc_centric_3        0.403502  0.695778  0.654692  0.584474      0.865116   \n",
       "2   doc_centric_5        0.395113  0.678131  0.646355  0.577801      0.841860   \n",
       "3   doc_centric_8        0.390897  0.676074  0.644773  0.578178      0.818605   \n",
       "4  doc_centric_10        0.386604  0.670099  0.634522  0.571584      0.860465   \n",
       "5  doc_centric_15        0.376311  0.668657  0.632718  0.564782      0.818798   \n",
       "6  doc_centric_20        0.369363  0.651749  0.611360  0.547354      0.846899   \n",
       "7  doc_centric_25        0.347077  0.635784  0.601294  0.538971      0.854651   \n",
       "8  doc_centric_30        0.340047  0.623785  0.586891  0.521498      0.825581   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.853659  0.939922  0.939922       0.695251    67.255814     75.909302   \n",
       "1   0.865290  0.951550  0.951550       0.695251    64.976744     75.209302   \n",
       "2   0.842045  0.943798  0.943798       0.695251    68.372093     73.232558   \n",
       "3   0.818789  0.934109  0.934109       0.695251    62.162791     70.541860   \n",
       "4   0.860656  0.939922  0.939922       0.695251    65.441860     69.848837   \n",
       "5   0.818999  0.959302  0.959302       0.695251    58.976744     66.506977   \n",
       "6   0.847103  0.937984  0.937984       0.695251    57.813953     63.383721   \n",
       "7   0.854831  0.934109  0.934109       0.695251    57.930233     61.118605   \n",
       "8   0.825834  0.914729  0.914729       0.695251    53.906977     59.611628   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      78.472326  75.997488  \n",
       "1      76.796047  73.997488  \n",
       "2      75.153256  71.997488  \n",
       "3      72.909767  68.997488  \n",
       "4      71.484186  66.997488  \n",
       "5      67.436279  61.997581  \n",
       "6      64.293256  56.999023  \n",
       "7      61.310930  52.009907  \n",
       "8      58.687674  47.047233  "
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doccentric_2019 = do_exp(\n",
    "    [prune_doccentric_random_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/random_doccentric_%s/\"\n",
    ")\n",
    "doccentric_2019"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "municipal-corner",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f55d9f30d314f74a19d1773ec41acad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/27 [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>doc_centric_1</td>\n",
       "      <td>0.456601</td>\n",
       "      <td>0.685930</td>\n",
       "      <td>0.655573</td>\n",
       "      <td>0.612020</td>\n",
       "      <td>0.831481</td>\n",
       "      <td>0.832997</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.924399</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>69.037037</td>\n",
       "      <td>73.529630</td>\n",
       "      <td>77.506481</td>\n",
       "      <td>75.977167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>doc_centric_3</td>\n",
       "      <td>0.457122</td>\n",
       "      <td>0.688187</td>\n",
       "      <td>0.657075</td>\n",
       "      <td>0.611438</td>\n",
       "      <td>0.843827</td>\n",
       "      <td>0.846212</td>\n",
       "      <td>0.922840</td>\n",
       "      <td>0.925224</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>67.648148</td>\n",
       "      <td>72.335185</td>\n",
       "      <td>75.827778</td>\n",
       "      <td>73.977167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>doc_centric_5</td>\n",
       "      <td>0.447944</td>\n",
       "      <td>0.685629</td>\n",
       "      <td>0.649351</td>\n",
       "      <td>0.606512</td>\n",
       "      <td>0.814991</td>\n",
       "      <td>0.816314</td>\n",
       "      <td>0.916049</td>\n",
       "      <td>0.917372</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>65.574074</td>\n",
       "      <td>71.533333</td>\n",
       "      <td>74.315926</td>\n",
       "      <td>71.977167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>doc_centric_8</td>\n",
       "      <td>0.456141</td>\n",
       "      <td>0.693776</td>\n",
       "      <td>0.653468</td>\n",
       "      <td>0.605945</td>\n",
       "      <td>0.866049</td>\n",
       "      <td>0.866524</td>\n",
       "      <td>0.933951</td>\n",
       "      <td>0.934425</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>64.129630</td>\n",
       "      <td>69.527778</td>\n",
       "      <td>72.097593</td>\n",
       "      <td>68.977167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>doc_centric_10</td>\n",
       "      <td>0.435066</td>\n",
       "      <td>0.677901</td>\n",
       "      <td>0.638176</td>\n",
       "      <td>0.594374</td>\n",
       "      <td>0.813007</td>\n",
       "      <td>0.814691</td>\n",
       "      <td>0.920370</td>\n",
       "      <td>0.922054</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>63.074074</td>\n",
       "      <td>68.255556</td>\n",
       "      <td>70.834444</td>\n",
       "      <td>66.977167</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>doc_centric_15</td>\n",
       "      <td>0.431837</td>\n",
       "      <td>0.671697</td>\n",
       "      <td>0.635915</td>\n",
       "      <td>0.592673</td>\n",
       "      <td>0.823045</td>\n",
       "      <td>0.824729</td>\n",
       "      <td>0.929527</td>\n",
       "      <td>0.931210</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>57.500000</td>\n",
       "      <td>66.344444</td>\n",
       "      <td>67.368519</td>\n",
       "      <td>61.977352</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>doc_centric_20</td>\n",
       "      <td>0.433849</td>\n",
       "      <td>0.671999</td>\n",
       "      <td>0.634266</td>\n",
       "      <td>0.586405</td>\n",
       "      <td>0.831349</td>\n",
       "      <td>0.831349</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.925485</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>55.000000</td>\n",
       "      <td>62.372222</td>\n",
       "      <td>63.981667</td>\n",
       "      <td>56.981333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>doc_centric_25</td>\n",
       "      <td>0.412176</td>\n",
       "      <td>0.653189</td>\n",
       "      <td>0.618099</td>\n",
       "      <td>0.570496</td>\n",
       "      <td>0.854034</td>\n",
       "      <td>0.854034</td>\n",
       "      <td>0.948942</td>\n",
       "      <td>0.948942</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>53.981481</td>\n",
       "      <td>60.385185</td>\n",
       "      <td>61.201481</td>\n",
       "      <td>52.002389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>doc_centric_30</td>\n",
       "      <td>0.384685</td>\n",
       "      <td>0.638039</td>\n",
       "      <td>0.599165</td>\n",
       "      <td>0.551604</td>\n",
       "      <td>0.839506</td>\n",
       "      <td>0.839506</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.950617</td>\n",
       "      <td>0.754593</td>\n",
       "      <td>50.500000</td>\n",
       "      <td>58.340741</td>\n",
       "      <td>58.346481</td>\n",
       "      <td>47.050130</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.456601  0.685930  0.655573  0.612020      0.831481   \n",
       "1   doc_centric_3        0.457122  0.688187  0.657075  0.611438      0.843827   \n",
       "2   doc_centric_5        0.447944  0.685629  0.649351  0.606512      0.814991   \n",
       "3   doc_centric_8        0.456141  0.693776  0.653468  0.605945      0.866049   \n",
       "4  doc_centric_10        0.435066  0.677901  0.638176  0.594374      0.813007   \n",
       "5  doc_centric_15        0.431837  0.671697  0.635915  0.592673      0.823045   \n",
       "6  doc_centric_20        0.433849  0.671999  0.634266  0.586405      0.831349   \n",
       "7  doc_centric_25        0.412176  0.653189  0.618099  0.570496      0.854034   \n",
       "8  doc_centric_30        0.384685  0.638039  0.599165  0.551604      0.839506   \n",
       "\n",
       "   RR(rel=2)     RR@10        RR  R(rel=2)@1000  AvgDoclen@1  AvgDoclen@10  \\\n",
       "0   0.832997  0.922840  0.924399       0.754593    69.037037     73.529630   \n",
       "1   0.846212  0.922840  0.925224       0.754593    67.648148     72.335185   \n",
       "2   0.816314  0.916049  0.917372       0.754593    65.574074     71.533333   \n",
       "3   0.866524  0.933951  0.934425       0.754593    64.129630     69.527778   \n",
       "4   0.814691  0.920370  0.922054       0.754593    63.074074     68.255556   \n",
       "5   0.824729  0.929527  0.931210       0.754593    57.500000     66.344444   \n",
       "6   0.831349  0.925485  0.925485       0.754593    55.000000     62.372222   \n",
       "7   0.854034  0.948942  0.948942       0.754593    53.981481     60.385185   \n",
       "8   0.839506  0.950617  0.950617       0.754593    50.500000     58.340741   \n",
       "\n",
       "   AvgDoclen@100  AvgDoclen  \n",
       "0      77.506481  75.977167  \n",
       "1      75.827778  73.977167  \n",
       "2      74.315926  71.977167  \n",
       "3      72.097593  68.977167  \n",
       "4      70.834444  66.977167  \n",
       "5      67.368519  61.977352  \n",
       "6      63.981667  56.981333  \n",
       "7      61.201481  52.002389  \n",
       "8      58.346481  47.050130  "
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doccentric_2020 = do_exp(\n",
    "    [prune_doccentric_random_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/random_doccentric_%s/\"\n",
    ")\n",
    "doccentric_2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "partial-helen",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78e6aadda396412388326f7a11e32c8d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pt.Experiment:   0%|          | 0/3141 [00:00<?, ?batches/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "doccentric_dev = do_exp(\n",
    "    [prune_doccentric_random_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/random_doccentric_%s/\"\n",
    ")\n",
    "doccentric_dev"
   ]
  }
 ],
 "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
}