Guide to using ToTER.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2fc6718a",
   "metadata": {},
   "source": [
    "We provide a step-by-step guide for applying the ToTER framework (last updated: 24.02.21)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fc0b830b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"3\"\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.simplefilter('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "049dff82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from Utils import *\n",
    "import numpy as np\n",
    "from beir.datasets.data_loader import GenericDataLoader\n",
    "from beir import util, LoggingHandler\n",
    "from beir.retrieval import models\n",
    "from beir.retrieval.evaluation import EvaluateRetrieval\n",
    "from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES\n",
    "\n",
    "import pathlib, os\n",
    "from tqdm.auto import tqdm, trange"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "133a31fd",
   "metadata": {},
   "source": [
    "## Data load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6a7bdb3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b81ebcea4c64f089f20b1d7ebe3e5aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25657 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = \"scidocs\"\n",
    "url = \"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip\".format(dataset)\n",
    "out_dir = os.path.join(\"datasets\")\n",
    "data_path = util.download_and_unzip(url, out_dir)\n",
    "\n",
    "corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split=\"test\")\n",
    "retriever = EvaluateRetrieval()\n",
    "\n",
    "k_values=[10,50,100,500, 1000, 2500]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fee15f76",
   "metadata": {},
   "source": [
    "## Backbone retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ab89120",
   "metadata": {},
   "source": [
    "Contriever-MS has been trained using massive training labels from general domain (MS-MARCO)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be0ae83f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.16524 , NDCG@100: 0.23594\n",
      "Recall@100: 0.37807 , Recall@500: 0.54943 , Recall@1000: 0.62155 , Recall@2500: 0.7239\n"
     ]
    }
   ],
   "source": [
    "from IPython.utils import io\n",
    "\n",
    "with torch.no_grad():\n",
    "    with io.capture_output() as captured:\n",
    "        text_encoder = DRES(models.SentenceBERT(\"facebook/contriever-msmarco\"), batch_size=32)\n",
    "        retriever = EvaluateRetrieval(text_encoder, score_function=\"dot\", k_values=k_values + [len(corpus)])\n",
    "        CTR_full_results = retriever.retrieve(corpus, queries)\n",
    "print_metric(retriever.evaluate(qrels, CTR_full_results, retriever.k_values))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a2f9ec",
   "metadata": {},
   "source": [
    "## Training Phase: Class relevance learning using Taxonomy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efce3a88",
   "metadata": {},
   "source": [
    "#### 1. Load pre-computed PLM representations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc7af7f",
   "metadata": {},
   "source": [
    "We encode each text and topic class using PLM. If we directly use the PLM representations for retrieval, the results are unsatisfactory, as they have not been fine-tuned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a6db692e",
   "metadata": {},
   "outputs": [],
   "source": [
    "c_id_list, q_id_list = np.load('resource/test_doc_id_list.npy'), np.load('resource/test_q_id_list.npy')\n",
    "c_emb_list, q_emb_list = np.load('resource/test_doc_emb.npy'), np.load('resource/test_q_emb.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c8c52812",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:17<00:00, 56.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.02657 , NDCG@100: 0.05212\n",
      "Recall@100: 0.10853 , Recall@500: 0.24578 , Recall@1000: 0.33232 , Recall@2500: 0.46243\n"
     ]
    }
   ],
   "source": [
    "score_mat = np.matmul(q_emb_list, c_emb_list.T)\n",
    "PLM_results = eval_full_score_mat(score_mat, q_id_list, c_id_list)\n",
    "print_metric(retriever.evaluate(qrels, PLM_results, k_values))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "062c3c30",
   "metadata": {},
   "source": [
    "#### 2. Conduct class relevance learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28bb0bb",
   "metadata": {},
   "source": [
    "This example uses a simple MLP classifier. More sophisticated choices can further improve the final result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "802ac118",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Class_relevance_learning import train_classifier\n",
    "train_classifier()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "102f2e8e",
   "metadata": {},
   "source": [
    "## Inference Phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97a4fc7f",
   "metadata": {},
   "source": [
    "We first compute topic distributions for queries and documents, and then apply filtering to retain topic classes that have a certain degree of relevance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5398cc0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Classifier(num_class=4028)\n",
    "model.load_state_dict(torch.load('resource/Classifier_model'))\n",
    "\n",
    "with torch.no_grad():\n",
    "    c_emb_tensor = torch.FloatTensor(c_emb_list)\n",
    "    c_clf_output = model(c_emb_tensor).numpy()\n",
    "    \n",
    "    q_emb_tensor = torch.FloatTensor(q_emb_list)\n",
    "    q_clf_output = model(q_emb_tensor).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3ee0aca7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "613ef28e0b674ef1b3a223750c3da222",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88871bb65403429ab9b100762b835686",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "filtered_X = filtering(q_clf_output)\n",
    "filtered_Y = filtering(c_clf_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44c0407",
   "metadata": {},
   "source": [
    "### 1. search space adjustment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "227c4392",
   "metadata": {},
   "source": [
    "Before applying retrievers, we filter out a large number of irrelevant documents that have little overlap in topic classes with the query. In this example, we retain 10% of the corpus size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "492ed269",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.contrib.concurrent import process_map\n",
    "\n",
    "def topic_filter(i):\n",
    "    num_query_topic = (filtered_X[i] > 0).sum()    \n",
    "    co_topic_rank = np.argsort(-(filtered_Y * (filtered_X[i] > 0) > 0).sum(1))\n",
    "    tmp = topical_relatedness_mat[i].copy()\n",
    "    is_filtered = co_topic_rank[int(len(co_topic_rank) * adjustment_percent)+1:]\n",
    "    tmp[is_filtered] = 0.\n",
    "    \n",
    "    return tmp  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c7f7fbf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "For each query, we retain 10.0 % of the documents in the corpus.\n"
     ]
    }
   ],
   "source": [
    "topical_relatedness_mat = compute_topical_relatedness(filtered_X, filtered_Y)\n",
    "\n",
    "with io.capture_output() as captured:\n",
    "    adjustment_percent = 0.1\n",
    "    r = process_map(topic_filter, range(q_clf_output.shape[0]), max_workers=40)\n",
    "    SSA_score_mat = np.asarray(r)\n",
    "\n",
    "non_zero_percent = (((SSA_score_mat > 0).sum(1) / (topical_relatedness_mat > 0).sum(1)).mean() * 100)\n",
    "print(\"For each query, we retain\", round(non_zero_percent, 2), \"% of the documents in the corpus.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3a63ac2a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a670c9a2e7bc47dc9894f4f9b99a8adb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall@2500: 0.82947\n"
     ]
    }
   ],
   "source": [
    "SSA_results = eval_full_score_mat(SSA_score_mat, q_id_list, c_id_list)\n",
    "SSA_results = return_topK_result(SSA_results, topk=2500)\n",
    "print('Recall@2500:', retriever.evaluate(qrels, SSA_results, [2500])[2]['Recall@2500'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "669bed5c",
   "metadata": {},
   "source": [
    "###  2. Class relevance matching (CRM) for retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf34a0ee",
   "metadata": {},
   "source": [
    "We can conduct retrieval on the reduced search space. This step can benefit subsequent retrieval by reducing the search\n",
    "space while preserving topically relevant documents that may otherwise be overlooked by PLM-based retrievers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0df0759c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f3467d67ab14f09bca1f5b50fafb0f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.16749 , NDCG@100: 0.24042\n",
      "Recall@100: 0.38867 , Recall@500: 0.58482 , Recall@1000: 0.6813 , Recall@2500: 0.82947\n"
     ]
    }
   ],
   "source": [
    "CTR_SSA_results = {}\n",
    "for _, qid in tqdm(enumerate(SSA_results)):\n",
    "    CTR_SSA_results[qid] = {}\n",
    "    for cid in SSA_results[qid]:\n",
    "        CTR_SSA_results[qid][cid] = CTR_full_results[qid][cid]\n",
    "        \n",
    "print_metric(retriever.evaluate(qrels, CTR_SSA_results, retriever.k_values))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881507c1",
   "metadata": {},
   "source": [
    "We retrieve documents based on both semantic similarity (retriever) and topical relatedness (our classifier). This can help to handle lexical mismatches and fill in missing contexts, providing a complementary aspect to semantic similarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "37cafc0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "normalized_CRM_results = z_normalize(SSA_results)\n",
    "normalized_CTR_results = z_normalize(CTR_full_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f8138f70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ebae8b21d1d4fafb02c8c88af019359",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.17636 , NDCG@100: 0.26056\n",
      "Recall@100: 0.4337 , Recall@500: 0.636 , Recall@1000: 0.71805 , Recall@2500: 0.82947\n"
     ]
    }
   ],
   "source": [
    "CRM_results = {}\n",
    "for _, qid in tqdm(enumerate(SSA_results)):\n",
    "    CRM_results[qid] = {}\n",
    "\n",
    "    for cid in SSA_results[qid]:\n",
    "        s_CRM = normalized_CRM_results[qid][cid]\n",
    "        s_de = normalized_CTR_results[qid][cid]\n",
    "\n",
    "        CRM_results[qid][cid] = s_de + s_CRM\n",
    "        \n",
    "print_metric(retriever.evaluate(qrels, CRM_results, retriever.k_values))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67c34910",
   "metadata": {},
   "source": [
    "### 3. Query enrichment by core phrases (QEP) for reranker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16bdffd7",
   "metadata": {},
   "source": [
    "In this last stage, a reranker reorder top-ranked candidates from the retriever. Here, we delve deeper into each topic by focusing on  class-related phrases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10133361",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_result = return_topK_result(CRM_results, topk=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "42c29f86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f9ac5f9e19a5419c965761eaa7c4fd76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.17992\n"
     ]
    }
   ],
   "source": [
    "from beir.reranking.models import MonoT5\n",
    "from beir.reranking import Rerank\n",
    "\n",
    "T5_model = MonoT5('castorini/monot5-base-msmarco-10k', token_false='▁false', token_true='▁true')\n",
    "reranker = Rerank(T5_model, batch_size=1024)\n",
    "\n",
    "T5_rerank_results = reranker.rerank(corpus, queries, top_result, top_k=100)\n",
    "print('NDCG@10:', retriever.evaluate(qrels, T5_rerank_results, [10])[0]['NDCG@10'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62a43fc3",
   "metadata": {},
   "source": [
    "Users familiar with a domain often omit contexts in their queries, which makes it difficult to find accurate relevance. We identify missing contexts (called core phrases) and enrich the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "af520bd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original query:\n",
      "Provable data possession at untrusted stores\n",
      "Topic-enriched query:\n",
      "Provable data possession at untrusted stores, relevant topics: data, key, server, access_control, security, system, encryption, allows, untrusted\n"
     ]
    }
   ],
   "source": [
    "with open('resource/topic_enriched_queries', 'rb') as f:\n",
    "    topic_enriched_queries = pickle.load(f)\n",
    "\n",
    "example_qid = \"2d52f69dd4686a3e66f5a8a1650a24bcea43530e\"\n",
    "print(\"Original query:\\n\" + queries[example_qid])\n",
    "print(\"Topic-enriched query:\\n\" + topic_enriched_queries[example_qid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "16a4f94e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NDCG@10: 0.18645\n"
     ]
    }
   ],
   "source": [
    "T5_rerank_results_QEP = reranker.rerank(corpus, topic_enriched_queries, top_result, top_k=100)\n",
    "print('NDCG@10:', retriever.evaluate(qrels, T5_rerank_results_QEP, [10])[0]['NDCG@10'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch-g",
   "language": "python",
   "name": "torch-g"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}