msmarco-psg-v1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "07cbb4bc",
   "metadata": {},
   "source": [
    "# SPLADE on MSMARCO v1 Passage Corpus using PyTerrier.\n",
    "\n",
    "This notebook demonstrates the creation of a SPLADE index using PyTerrier."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f939e2d1",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "Installation is using PIP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "heavy-methodology",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q git+https://github.com/naver/splade.git\n",
    "%pip install -q git+https://github.com/cmacdonald/pyt_splade.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c6943a4",
   "metadata": {},
   "source": [
    "## SPLADE setup\n",
    "\n",
    "We create a factory object that gives us access to the appropriate transformers to use SPLADE. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "municipal-arizona",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pyterrier as pt\n",
    "import pyt_splade\n",
    "splade = pyt_splade.Splade()\n",
    "doc_encoder = splade.doc_encoder()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65c47343",
   "metadata": {},
   "source": [
    "## Indexing demonstration\n",
    "\n",
    "Lets see what terms are generated by the SPLADE model during indexing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "executed-latest",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'w': 199,\n",
       " '##2': 193,\n",
       " 'war': 167,\n",
       " 'wwii': 150,\n",
       " '##w': 130,\n",
       " 'ii': 110,\n",
       " '2': 94,\n",
       " 'germany': 86,\n",
       " 'army': 76,\n",
       " 'battle': 70,\n",
       " 'was': 66,\n",
       " 'bomb': 48,\n",
       " 'event': 43,\n",
       " 'wilson': 43,\n",
       " 'conflict': 38,\n",
       " 'marshall': 33,\n",
       " 'allied': 23,\n",
       " 'surrender': 22,\n",
       " 'peace': 16,\n",
       " 'military': 12,\n",
       " 'era': 10,\n",
       " 'alliance': 10,\n",
       " 'weapon': 10,\n",
       " 'wars': 8,\n",
       " 'camp': 7,\n",
       " 'were': 6,\n",
       " 'france': 6,\n",
       " 'invasion': 6,\n",
       " 'nazi': 4,\n",
       " 'zombie': 2,\n",
       " 'patton': 2,\n",
       " 'german': 1}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = doc_encoder([{'docno' : 'd1', 'text' : 'ww2'}])\n",
    "df.iloc[0].toks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f37f3237",
   "metadata": {},
   "source": [
    "## Indexing MSMARCO\n",
    "\n",
    "Lets go and create an index for the MSMARCO v1 passage corpus. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "functioning-southwest",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this will provide access to the corpus\n",
    "dataset = pt.get_dataset('irds:msmarco-passage')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb1bec5",
   "metadata": {},
   "source": [
    "This is the actual indexing code. We use the SPLADE model to transform the passages into tokens and weights.\n",
    "\n",
    "Indexing the 8.8M passages of MSMARCO using a GeForce RTX 3090 took 5 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "auburn-photography",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('./msmarco_psg'): # skip if alraedy created\n",
    "    indexer = pt.IterDictIndexer('./msmarco_psg', pretokenised=True)\n",
    "    indexer.setProperty(\"termpipelines\", \"\")\n",
    "    indexer.setProperty(\"tokeniser\", \"WhitespaceTokeniser\")\n",
    "    \n",
    "    indexer_pipe = doc_encoder >> indexer\n",
    "    index_ref = indexer_pipe.index(dataset.get_corpus_iter())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1877b30b",
   "metadata": {},
   "source": [
    "## Retrieval\n",
    "\n",
    "We can now conduct retrieval using PyTerrier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "deluxe-nature",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Java started (triggered by Retriever.__init__) and loaded: pyterrier.java, pyterrier.terrier.java [version=5.10 (build: craigm 2024-08-22 17:33), helper_version=0.0.8]\n"
     ]
    }
   ],
   "source": [
    "retr = pt.terrier.Retriever('./msmarco_psg', wmodel='Tf', verbose=True)\n",
    "\n",
    "retr_pipe = splade.query_encoder() >> retr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f48d770",
   "metadata": {},
   "source": [
    "Let's check retrieval works, and we can see the generated query. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "novel-spine",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TerrierRetr(Tf): 100%|██████████| 1/1 [00:00<00:00,  1.11q/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qid</th>\n",
       "      <th>docid</th>\n",
       "      <th>docno</th>\n",
       "      <th>rank</th>\n",
       "      <th>score</th>\n",
       "      <th>query</th>\n",
       "      <th>query_toks</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>758284</td>\n",
       "      <td>758284</td>\n",
       "      <td>0</td>\n",
       "      <td>759.955493</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>5913794</td>\n",
       "      <td>5913794</td>\n",
       "      <td>1</td>\n",
       "      <td>758.944513</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>742206</td>\n",
       "      <td>742206</td>\n",
       "      <td>2</td>\n",
       "      <td>757.742353</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>8572191</td>\n",
       "      <td>8572191</td>\n",
       "      <td>3</td>\n",
       "      <td>751.762260</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>129901</td>\n",
       "      <td>129901</td>\n",
       "      <td>4</td>\n",
       "      <td>749.707581</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>995</th>\n",
       "      <td>1</td>\n",
       "      <td>6342771</td>\n",
       "      <td>6342771</td>\n",
       "      <td>995</td>\n",
       "      <td>543.450242</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>996</th>\n",
       "      <td>1</td>\n",
       "      <td>3355651</td>\n",
       "      <td>3355651</td>\n",
       "      <td>996</td>\n",
       "      <td>543.413003</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>997</th>\n",
       "      <td>1</td>\n",
       "      <td>5094605</td>\n",
       "      <td>5094605</td>\n",
       "      <td>997</td>\n",
       "      <td>543.349945</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>998</th>\n",
       "      <td>1</td>\n",
       "      <td>2226467</td>\n",
       "      <td>2226467</td>\n",
       "      <td>998</td>\n",
       "      <td>543.343830</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999</th>\n",
       "      <td>1</td>\n",
       "      <td>7866969</td>\n",
       "      <td>7866969</td>\n",
       "      <td>999</td>\n",
       "      <td>543.317541</td>\n",
       "      <td>chemical reactions</td>\n",
       "      <td>{'reactions': 269.44158935546875, 'reaction': ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000 rows × 7 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    qid    docid    docno  rank       score               query  \\\n",
       "0     1   758284   758284     0  759.955493  chemical reactions   \n",
       "1     1  5913794  5913794     1  758.944513  chemical reactions   \n",
       "2     1   742206   742206     2  757.742353  chemical reactions   \n",
       "3     1  8572191  8572191     3  751.762260  chemical reactions   \n",
       "4     1   129901   129901     4  749.707581  chemical reactions   \n",
       "..   ..      ...      ...   ...         ...                 ...   \n",
       "995   1  6342771  6342771   995  543.450242  chemical reactions   \n",
       "996   1  3355651  3355651   996  543.413003  chemical reactions   \n",
       "997   1  5094605  5094605   997  543.349945  chemical reactions   \n",
       "998   1  2226467  2226467   998  543.343830  chemical reactions   \n",
       "999   1  7866969  7866969   999  543.317541  chemical reactions   \n",
       "\n",
       "                                            query_toks  \n",
       "0    {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "1    {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "2    {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "3    {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "4    {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "..                                                 ...  \n",
       "995  {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "996  {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "997  {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "998  {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "999  {'reactions': 269.44158935546875, 'reaction': ...  \n",
       "\n",
       "[1000 rows x 7 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retr_pipe.search('chemical reactions')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65710e19",
   "metadata": {},
   "source": [
    "Finally, lets run the experiment and see the resulting performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "simplified-planner",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "TerrierRetr(Tf): 100%|██████████| 43/43 [01:07<00:00,  1.57s/q]\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>RR(rel=2)</th>\n",
       "      <th>nDCG@10</th>\n",
       "      <th>nDCG@100</th>\n",
       "      <th>AP(rel=2)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>splade</td>\n",
       "      <td>0.918605</td>\n",
       "      <td>0.729712</td>\n",
       "      <td>0.671791</td>\n",
       "      <td>0.50429</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     name  RR(rel=2)   nDCG@10  nDCG@100  AP(rel=2)\n",
       "0  splade   0.918605  0.729712  0.671791    0.50429"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyterrier.measures import *\n",
    "pt.Experiment(\n",
    "    [retr_pipe],\n",
    "    pt.get_dataset('irds:msmarco-passage/trec-dl-2019/judged').get_topics(),\n",
    "    pt.get_dataset('irds:msmarco-passage/trec-dl-2019/judged').get_qrels(),\n",
    "    eval_metrics=[RR(rel=2), nDCG@10, nDCG@100, AP(rel=2)],\n",
    "    names=['splade']\n",
    ")        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c27407b9",
   "metadata": {},
   "source": [
    "## Exploring the Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "reflected-little",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = pt.java.cast(\"org.terrier.querying.LocalManager\", retr.manager).index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9243f907",
   "metadata": {},
   "source": [
    "Lets explore the lexicon - what tokens were used? (First 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "systematic-disclosure",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "! term8767 Nt=35620 TF=1146138 maxTF=2147483647 @{0 0 0}\n",
      "\" term909 Nt=860880 TF=35393854 maxTF=2147483647 @{0 193938 6}\n",
      "# term5228 Nt=69447 TF=3962047 maxTF=2147483647 @{0 5070117 6}\n",
      "##0 term5242 Nt=68165 TF=4325474 maxTF=2147483647 @{0 5657689 4}\n",
      "##00 term19503 Nt=14652 TF=971888 maxTF=2147483647 @{0 6283938 7}\n",
      "##01 term12384 Nt=7847 TF=535341 maxTF=2147483647 @{0 6430337 3}\n",
      "##0s term26590 Nt=389 TF=27046 maxTF=2147483647 @{0 6512277 3}\n",
      "##1 term5497 Nt=105121 TF=5875946 maxTF=2147483647 @{0 6516783 2}\n",
      "##10 term17386 Nt=21140 TF=1573665 maxTF=2147483647 @{0 7368646 7}\n",
      "##100 term12385 Nt=13618 TF=806019 maxTF=2147483647 @{0 7599660 2}\n",
      "##11 term9508 Nt=17101 TF=1082393 maxTF=2147483647 @{0 7723654 3}\n",
      "##12 term8683 Nt=8393 TF=695082 maxTF=2147483647 @{0 7887210 5}\n",
      "##13 term12421 Nt=13616 TF=714315 maxTF=2147483647 @{0 7989915 2}\n",
      "##14 term17858 Nt=6103 TF=427235 maxTF=2147483647 @{0 8102904 7}\n",
      "##15 term13479 Nt=14605 TF=926123 maxTF=2147483647 @{0 8168491 1}\n",
      "##16 term8682 Nt=6139 TF=441595 maxTF=2147483647 @{0 8309519 5}\n",
      "##17 term24327 Nt=2636 TF=230487 maxTF=2147483647 @{0 8376846 5}\n",
      "##18 term16765 Nt=3899 TF=290747 maxTF=2147483647 @{0 8411696 4}\n",
      "##19 term10533 Nt=4092 TF=245854 maxTF=2147483647 @{0 8456419 6}\n",
      "##2 term8648 Nt=105405 TF=9244655 maxTF=2147483647 @{0 8496024 4}\n",
      "##20 term18612 Nt=7489 TF=659393 maxTF=2147483647 @{0 9766749 4}\n",
      "##200 term25236 Nt=2388 TF=231039 maxTF=2147483647 @{0 9863810 6}\n",
      "##21 term23253 Nt=4138 TF=373183 maxTF=2147483647 @{0 9898291 7}\n",
      "##22 term18179 Nt=8746 TF=661821 maxTF=2147483647 @{0 9953516 6}\n",
      "##23 term17559 Nt=3047 TF=281560 maxTF=2147483647 @{0 10052645 7}\n",
      "##24 term17724 Nt=3376 TF=278105 maxTF=2147483647 @{0 10094475 6}\n",
      "##25 term21612 Nt=5195 TF=400257 maxTF=2147483647 @{0 10136575 3}\n",
      "##26 term5510 Nt=9452 TF=451537 maxTF=2147483647 @{0 10196996 1}\n",
      "##27 term21849 Nt=4385 TF=308564 maxTF=2147483647 @{0 10271178 6}\n",
      "##28 term17723 Nt=3057 TF=203644 maxTF=2147483647 @{0 10318874 1}\n",
      "##29 term18150 Nt=3254 TF=218930 maxTF=2147483647 @{0 10351066 0}\n",
      "##3 term9090 Nt=100443 TF=6162406 maxTF=2147483647 @{0 10385492 4}\n",
      "##30 term8700 Nt=10287 TF=676645 maxTF=2147483647 @{0 11262879 7}\n",
      "##31 term8107 Nt=11102 TF=522073 maxTF=2147483647 @{0 11366143 5}\n",
      "##32 term22738 Nt=6039 TF=704956 maxTF=2147483647 @{0 11451578 4}\n",
      "##33 term12019 Nt=3189 TF=240815 maxTF=2147483647 @{0 11550979 3}\n",
      "##34 term5509 Nt=4681 TF=349799 maxTF=2147483647 @{0 11588002 7}\n",
      "##35 term15229 Nt=3597 TF=374157 maxTF=2147483647 @{0 11641339 7}\n",
      "##36 term8752 Nt=2978 TF=301235 maxTF=2147483647 @{0 11695555 1}\n",
      "##37 term8748 Nt=5195 TF=305722 maxTF=2147483647 @{0 11739746 4}\n",
      "##38 term8736 Nt=10589 TF=566708 maxTF=2147483647 @{0 11788302 7}\n",
      "##39 term21617 Nt=2366 TF=191109 maxTF=2147483647 @{0 11878110 2}\n",
      "##3d term23449 Nt=824 TF=71289 maxTF=2147483647 @{0 11907285 3}\n",
      "##4 term5220 Nt=73412 TF=5356102 maxTF=2147483647 @{0 11918282 2}\n",
      "##40 term8696 Nt=6188 TF=513084 maxTF=2147483647 @{0 12676448 4}\n",
      "##400 term23344 Nt=869 TF=103966 maxTF=2147483647 @{0 12752579 2}\n",
      "##41 term8695 Nt=7573 TF=366839 maxTF=2147483647 @{0 12767777 3}\n",
      "##42 term10284 Nt=2446 TF=223238 maxTF=2147483647 @{0 12828293 7}\n",
      "##43 term17560 Nt=3712 TF=264580 maxTF=2147483647 @{0 12861751 3}\n",
      "##44 term5554 Nt=18204 TF=710017 maxTF=2147483647 @{0 12902480 5}\n",
      "##45 term14730 Nt=5030 TF=467079 maxTF=2147483647 @{0 13020075 6}\n",
      "##46 term5493 Nt=3744 TF=273381 maxTF=2147483647 @{0 13088753 5}\n",
      "##47 term5492 Nt=5004 TF=331699 maxTF=2147483647 @{0 13130898 2}\n",
      "##48 term20854 Nt=1245 TF=110909 maxTF=2147483647 @{0 13182433 7}\n",
      "##49 term5491 Nt=3089 TF=209450 maxTF=2147483647 @{0 13199345 5}\n",
      "##4th term26677 Nt=98 TF=7839 maxTF=2147483647 @{0 13232227 0}\n",
      "##5 term9688 Nt=59317 TF=3827195 maxTF=2147483647 @{0 13233555 1}\n",
      "##50 term21092 Nt=4266 TF=387278 maxTF=2147483647 @{0 13789398 1}\n",
      "##500 term23143 Nt=1114 TF=149132 maxTF=2147483647 @{0 13846451 1}\n",
      "##51 term23512 Nt=2675 TF=236230 maxTF=2147483647 @{0 13867765 3}\n",
      "##52 term24259 Nt=3477 TF=262632 maxTF=2147483647 @{0 13903210 0}\n",
      "##53 term14969 Nt=2714 TF=246428 maxTF=2147483647 @{0 13943259 3}\n",
      "##54 term18395 Nt=2610 TF=223220 maxTF=2147483647 @{0 13979895 3}\n",
      "##55 term19087 Nt=3263 TF=244554 maxTF=2147483647 @{0 14013648 7}\n",
      "##56 term19086 Nt=972 TF=82586 maxTF=2147483647 @{0 14051191 4}\n",
      "##57 term19085 Nt=1950 TF=174913 maxTF=2147483647 @{0 14063950 2}\n",
      "##58 term23813 Nt=1396 TF=111250 maxTF=2147483647 @{0 14090363 5}\n",
      "##59 term12642 Nt=4046 TF=250890 maxTF=2147483647 @{0 14107650 7}\n",
      "##6 term8738 Nt=61284 TF=4087261 maxTF=2147483647 @{0 14147606 1}\n",
      "##60 term21093 Nt=6589 TF=498311 maxTF=2147483647 @{0 14738015 2}\n",
      "##64 term14967 Nt=3579 TF=327918 maxTF=2147483647 @{0 14812830 0}\n",
      "##65 term8770 Nt=1823 TF=186589 maxTF=2147483647 @{0 14861289 5}\n",
      "##66 term20390 Nt=1647 TF=165527 maxTF=2147483647 @{0 14888865 7}\n",
      "##6th term21814 Nt=166 TF=11344 maxTF=2147483647 @{0 14913286 7}\n",
      "##7 term5892 Nt=47423 TF=3061766 maxTF=2147483647 @{0 14915233 3}\n",
      "##70 term375 Nt=5422 TF=419074 maxTF=2147483647 @{0 15361988 0}\n",
      "##75 term12404 Nt=2318 TF=219353 maxTF=2147483647 @{0 15425064 0}\n",
      "##7th term22520 Nt=133 TF=12471 maxTF=2147483647 @{0 15457839 7}\n",
      "##8 term5241 Nt=62666 TF=3415246 maxTF=2147483647 @{0 15459779 3}\n",
      "##80 term19504 Nt=4460 TF=354197 maxTF=2147483647 @{0 15966939 5}\n",
      "##85 term24749 Nt=1705 TF=183597 maxTF=2147483647 @{0 16020244 2}\n",
      "##86 term20582 Nt=4161 TF=387450 maxTF=2147483647 @{0 16047107 2}\n",
      "##8th term22519 Nt=209 TF=9788 maxTF=2147483647 @{0 16104018 7}\n",
      "##9 term8769 Nt=37731 TF=2437626 maxTF=2147483647 @{0 16105893 4}\n",
      "##90 term12383 Nt=7155 TF=505851 maxTF=2147483647 @{0 16464046 7}\n",
      "##9th term27301 Nt=129 TF=10535 maxTF=2147483647 @{0 16541014 5}\n",
      "##a term1821 Nt=346105 TF=22046692 maxTF=2147483647 @{0 16542754 7}\n",
      "##a1 term25025 Nt=765 TF=67593 maxTF=2147483647 @{0 19555482 4}\n",
      "##aa term13734 Nt=10360 TF=1261511 maxTF=2147483647 @{0 19565793 4}\n",
      "##aan term19779 Nt=3116 TF=290401 maxTF=2147483647 @{0 19740723 5}\n",
      "##aar term23540 Nt=435 TF=78231 maxTF=2147483647 @{0 19782740 6}\n",
      "##ab term11711 Nt=31941 TF=2737720 maxTF=2147483647 @{0 19793592 2}\n",
      "##aba term19539 Nt=2463 TF=400307 maxTF=2147483647 @{0 20179524 5}\n",
      "##abad term26450 Nt=349 TF=47897 maxTF=2147483647 @{0 20234153 7}\n",
      "##abas term27079 Nt=510 TF=117880 maxTF=2147483647 @{0 20240969 7}\n",
      "##abe term21927 Nt=3437 TF=480491 maxTF=2147483647 @{0 20256883 3}\n",
      "##abi term18987 Nt=9253 TF=883705 maxTF=2147483647 @{0 20323049 1}\n",
      "##abia term28034 Nt=193 TF=35780 maxTF=2147483647 @{0 20448202 3}\n",
      "##ability term12378 Nt=18884 TF=1335655 maxTF=2147483647 @{0 20453179 4}\n",
      "##able term3684 Nt=180340 TF=12313273 maxTF=2147483647 @{0 20652456 3}\n"
     ]
    }
   ],
   "source": [
    "for i, entry in enumerate(index.getLexicon()):\n",
    "    if i == 100:\n",
    "        break\n",
    "    print(entry.getKey() + \" \" + entry.getValue().toString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ordered-panama",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of documents: 8841823\n",
      "Number of terms: 28679\n",
      "Number of postings: 1037725680\n",
      "Number of fields: 0\n",
      "Number of tokens: 51855220573\n",
      "Field names: []\n",
      "Positions:   false\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(index.getCollectionStatistics().toString())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "475d4a65",
   "metadata": {},
   "source": [
    "We can even look into particular document in the index. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "promising-breath",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\" 13\n",
      "##uation 37\n",
      "000 59\n",
      "25 43\n",
      "30 70\n",
      "35 72\n",
      "40 35\n",
      "accountant 39\n",
      "accounting 70\n",
      "advice 11\n",
      "amount 71\n",
      "amounts 88\n",
      "applicants 12\n",
      "ask 108\n",
      "asked 86\n",
      "asking 33\n",
      "asks 6\n",
      "assessment 19\n",
      "average 38\n",
      "bargaining 48\n",
      "bart 20\n",
      "bottom 94\n",
      "briggs 20\n",
      "burke 10\n",
      "business 13\n",
      "businesses 28\n",
      "calculate 62\n",
      "calculated 2\n",
      "candidacy 30\n",
      "candidate 109\n",
      "candidates 115\n",
      "chart 26\n",
      "companies 43\n",
      "company 15\n",
      "considered 36\n",
      "corporate 10\n",
      "dave 17\n",
      "davis 5\n",
      "desk 7\n",
      "diversity 2\n",
      "employee 95\n",
      "employees 17\n",
      "employment 40\n",
      "engineer 20\n",
      "example 107\n",
      "examples 98\n",
      "excel 14\n",
      "executive 5\n",
      "finance 43\n",
      "fisher 8\n",
      "flat 134\n",
      "flex 50\n",
      "flexibility 129\n",
      "flexible 92\n",
      "gage 6\n",
      "give 98\n",
      "given 17\n",
      "giving 29\n",
      "highest 1\n",
      "hr 62\n",
      "improvisation 9\n",
      "include 57\n",
      "included 28\n",
      "income 15\n",
      "interview 19\n",
      "job 80\n",
      "jobs 33\n",
      "kelly 1\n",
      "letter 9\n",
      "low 76\n",
      "management 4\n",
      "marketing 14\n",
      "matching 23\n",
      "math 56\n",
      "max 44\n",
      "maximum 84\n",
      "median 101\n",
      "mid 102\n",
      "middle 32\n",
      "money 31\n",
      "murray 21\n",
      "negotiate 98\n",
      "negotiating 49\n",
      "negotiation 82\n",
      "negotiations 59\n",
      "normal 13\n",
      "numbers 9\n",
      "often 19\n",
      "pay 139\n",
      "payroll 7\n",
      "point 59\n",
      "points 50\n",
      "post 50\n",
      "posted 72\n",
      "posting 95\n",
      "practical 5\n",
      "price 123\n",
      "provide 54\n",
      "provided 77\n",
      "qualification 14\n",
      "range 233\n",
      "ranges 149\n",
      "rather 52\n",
      "recruit 36\n",
      "required 70\n",
      "requirements 125\n",
      "respect 4\n",
      "resume 37\n",
      "salaries 22\n",
      "salary 224\n",
      "say 93\n",
      "smith 43\n",
      "statistics 1\n",
      "sum 14\n",
      "survey 35\n",
      "thirty 24\n",
      "thousand 11\n",
      "wage 40\n",
      "when 61\n"
     ]
    }
   ],
   "source": [
    "di = index.getDirectIndex()\n",
    "doi = index.getDocumentIndex()\n",
    "lex = index.getLexicon()\n",
    "docid = 7700000 #docids are 0-based\n",
    "#NB: postings will be null if the document is empty\n",
    "dictrep = {}\n",
    "for posting in di.getPostings(doi.getDocumentEntry(docid)):\n",
    "    termid = posting.getId()\n",
    "    lee = lex.getLexiconEntry(termid)\n",
    "    dictrep[lee.getKey()] = posting.getFrequency()\n",
    "\n",
    "for k in sorted(dictrep.keys()):\n",
    "    print(k, dictrep[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "superior-oxygen",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}