Eval BERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT Evalscript\n",
    "\n",
    "In order to evaluate a BERT model, we follow the same procedure as training. Retrieve top-10 (100? 1000?) from Anserini and re-rank accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:06.304736Z",
     "start_time": "2020-05-24T17:04:06.261991Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from os.path import expanduser\n",
    "home = expanduser(\"~\")\n",
    "os.environ[\"JAVA_HOME\"] = f\"{home}/.sdkman/candidates/java/11.0.7.hs-adpt\"  #Set right JAVA version\n",
    "data_home = \"/ssd2/arthur/MsMarcoTREC/\"\n",
    "def path(x):\n",
    "    return os.path.join(data_home, x)\n",
    "\n",
    "try:\n",
    "    import pyserini\n",
    "except:\n",
    "    !pip install pyserini==0.9.2.0 # install pyserini\n",
    "    import pyserini\n",
    "try:\n",
    "    import tqdm\n",
    "except:\n",
    "    !pip install tqdm # Good for progress bars!\n",
    "    import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:08.194230Z",
     "start_time": "2020-05-24T17:04:06.962466Z"
    }
   },
   "outputs": [],
   "source": [
    "import jnius_config\n",
    "jnius_config.add_options('-Xmx32G') # Adjust to your machine.\n",
    "from pyserini.search import pysearch\n",
    "import subprocess\n",
    "from tqdm.auto import tqdm\n",
    "import random\n",
    "import pickle\n",
    "import sys\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import os\n",
    "from collections import defaultdict\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `pytrec_eval` setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T22:38:23.799097Z",
     "start_time": "2020-05-24T22:38:23.751122Z"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import pytrec_eval\n",
    "except:\n",
    "    !pip install pytrec_eval\n",
    "    import pytrec_eval\n",
    "    \n",
    "qrel = defaultdict(lambda:dict())\n",
    "qrel_path = path(\"qrels/2019qrels-docs.txt\")\n",
    "for line in open(qrel_path):\n",
    "    query_id, _, doc_id, rel = line.split()\n",
    "    qrel[query_id][doc_id] = int(rel)\n",
    "qrel = dict(qrel)\n",
    "metrics = {\"map\", \"ndcg\", \"recip_rank\"}\n",
    "evaluator = pytrec_eval.RelevanceEvaluator(qrel, metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:11.353988Z",
     "start_time": "2020-05-24T17:04:10.949429Z"
    }
   },
   "outputs": [],
   "source": [
    "index_path = path(\"lucene-index.msmarco-doc.pos+docvectors+rawdocs\")\n",
    "searcher = pysearch.SimpleSearcher(index_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Anserini top-K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:29.304350Z",
     "start_time": "2020-05-24T17:04:16.084872Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running retrieval step for 200 queries, on top-1000\n",
      "dumping results...\n"
     ]
    }
   ],
   "source": [
    "pattern = re.compile('([^\\s\\w]|_)+')\n",
    "\n",
    "searcher.set_bm25_similarity(0.9, 0.4)\n",
    "pairs = []\n",
    "threads = 42 # Number of Threads to use when retrieving\n",
    "k = 1000       # Number of documents to retrieve \n",
    "triples = []\n",
    "\n",
    "query_texts = dict()\n",
    "file_path = path(f\"queries/msmarco-test2019-queries.tsv\")\n",
    "run_search=True\n",
    "if os.path.isfile(\"test_triples.pkl\"):\n",
    "    print(f\"Already found file {file_path}. Cowardly refusing to run this again. Will only load querytexts.\")\n",
    "    pairs = pickle.load(open(path(\"test_triples.pkl\"), 'rb'))\n",
    "    run_search = False\n",
    "number_of_queries = int(subprocess.run(f\"wc -l {file_path}\".split(), capture_output=True).stdout.split()[0])\n",
    "print(f\"Running retrieval step for {number_of_queries} queries, on top-{k}\")\n",
    "queries = []\n",
    "query_ids = []\n",
    "for idx, line in enumerate(open(file_path, encoding=\"utf-8\")):\n",
    "    query_id, query = line.strip().split(\"\\t\")\n",
    "    query_ids.append(query_id)\n",
    "    query = unicodedata.normalize(\"NFKD\", query) # Force queries into UTF-8\n",
    "    query = pattern.sub(' ',query) # Remove non-ascii characters. It clears up most of the issues we may find on the query datasets\n",
    "    query_texts[query_id] = query\n",
    "    if run_search is False:\n",
    "        continue\n",
    "    queries.append(query)\n",
    "scores = defaultdict(lambda:dict())\n",
    "results = searcher.batch_search(queries, query_ids, k=k, threads=threads)\n",
    "\n",
    "#There is probably a one-liner for this...\n",
    "for qid in results.keys():\n",
    "    for result in results[qid]:\n",
    "        scores[qid][result.docid] = result.score\n",
    "        triples.append((qid, result.docid))\n",
    "\n",
    "print(f\"dumping results...\")\n",
    "pickle.dump(dict(scores), open(path(f\"test_triples.pkl\"), 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:35.838813Z",
     "start_time": "2020-05-24T17:04:35.703990Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Anserini MAP: 0.33099106298835235 nDCG: 0.5942674768376021\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "avg_map = np.mean([x[\"map\"] for x in evaluator.evaluate(scores).values()])\n",
    "avg_ndcg = np.mean([x[\"ndcg\"] for x in evaluator.evaluate(scores).values()])\n",
    "print(f\"Anserini MAP: {avg_map} nDCG: {avg_ndcg}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T21:59:01.751262Z",
     "start_time": "2020-05-24T21:59:01.697170Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# This is a copy from the MsMarcoDataset Class.\n",
    "from torch.utils.data import Dataset\n",
    "import torch\n",
    "\n",
    "# This is our main Dataset class.\n",
    "class MsMarcoDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 samples,\n",
    "                 tokenizer,\n",
    "                 searcher,\n",
    "                 split,\n",
    "                 tokenizer_batch=8000):\n",
    "        '''Initialize a Dataset object. \n",
    "        Arguments:\n",
    "            samples: A list of samples. Each sample should be a tuple with (query_id, doc_id, <label>), where label is optional\n",
    "            tokenizer: A tokenizer object from Hugging Face's Tokenizer lib. (need to implement encode_batch())\n",
    "            searcher: A PySerini Simple Searcher object. Should implement the .doc() method\n",
    "            split: A strong indicating if we are in a train, dev or test dataset.\n",
    "            tokenizer_batch: How many samples to be tokenized at once by the tokenizer object.\n",
    "            The biggest bottleneck is the searcher, not the tokenizer.\n",
    "        '''\n",
    "        self.searcher = searcher\n",
    "        self.split = split\n",
    "        # If we already have the data pre-computed, we shouldn't need to re-compute it.\n",
    "        self.split = split\n",
    "        if (os.path.isfile(path(f\"{split}_msmarco_samples.tsv\"))\n",
    "                and os.path.isfile(path(f\"{split}_msmarco_offset.pkl\"))\n",
    "                and os.path.isfile(path(f\"{split}_msmarco_index.pkl\"))):\n",
    "            print(\"Already found every meaningful file. Cowardly refusing to re-compute.\")\n",
    "            self.samples_offset_dict = pickle.load(open(path(f\"{split}_msmarco_offset.pkl\"), 'rb'))\n",
    "            self.index_dict = pickle.load(open(path(f\"{split}_msmarco_index.pkl\"), 'rb'))\n",
    "            return\n",
    "        self.tokenizer = tokenizer\n",
    "        print(\"Loading and tokenizing dataset...\")\n",
    "        self.samples_offset_dict = dict()\n",
    "        self.index_dict = dict()\n",
    "\n",
    "        self.samples_file = open(path(f\"{split}_msmarco_samples.tsv\"),'w',encoding=\"utf-8\")\n",
    "        self.processed_samples = 0\n",
    "        query_batch = []\n",
    "        doc_batch = []\n",
    "        sample_ids_batch = []\n",
    "        labels_batch = []\n",
    "        number_of_batches = math.ceil(len(samples) // tokenizer_batch)\n",
    "        print(number_of_batches)\n",
    "        # A progress bar to display how far we are.\n",
    "        batch_pbar = tqdm(total=number_of_batches, desc=\"Tokenizer batches\")\n",
    "        for i, sample in enumerate(samples):\n",
    "            if split==\"train\" or split == \"dev\":\n",
    "                label = sample[2]\n",
    "                labels_batch.append(label)\n",
    "            query_batch.append(query_texts[sample[0]])\n",
    "            doc_batch.append(self._get_document_content_from_id(sample[1]))\n",
    "            sample_ids_batch.append(f\"{sample[0]}_{sample[1]}\")\n",
    "            #If we hit the number of samples for this batch OR this is the last sample\n",
    "            if len(query_batch) == tokenizer_batch or i == len(samples) - 1:\n",
    "                self._tokenize_and_dump_batch(doc_batch, query_batch, labels_batch, sample_ids_batch)\n",
    "                batch_pbar.update()\n",
    "                query_batch = []\n",
    "                doc_batch = []\n",
    "                sample_ids_batch = []\n",
    "                if split == \"train\" or split == \"dev\":\n",
    "                    labels_batch = []\n",
    "        batch_pbar.close()\n",
    "        # Dump files in disk, so we don't need to go over it again.\n",
    "        self.samples_file.close()\n",
    "        pickle.dump(self.index_dict, open(path(f\"{self.split}_msmarco_index.pkl\"), 'wb'))\n",
    "        pickle.dump(self.samples_offset_dict, open(path(f\"{self.split}_msmarco_offset.pkl\"), 'wb'))\n",
    "\n",
    "    def _tokenize_and_dump_batch(self, doc_batch, query_batch, labels_batch,\n",
    "                                 sample_ids_batch):\n",
    "        '''tokenizes and dumps the samples in the current batch\n",
    "        It also store the positions from the current file into the samples_offset_dict.\n",
    "        '''\n",
    "        # Use the tokenizer object\n",
    "        tokens = self.tokenizer.encode_batch(list(zip(query_batch, doc_batch)))\n",
    "        for idx, (sample_id, token) in enumerate(zip(sample_ids_batch, tokens)):\n",
    "            #BERT supports up to 512 tokens. If we have more than that, we need to remove some tokens from the document\n",
    "            if len(token.ids) >= 512:\n",
    "                token_ids = token.ids[:511]\n",
    "                token_ids.append(tokenizer.token_to_id(\"[SEP]\"))\n",
    "                segment_ids = token.type_ids[:512]\n",
    "            # With less tokens, we need to \"pad\" the vectors up to 512.\n",
    "            else:\n",
    "                padding = [0] * (512 - len(token.ids))\n",
    "                token_ids = token.ids + padding\n",
    "                segment_ids = token.type_ids + padding\n",
    "            # How far in the file are we? This is where we need to go to find the documents later.\n",
    "            file_location = self.samples_file.tell()\n",
    "            # If we have labels\n",
    "            if self.split==\"train\" or self.split == \"dev\":\n",
    "                self.samples_file.write(f\"{sample_id}\\t{token_ids}\\t{segment_ids}\\t{labels_batch[idx]}\\n\")\n",
    "            else:\n",
    "                self.samples_file.write(f\"{sample_id}\\t{token_ids}\\t{segment_ids}\\n\")\n",
    "            self.samples_offset_dict[sample_id] = file_location\n",
    "            self.index_dict[self.processed_samples] = sample_id\n",
    "            self.processed_samples += 1\n",
    "\n",
    "    def _get_document_content_from_id(self, doc_id):\n",
    "        '''Get the raw text value from the doc_id\n",
    "        There is probably an easier way to do that, but this works.\n",
    "        '''\n",
    "        doc_text = self.searcher.doc(doc_id).lucene_document().getField(\"raw\").stringValue()\n",
    "        return doc_text[7:-8]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        '''Returns a sample with index idx\n",
    "        DistilBERT does not take into account segment_ids. (indicator if the token comes from the query or the document) \n",
    "        However, for the sake of completness, we are including it here, together with the attention mask\n",
    "        position_ids, with the positional encoder, is not needed. It's created for you inside the model.\n",
    "        '''\n",
    "        if isinstance(idx, int):\n",
    "            idx = self.index_dict[idx]\n",
    "        with open(path(f\"{self.split}_msmarco_samples.tsv\"), 'r', encoding=\"utf-8\") as inf:\n",
    "            inf.seek(self.samples_offset_dict[idx])\n",
    "            line = inf.readline().split(\"\\t\")\n",
    "            try:\n",
    "                sample_id = line[0]\n",
    "                qid, did = sample_id.split(\"_\")\n",
    "                input_ids = eval(line[1])\n",
    "                token_type_ids = eval(line[2])\n",
    "                input_mask = [1] * 512\n",
    "            except:\n",
    "                print(line, idx)\n",
    "                raise IndexError\n",
    "            # If it's a training dataset, we also have a label tag.\n",
    "            if self.split==\"train\" or self.split == \"dev\":\n",
    "                label = int(line[3])\n",
    "                return (torch.tensor(input_ids, dtype=torch.long),\n",
    "                        torch.tensor(input_mask, dtype=torch.long),\n",
    "                        torch.tensor(token_type_ids, dtype=torch.long),\n",
    "                        torch.tensor([label], dtype=torch.long),\n",
    "                       qid,\n",
    "                       did)\n",
    "            return (torch.tensor(input_ids, dtype=torch.long),\n",
    "                    torch.tensor(input_mask, dtype=torch.long),\n",
    "                    torch.tensor(token_type_ids, dtype=torch.long),\n",
    "                   qid,\n",
    "                   did)\n",
    "    def __len__(self):\n",
    "        return len(self.samples_offset_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:04:58.627430Z",
     "start_time": "2020-05-24T17:04:57.876342Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import DistilBertForSequenceClassification\n",
    "from torch.utils.data import DataLoader\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "\n",
    "tokenizer = BertWordPieceTokenizer(\"/ssd2/arthur/bert-axioms/tokenizer/bert-base-uncased-vocab.txt\", lowercase=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T21:59:04.245596Z",
     "start_time": "2020-05-24T21:59:04.045957Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Already found every meaningful file. Cowardly refusing to re-compute.\n"
     ]
    }
   ],
   "source": [
    "eval_dataset = MsMarcoDataset(triples, tokenizer, searcher, split=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T21:59:08.151187Z",
     "start_time": "2020-05-24T21:59:08.120459Z"
    },
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([  101, 12098,  7877,  6210,   102, 16770,  1024,  1013,  1013, 29393,\n",
       "         15007, 12782,  1012,  4012,  1013, 25022, 11365,  1013,  3193, 14399,\n",
       "          2594,  1012, 25718,  1029,  1056,  1027, 19843,  2692, 14142, 22166,\n",
       "         15072,  1024,  2024,  2027,  4276,  9343,  1029, 10682,  2332, 29393,\n",
       "         10975,  6679, 25969,  2271,  5068,  1024, 11703,  2538,  1010,  2456,\n",
       "         19894,  2015,  1024, 21211,  2575, 19894,  2098,  1024, 12256, 12022,\n",
       "          2324,  1010,  2289,  1016,  1024,  2459,  2572,  1045,  1005,  1049,\n",
       "          2349,  2005,  1037,  2047,  3940,  1997, 20422,  7877,  1998,  1045,\n",
       "          1005,  2310,  2042,  4994,  2070,  2204,  2477,  2055,  6653, 15072,\n",
       "          1012,  2129,  2024,  2027,  1029,  2157,  2085,  1045,  2031, 20422,\n",
       "          7877,  1998, 20422, 17072,  2021,  2049,  2428,  1037,  3255,  2000,\n",
       "          4287,  2105,  2048,  7689,  1997,  7877,  2043,  1045,  2175,  2000,\n",
       "          1037,  4323,  2380,  1006,  2073,  1045,  1005,  1049,  2648,  1998,\n",
       "          2503,  2035,  2154,  1007,  1012,  2009,  4165,  2066, 22166, 15072,\n",
       "          2052,  9611,  2035,  1997,  2122,  3471,  2021,  1045,  4687,  2129,\n",
       "          2092,  2027,  2147,  2043,  4102,  2000,  1037,  4056,  2275,  1997,\n",
       "          3180,  1009, 17072,  1012,  2151, 10740,  1029,  2821,  1010,  1998,\n",
       "          2129,  2172,  4469,  2079, 22166, 15072,  3465,  1029,  3746,  3478,\n",
       "          2000,  7170,  1011,  3046, 12106,  1999,  2047,  3332,  4349, 23301,\n",
       "         29393,  4190, 15590,  8009,  2483,  3802,  4942, 22483,  2953, 13012,\n",
       "          8286,  1024,  3438, 20702,  9442,  1998,  4175,  2378,  1005,   999,\n",
       "          5068,  1024, 13292,  2423,  1010,  2494, 19894,  2015,  1024, 11502,\n",
       "          2575,  2475, 19894,  2098,  1024, 12256, 12022,  2324,  1010,  2289,\n",
       "          1016,  1024,  2539,  2572,  1045,  2145,  2228,  2027,  2298,  2066,\n",
       "         25577,  7877,  1011,  2130,  2295,  2111, 18292,  2027,  2123,  1005,\n",
       "          1056,  1012,  1037,  2450,  2012,  2026,  2283,  2197,  2305,  2018,\n",
       "          2068,  1011,  5791,  1037, 25577,  2298,  2000,  2068,  1012,  1045,\n",
       "          2123,  1005,  1056,  2066,  2068,  2172,  1010,  2000,  2022,  7481,\n",
       "          1012,  2028,  2173,  2073,  2027,  2052,  2022,  6179,  1010,  2043,\n",
       "          4439,  1999,  1996,  2482,  1010,  2027,  2123,  1005,  1056,  2147,\n",
       "          1012,  8825, 10556, 13620,  1000,  9543,  5484,  2121,  2013,  3109,\n",
       "          1000, 29393, 10975,  6679, 25969,  2271,  3802,  4942, 22483,  2953,\n",
       "          5068,  1024,  5553,  2184,  1010,  2541, 19894,  2015,  1024,  4868,\n",
       "          2581,  2509, 19894,  2098,  1024, 12256, 12022,  2324,  1010,  2289,\n",
       "          1016,  1024,  2382,  2572,  1045,  1005,  2310,  6247, 22166,  2005,\n",
       "          1996,  2197,  2093,  2086,  1998,  2027,  2031,  2042, 10392,  1012,\n",
       "          1999,  2755,  1045,  2074,  3641,  1037,  2047,  3940,  1997,  7877,\n",
       "          1998,  2027,  2205,  2097,  2031,  1996,  6653, 18898,  1012,  2009,\n",
       "          1005,  1055,  2428,  3243,  3835,  2025,  2383,  2000,  4737,  2055,\n",
       "          3239,  3860,  2076,  1996,  2154,  2043,  3048,  2013, 24274,  2000,\n",
       "         19350,  1012,  2115,  2159,  2131, 17730,  2000,  1996,  2689,  2107,\n",
       "          2008,  2017,  2180,  1005,  1056,  5060,  1996,  6653,  4902,  1012,\n",
       "          1996,  2069,  2051,  2009,  1005,  1055, 16880,  2003,  3788,  2013,\n",
       "          2648,  2000,  2503,  2021,  1996,  2689,  2067,  2000,  3154,  1008,\n",
       "          2069,  3138,  1037,  2261,  2781,  1012,  4009, 12221,  1024,  4439,\n",
       "          1011,  2115, 19521,  2038,  1037, 23068, 18898,  2061,  2027,  2123,\n",
       "          1005,  1056,  2147,  2005,  4439,  1012,  5855,  1011,  1045,  5293,\n",
       "          2008,  2026,  7877,  2031, 13755,  1998,  2947,  2411,  2058, 14451,\n",
       "          2026,  4620,  1012,  1045,  2031,  2000, 10825,  2870,  2000,  4638,\n",
       "          7524,  2007,  2026,  6248,  3239,  2738,  2084,  2083,  1996, 25577,\n",
       "         15072,  1012,  1008,  3154,  1024,  2043,  1996, 18898,  2003,  2047,\n",
       "          1996,   102]),\n",
       " tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1]),\n",
       " tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1]),\n",
       " '20455',\n",
       " 'D840434')"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T17:05:02.536459Z",
     "start_time": "2020-05-24T17:05:02.516321Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# In a real-world scenario, this would be in a separate file, and you would just import this.\n",
    "import torch\n",
    "from torch import nn\n",
    "from transformers import DistilBertModel, BertModel\n",
    "\n",
    "class BertRelevanceRanker(nn.Module):\n",
    "    def __init__(self, model=\"distilbert-base-uncased\"):\n",
    "        \"\"\"Creates an instance of Bert Relevance Ranker. \n",
    "        It feeds two senteces into a pre-trained BERT model, extracts the [CLS] token and feeds it into a one-layer FFNN\"\"\"\n",
    "        super().__init__()\n",
    "        self.distil = False\n",
    "        self.loss_fct = nn.CrossEntropyLoss()\n",
    "        if \"distil\" in model:\n",
    "            self.distil = True\n",
    "            self.bert = DistilBertModel.from_pretrained(model)\n",
    "        else:\n",
    "            self.bert = BertModel.from_pretrained(model)\n",
    "        self.config = self.bert.config\n",
    "        self.linear1 = nn.Linear(self.bert.config.dim, self.bert.config.dim)\n",
    "        self.linear2 = nn.Linear(self.bert.config.dim, 2)\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "    \n",
    "    def forward(self, input_ids, attention_mask, token_type_ids=None, labels=None):\n",
    "        if not self.distil and token_type_ids is None:\n",
    "            raise ValueError(\"Model is not distilBERT and it did not received token_type_ids!\")\n",
    "        if not self.distil:\n",
    "            outputs = self.bert(input_ids, attention_mask, token_type_ids)\n",
    "        else:\n",
    "            pooled_output = self.bert(input_ids, attention_mask)[0][:, 0]\n",
    "        pooled_output = self.linear1(pooled_output)\n",
    "        pooled_output = nn.ReLU()(pooled_output)\n",
    "        pooled_output = self.dropout(pooled_output)\n",
    "        logits = self.linear2(pooled_output)\n",
    "        outputs = (logits,)\n",
    "        if labels is not None:\n",
    "            loss = self.loss_fct(logits.view(-1, 2), labels.view(-1))\n",
    "            outputs = (loss, ) + outputs\n",
    "        return outputs\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T22:03:17.080386Z",
     "start_time": "2020-05-24T22:03:13.507913Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running on 2 GPUS, on 256-sized batches\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "try:\n",
    "    del model\n",
    "    torch.cuda.empty_cache() # Make sure we have a clean slate. Usefull in a Notebook.\n",
    "except:\n",
    "    pass\n",
    "\n",
    "GPUS_TO_USE = [0, 1] # If you have multiple GPUs, pick the ones you want to use.\n",
    "number_of_cpus = 24 # Number of CPUS to use when loading your dataset.\n",
    "model = BertRelevanceRanker()\n",
    "model.load_state_dict(torch.load(path(f\"models/distilBERT-2020-05-21/pytorch_model.bin\"))) #load last model saved\n",
    "model.eval()\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model = torch.nn.DataParallel(model, device_ids=GPUS_TO_USE)\n",
    "    device = torch.device(f\"cuda:{GPUS_TO_USE[0]}\") \n",
    "    model.to(device)\n",
    "    batch_size = len(GPUS_TO_USE) * 128 #Eval is WAAY smaller than train. We can load a fairly large batch here.\n",
    "    print(f\"running on {len(GPUS_TO_USE)} GPUS, on {batch_size}-sized batches\")\n",
    "else:\n",
    "    print(\"Are you sure about it? We will try to run this in CPU, but it's a BAD idea...\")\n",
    "    device = torch.device(\"cpu\")\n",
    "    batch_size = 16\n",
    "    model.to(device)\n",
    "\n",
    "data_loader = DataLoader(eval_dataset, batch_size=batch_size, num_workers=number_of_cpus, shuffle=False) # Don't shuffle!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T22:18:08.973671Z",
     "start_time": "2020-05-24T22:03:19.338815Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "829ac7ea63224258b446ad5321bf62e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Batches', max=782.0, style=ProgressStyle(description_widt…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "bert_scores = defaultdict(lambda:dict())\n",
    "for batch in tqdm(data_loader, desc=\"Batches\", total=len(data_loader)):\n",
    "    with torch.no_grad():\n",
    "        inputs = {'input_ids': batch[0].to(device),\n",
    "          'attention_mask': batch[1].to(device)}\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter(\"ignore\")\n",
    "            outputs = model(**inputs)\n",
    "        predict_scores = nn.Softmax(dim=1)(outputs[0])[:,1].detach().cpu().numpy().flatten()\n",
    "        for qid, did, score in zip(batch[3], batch[4], predict_scores):\n",
    "            bert_scores[qid][did] = score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-24T22:47:13.294709Z",
     "start_time": "2020-05-24T22:47:01.753345Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.0 MAP@1000: 0.16087 nDCG@1000: 0.45354 MRR@1000: 0.35727\n",
      " 0.05 MAP@1000: 0.18849 nDCG@1000: 0.47773 MRR@1000: 0.56659\n",
      "  0.1 MAP@1000: 0.20094 nDCG@1000: 0.48724 MRR@1000: 0.57985\n",
      " 0.15 MAP@1000: 0.21243 nDCG@1000: 0.50081 MRR@1000: 0.63296\n",
      "  0.2 MAP@1000: 0.22177 nDCG@1000: 0.50929 MRR@1000: 0.67885\n",
      " 0.25 MAP@1000: 0.23293 nDCG@1000: 0.52206 MRR@1000: 0.72644\n",
      "  0.3 MAP@1000: 0.24816 nDCG@1000: 0.54091 MRR@1000: 0.78575\n",
      " 0.35 MAP@1000: 0.26332 nDCG@1000: 0.55697 MRR@1000: 0.82193\n",
      "  0.4 MAP@1000: 0.27895 nDCG@1000: 0.57372 MRR@1000: 0.87016\n",
      " 0.45 MAP@1000: 0.29465 nDCG@1000: 0.58715 MRR@1000: 0.91805\n",
      "  0.5 MAP@1000: 0.30808 nDCG@1000: 0.59671 MRR@1000: 0.91938\n",
      " 0.55 MAP@1000: 0.3215 nDCG@1000: 0.60492 MRR@1000: 0.9186\n",
      "  0.6 MAP@1000: 0.33528 nDCG@1000: 0.61219 MRR@1000: 0.9186\n",
      " 0.65 MAP@1000: 0.34472 nDCG@1000: 0.61728 MRR@1000: 0.89903\n",
      "  0.7 MAP@1000: 0.35086 nDCG@1000: 0.61895 MRR@1000: 0.88566\n",
      " 0.75 MAP@1000: 0.35315 nDCG@1000: 0.6192 MRR@1000: 0.88218\n",
      "  0.8 MAP@1000: 0.35119 nDCG@1000: 0.61505 MRR@1000: 0.86217\n",
      " 0.85 MAP@1000: 0.34703 nDCG@1000: 0.60861 MRR@1000: 0.84199\n",
      "  0.9 MAP@1000: 0.34127 nDCG@1000: 0.60372 MRR@1000: 0.83367\n",
      " 0.95 MAP@1000: 0.33665 nDCG@1000: 0.5997 MRR@1000: 0.80644\n",
      "  1.0 MAP@1000: 0.33099 nDCG@1000: 0.59427 MRR@1000: 0.80447\n",
      "****Best alpha is   0.5, with nDCG@1000: 0.59671 MAP@1000: 0.30808 MRR@1000: 0.91938****\n"
     ]
    }
   ],
   "source": [
    "alphas = np.arange(0.0, 1.05, 0.05) # It's a nice idea to combine BM25 scores with BERT scores. Let's see which value is the best here.\n",
    "best_alpha = -1\n",
    "best_ndcg = 0.0\n",
    "best_map = 0.0\n",
    "best_mrr = 0.0\n",
    "top_k = 1000\n",
    "for alpha in alphas:\n",
    "    final_scores = defaultdict(lambda: dict())\n",
    "    for qid in bert_scores:\n",
    "        max_score = max(scores[qid].values())\n",
    "        min_score = min(scores[qid].values())\n",
    "        normalized_scores_anserini = {k : (v-min_score)/(max_score-min_score) for k, v in scores[qid].items()}\n",
    "        max_score = max(bert_scores[qid].values())\n",
    "        min_score = min(bert_scores[qid].values())\n",
    "        normalized_scores_bert = {k : (v-min_score)/(max_score-min_score) for k, v in bert_scores[qid].items()}\n",
    "        for did in normalized_scores_anserini:\n",
    "            final_scores[qid][did] = (alpha * normalized_scores_anserini[did]  + (1-alpha) * normalized_scores_bert[did])\n",
    "        final_scores[qid] = dict(sorted(final_scores[qid].items(), key=lambda x: x[1], reverse=True)[:top_k])\n",
    "    avg_map = np.mean([x[\"map\"] for x in evaluator.evaluate(final_scores).values()])\n",
    "    avg_ndcg = np.mean([x[\"ndcg\"] for x in evaluator.evaluate(final_scores).values()])\n",
    "    avg_mrr = np.mean([x[\"recip_rank\"] for x in evaluator.evaluate(final_scores).values()])\n",
    "    if avg_mrr > best_mrr:\n",
    "        best_ndcg = avg_ndcg\n",
    "        best_alpha = alpha\n",
    "        best_map = avg_map\n",
    "        best_mrr = avg_mrr\n",
    "    print(f\"{alpha:5.2} MAP@{top_k}: {avg_map:.5} nDCG@{top_k}: {avg_ndcg:.5} MRR@{top_k}: {avg_mrr:.5}\")\n",
    "print(f\"****Best alpha is {best_alpha:5.2}, with nDCG@{top_k}: {best_ndcg:.5} MAP@{top_k}: {best_map:.5f} MRR@{top_k}: {best_mrr:.5f}****\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7 (Bert4IR)",
   "language": "python",
   "name": "myenv"
  },
  "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.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}