01_module_passage.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a19ce1-bb8a-40c1-8ef3-ff78f60b71d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp passage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885c918b-5721-4a0d-947e-62451a6d4c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec96dcb-56ee-412f-bc63-253c5b1a19b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from nbdev import *\n",
    "from fastcore.test import *\n",
    "from fastcore.utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f00408-3ef3-475c-8ab4-eea43088ecae",
   "metadata": {},
   "source": [
    "# passage\n",
    "> Reference API related to the passage ranking use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "806e548c-5588-4a0a-b4d5-e744bbe1fc1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "import ir_datasets \n",
    "import json\n",
    "import requests\n",
    "import os \n",
    "import pandas as pd\n",
    "\n",
    "from fastcore.utils import patch, patch_to\n",
    "from typing import Dict, Tuple, Optional, List\n",
    "import random\n",
    "from vespa.package import (\n",
    "    ApplicationPackage, \n",
    "    Field, \n",
    "    FieldSet, \n",
    "    RankProfile, \n",
    "    QueryField\n",
    ")\n",
    "from learntorank.query import QueryModel\n",
    "from vespa.deployment import VespaDocker\n",
    "from learntorank.evaluation import EvalMetric, evaluate\n",
    "from learntorank.stats import compute_evaluation_estimates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e108d22a-2827-4fd0-be8d-2c33e303fa4b",
   "metadata": {},
   "source": [
    "## Data manipulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6ba44d-8b70-4bde-a7b0-126df311773f",
   "metadata": {},
   "source": [
    "Code related to the manipulation of passage ranking data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cf719ad-e880-4100-9db3-75aff28a64f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def sample_dict_items(\n",
    "    d: Dict, # dict to be samples from.                     \n",
    "    n: int   # Number of samples \n",
    ") -> Dict: # dict with sampled values\n",
    "    \"Sample items from a dict.\"\n",
    "    n = min(len(d), n)\n",
    "    return {\n",
    "        k: d[k]\n",
    "        for k in random.sample(\n",
    "            k=n, population=sorted(d)\n",
    "        )\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed1b5ab6-091e-4bc3-a37c-476fc2a5c075",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3ea1a33-3224-4f4c-937f-6f245bbca954",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = {\"a\": 1, \"b\":2, \"c\":3}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e864ed9-0fcc-4a90-ab4d-4a3b501d6f34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': 1}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_dict_items(d, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6ff59d1-155f-4f61-922a-a62ff3292274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'c': 3, 'b': 2}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_dict_items(d, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61247c74-3fca-49f0-8567-a822773351f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': 1, 'c': 3, 'b': 2}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_dict_items(d, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08fdcfce-4877-4643-a972-caa4e0a22a72",
   "metadata": {},
   "source": [
    "Return full dict in case number of samples is higher than length of the dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "923f2213-b8d7-4d40-b080-af6aa3c5988f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'a': 1, 'c': 3, 'b': 2}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_dict_items(d, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f6dd73f-fd43-4970-80d5-3dcc42fb079f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(len(sample_dict_items(d, 4)), 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b31d818-5173-4163-90a2-2c7f1d7735ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def save_data(\n",
    "    corpus: Dict,  # Document corpus, see usage example below. \n",
    "    train_qrels: Dict,  # Training relevance scores, see usage example below.  \n",
    "    train_queries: Dict,  # Training queries, see usage example below.   \n",
    "    dev_qrels: Dict,  #  Development relevance scores, see usage example below.  \n",
    "    dev_queries: Dict,  # Development queries, see usage example below.   \n",
    "    file_path: str = \"passage_sample.json\"  # Valid JSON file path.\n",
    ") -> None:  # Side-effect: data is saved to `file_path`.\n",
    "    \"\"\"Save data to disk.\n",
    "    \n",
    "    The main goal is to save sample data to disk.\n",
    "    \"\"\"\n",
    "    data = {\n",
    "        \"corpus\": corpus, \n",
    "        \"train_qrels\": train_qrels, \n",
    "        \"train_queries\": train_queries, \n",
    "        \"dev_qrels\": dev_qrels, \n",
    "        \"dev_queries\": dev_queries\n",
    "    }\n",
    "    with open(file_path, \"w\") as f:\n",
    "        json.dump(data, f)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c039b4d7-3675-4ad8-84b2-b30cc90167b7",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33aca50d-6154-4005-a075-0cf7a1d532c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = {\n",
    "    \"0\": \"sentence 0\", \n",
    "    \"1\": \"sentence 1\", \n",
    "    \"2\": \"sentence 2\", \n",
    "    \"3\": \"sentence 3\"\n",
    "}\n",
    "train_queries = {\n",
    "    \"10\": \"train query 10\",\n",
    "    \"11\": \"train query 11\"\n",
    "}\n",
    "train_qrels = {\n",
    "    \"10\": {\"0\": 1},\n",
    "    \"11\": {\"2\": 1}\n",
    "}\n",
    "dev_queries = {\n",
    "    \"20\": \"train query 20\",\n",
    "    \"21\": \"train query 21\"\n",
    "}\n",
    "dev_qrels = {\n",
    "    \"20\": {\"1\": 1},\n",
    "    \"21\": {\"3\": 1}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d655ae42-1d98-4d9e-af89-b546a8688d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_data(\n",
    "    corpus, \n",
    "    train_qrels, train_queries, \n",
    "    dev_qrels, dev_queries, \n",
    "    file_path=\"passage_sample.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9340d0b6-fd8e-446a-a714-27605e1aec9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def load_data(\n",
    "    file_path: Optional[str] = None  # valid JSON file path contain data saved by `save_data`. If `None`, a pre-generated sample will be downloaded.\n",
    ") -> Dict:  # See usage example below for expected format.  \n",
    "    \"\"\"Load data.\n",
    "    \n",
    "    The main goal is to load sample data from disk.\n",
    "    If a `file_path` is not provided, a pre-generated data sample \n",
    "    will be downloaded.\n",
    "    \"\"\"\n",
    "    if not file_path:\n",
    "        file_url = \"https://data.vespa.oath.cloud/learntorank/passage/passage_sample.json\"\n",
    "        data_json = json.loads(\n",
    "            requests.get(file_url).text\n",
    "        )\n",
    "    else:\n",
    "        with open(file_path, \"r\") as f:\n",
    "            data_json = json.load(f)\n",
    "        \n",
    "    return {\n",
    "        \"corpus\": data_json[\"corpus\"], \n",
    "        \"train_qrels\": data_json[\"train_qrels\"], \n",
    "        \"train_queries\": data_json[\"train_queries\"],\n",
    "        \"dev_qrels\": data_json[\"dev_qrels\"],\n",
    "        \"dev_queries\": data_json[\"dev_queries\"]\n",
    "    }\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb06c096-9013-4232-a720-649897740f4c",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a6f1e9f-3cc3-4f51-b233-81756fc00686",
   "metadata": {},
   "source": [
    "* With `file_path`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d2b439e-e992-4a62-b826-82ab173aca10",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_data(\"passage_sample.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da42fec9-3393-4e58-8d14-5d18c0eb129e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'corpus': {'0': 'sentence 0',\n",
       "  '1': 'sentence 1',\n",
       "  '2': 'sentence 2',\n",
       "  '3': 'sentence 3'},\n",
       " 'train_qrels': {'10': {'0': 1}, '11': {'2': 1}},\n",
       " 'train_queries': {'10': 'train query 10', '11': 'train query 11'},\n",
       " 'dev_qrels': {'20': {'1': 1}, '21': {'3': 1}},\n",
       " 'dev_queries': {'20': 'train query 20', '21': 'train query 21'}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb4ebd8f-5cda-4444-8b94-a29ed7d3b5e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    data, \n",
    "    {\n",
    "        \"corpus\": corpus, \n",
    "        \"train_qrels\": train_qrels, \n",
    "        \"train_queries\": train_queries, \n",
    "        \"dev_qrels\": dev_qrels, \n",
    "        \"dev_queries\": dev_queries\n",
    "    }\n",
    ")\n",
    "os.remove(\"passage_sample.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12fb74d9-06a2-4463-a866-389cf6fd0993",
   "metadata": {},
   "source": [
    "* Without `file_path` specified, a pre-generated sample data will be downloaded:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa44928-8ff4-4ab9-b63a-0fc67240cbe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac49a77f-91cd-4f83-a8a4-3113345b73dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['corpus', 'train_qrels', 'train_queries', 'dev_qrels', 'dev_queries'])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1376bc5d-2bcd-4711-8a1b-d0a5d5c5401a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data[\"corpus\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7dd90e1-3733-49a8-874c-e2fa0e21855d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "class PassageData(object):\n",
    "    def __init__(\n",
    "        self, \n",
    "        corpus: Optional[Dict] = None,  # Document corpus, see usage example below.\n",
    "        train_qrels: Optional[Dict] = None,  # Training relevance scores, see usage example below.  \n",
    "        train_queries: Optional[Dict] = None,  # Training queries, see usage example below.   \n",
    "        dev_qrels: Optional[Dict] = None,  #  Development relevance scores, see usage example below.  \n",
    "        dev_queries: Optional[Dict] = None,  # Development queries, see usage example below.   \n",
    "    ):\n",
    "        \"Container for passage data\"\n",
    "        self.corpus = corpus\n",
    "        self.train_qrels = train_qrels\n",
    "        self.train_queries = train_queries\n",
    "        self.dev_qrels = dev_qrels\n",
    "        self.dev_queries = dev_queries\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return \"PassageData({}, {}, {}, {}, {})\".format(\n",
    "            \"corpus\" if self.corpus else None,\n",
    "            \"train_qrels\" if self.train_qrels else None,\n",
    "            \"train_queries\" if self.train_queries else None,\n",
    "            \"dev_qrels\" if self.dev_qrels else None,\n",
    "            \"dev_queries\" if self.dev_queries else None,            \n",
    "        )\n",
    "    \n",
    "    def __eq__(self, obj: \"PassageData\"):\n",
    "        return (self.corpus == obj.corpus and \n",
    "                self.train_qrels == obj.train_qrels and \n",
    "                self.train_queries == obj.train_queries and \n",
    "                self.dev_qrels == obj.dev_qrels and \n",
    "                self.dev_queries == obj.dev_queries)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b20ca01-e845-40cc-bd97-efee2a406c8d",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef64e95-b4b8-400e-b806-fa17f562c5c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = {\n",
    "    \"0\": \"sentence 0\", \n",
    "    \"1\": \"sentence 1\", \n",
    "    \"2\": \"sentence 2\", \n",
    "    \"3\": \"sentence 3\"\n",
    "}\n",
    "train_queries = {\n",
    "    \"10\": \"train query 10\",\n",
    "    \"11\": \"train query 11\"\n",
    "}\n",
    "train_qrels = {\n",
    "    \"10\": {\"0\": 1},\n",
    "    \"11\": {\"2\": 1}\n",
    "}\n",
    "dev_queries = {\n",
    "    \"20\": \"train query 20\",\n",
    "    \"21\": \"train query 21\"\n",
    "}\n",
    "dev_qrels = {\n",
    "    \"20\": {\"1\": 1},\n",
    "    \"21\": {\"3\": 1}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16845c5d-d792-4a3d-bb83-95db87f1407f",
   "metadata": {},
   "outputs": [],
   "source": [
    "passage_data = PassageData(\n",
    "    corpus=corpus, \n",
    "    train_queries = train_queries, \n",
    "    train_qrels=train_qrels,\n",
    "    dev_queries = dev_queries,\n",
    "    dev_qrels = dev_qrels\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e5e339-ba5b-4205-94b3-5f34c47f358c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PassageData(corpus, train_qrels, train_queries, dev_qrels, dev_queries)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "507e971f-bb3b-499c-8507-9809e7191c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def save(self: PassageData, file_path: str = 'passage_sample.json'):\n",
    "    save_data(\n",
    "        corpus=self.corpus, \n",
    "        train_qrels=self.train_qrels, \n",
    "        train_queries=self.train_queries, \n",
    "        dev_qrels=self.dev_qrels, \n",
    "        dev_queries=self.dev_queries,\n",
    "        file_path=file_path\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d404f4-179d-42f3-a341-b038dac7cdc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "assert os.path.exists(\"passage_sample.json\") == False, \"File exists.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae056b0b-2256-4376-9d1b-6642b93b2a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "passage_data.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e03cf3e3-28d7-4421-b654-cdb2e428ef4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "assert os.path.exists(\"passage_sample.json\") == True, \"File does not exists.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410663bc-52a1-437c-a40f-d08009d3d5b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch_to(PassageData, cls_method=True)\n",
    "def load(\n",
    "    cls, file_path: Optional[str] = None  # valid JSON file path contain data saved by save_data. If None, a pre-generated sample will be downloaded.\n",
    ") -> PassageData:\n",
    "    \"Load passage data from disk.\"\n",
    "    data = load_data(file_path=file_path)\n",
    "    return PassageData(\n",
    "        corpus=data.get(\"corpus\", None), \n",
    "        train_qrels=data.get(\"train_qrels\", None), \n",
    "        train_queries=data.get(\"train_queries\", None),\n",
    "        dev_qrels=data.get(\"dev_qrels\", None),\n",
    "        dev_queries=data.get(\"dev_queries\", None),        \n",
    "    )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87e6505d-4363-4a5f-a712-48ca6aebd0c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = PassageData.load(file_path=\"passage_sample.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba41f4e7-7f6a-47a7-ad58-cb9e57a96307",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(data, passage_data)\n",
    "os.remove(\"passage_sample.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37fb1a5-8970-4b99-bc41-e5d14c54759f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PassageData(corpus, train_qrels, train_queries, dev_qrels, dev_queries)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1680f68c-c462-432f-bd65-a3ad36f5f2c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch(as_prop=True)\n",
    "def summary(self: PassageData):\n",
    "    \"Summary of the size of the dataset components.\"\n",
    "    print(\n",
    "f\"Number of documents: {len(self.corpus)}\\n\\\n",
    "Number of train queries: {len(self.train_queries)}\\n\\\n",
    "Number of train relevance judgments: {len(self.train_qrels)}\\n\\\n",
    "Number of dev queries: {len(self.dev_queries)}\\n\\\n",
    "Number of dev relevance judgments: {len(self.dev_qrels)}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da3d744-fabe-4580-950e-10430d223f5a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of documents: 4\n",
      "Number of train queries: 2\n",
      "Number of train relevance judgments: 2\n",
      "Number of dev queries: 2\n",
      "Number of dev relevance judgments: 2\n"
     ]
    }
   ],
   "source": [
    "data.summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a672c192-75f7-4504-915f-bdb2dadd92ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide \n",
    "test_stdout(lambda: data.summary, \"Number of documents: 4\\n\\\n",
    "Number of train queries: 2\\n\\\n",
    "Number of train relevance judgments: 2\\n\\\n",
    "Number of dev queries: 2\\n\\\n",
    "Number of dev relevance judgments: 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d1b9a29-8087-43ae-bb3c-7e75bbdc6948",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def get_corpus(self: PassageData) -> pd.DataFrame:\n",
    "    return pd.DataFrame([{\"doc_id\": k, \"text\": v} for k, v in self.corpus.items()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce841a5f-0297-41b2-bc87-1d01a9f14e46",
   "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>doc_id</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>sentence 0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>sentence 1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>sentence 2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>sentence 3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  doc_id        text\n",
       "0      0  sentence 0\n",
       "1      1  sentence 1\n",
       "2      2  sentence 2\n",
       "3      3  sentence 3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data.get_corpus()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cce9e331-2aa3-41b4-a534-f65347ccb4b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "pd.testing.assert_frame_equal(\n",
    "    passage_data.get_corpus(), \n",
    "    pd.DataFrame(data={\"doc_id\": [str(x) for x in range(4)], \"text\":[f\"sentence {x}\" for x in range(4)]})\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d947777-7a0e-4867-8de7-a8d1e81a13fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def get_queries(\n",
    "    self: PassageData, \n",
    "    type: str  # Either 'train' or 'dev'.\n",
    ") -> pd.DataFrame:  # DataFrame conaining 'query_id' and 'query'.\n",
    "    \"Get query data.\"\n",
    "    assert type in ['train', 'dev'], \"'type' should be either 'train' or 'dev'.\"\n",
    "    if type == 'train':\n",
    "        data = self.train_queries\n",
    "    elif type == 'dev':\n",
    "        data = self.dev_queries\n",
    "    return pd.DataFrame([{\"query_id\": k, \"query\": v} for k, v in data.items()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57794b8c-4dde-47d3-8c23-d0880eb3fea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(passage_data.get_queries, contains=\"get_queries() missing 1 required positional argument: 'type'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75b9907b-86c4-4801-83de-a827205dd142",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(passage_data.get_queries, kwargs=dict(type=\"wrong_type\"), contains=\"'type' should be either 'train' or 'dev'.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0babab70-7106-4566-9fab-c69a152063e8",
   "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>query_id</th>\n",
       "      <th>query</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10</td>\n",
       "      <td>train query 10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>11</td>\n",
       "      <td>train query 11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  query_id           query\n",
       "0       10  train query 10\n",
       "1       11  train query 11"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data.get_queries(type=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b3a04f-8be3-48e2-9ebf-df81055a0184",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "pd.testing.assert_frame_equal(\n",
    "    passage_data.get_queries(type=\"train\"), \n",
    "    pd.DataFrame(data={\"query_id\": [str(x + 10) for x in range(2)], \"query\":[f\"train query {x + 10}\" for x in range(2)]})\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85560ac-ff80-4127-81b1-f3cf35a8f465",
   "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>query_id</th>\n",
       "      <th>query</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>20</td>\n",
       "      <td>train query 20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>21</td>\n",
       "      <td>train query 21</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  query_id           query\n",
       "0       20  train query 20\n",
       "1       21  train query 21"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data.get_queries(type=\"dev\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31009cb6-121e-4525-9707-d6cd97dd2bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "pd.testing.assert_frame_equal(\n",
    "    passage_data.get_queries(type=\"dev\"), \n",
    "    pd.DataFrame(data={\"query_id\": [str(x + 20) for x in range(2)], \"query\":[f\"train query {x + 20}\" for x in range(2)]})\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c84b47d0-f43e-49e5-b819-450e3dcc9010",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def get_labels(\n",
    "    self: PassageData, \n",
    "    type: str  # Either 'train' or 'dev'.\n",
    ") -> Dict:  # pyvespa-formatted labeled data \n",
    "    \"Get labeled data\"\n",
    "    assert type in ['train', 'dev'], \"'type' should be either 'train' or 'dev'.\"\n",
    "    if type == 'train':\n",
    "        queries = self.train_queries\n",
    "        qrels = self.train_qrels\n",
    "    elif type == 'dev':\n",
    "        queries = self.dev_queries\n",
    "        qrels = self.dev_qrels        \n",
    "    return [\n",
    "        {\n",
    "            \"query_id\": query_id, \n",
    "            \"query\": query, \n",
    "            \"relevant_docs\": [{\"id\": doc_id, \"score\": score} for doc_id, score in qrels[query_id].items()]\n",
    "        } \n",
    "        for query_id, query in queries.items()\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c0d300-b532-4dd2-ab15-11c72e5939fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(passage_data.get_labels, contains=\"get_labels() missing 1 required positional argument: 'type'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e12aa382-3f3a-493c-ab90-55f23739683b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(passage_data.get_labels, contains=\"'type' should be either 'train' or 'dev'.\", kwargs=dict(type=\"wrong_type\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4862db9a-36e6-490c-9476-82cb21111f6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'query_id': '10',\n",
       "  'query': 'train query 10',\n",
       "  'relevant_docs': [{'id': '0', 'score': 1}]},\n",
       " {'query_id': '11',\n",
       "  'query': 'train query 11',\n",
       "  'relevant_docs': [{'id': '2', 'score': 1}]}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data.get_labels(type=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bae2e33-d046-45b2-bfc0-62b68b2105db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'query_id': '20',\n",
       "  'query': 'train query 20',\n",
       "  'relevant_docs': [{'id': '1', 'score': 1}]},\n",
       " {'query_id': '21',\n",
       "  'query': 'train query 21',\n",
       "  'relevant_docs': [{'id': '3', 'score': 1}]}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "passage_data.get_labels(type=\"dev\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b904d173-6dde-416a-8012-eac329dd42af",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    passage_data.get_labels(type=\"dev\"), [\n",
    "        {'query_id': '20',\n",
    "         'query': 'train query 20',\n",
    "         'relevant_docs': [{'id': '1', 'score': 1}]\n",
    "        },\n",
    "        {'query_id': '21',\n",
    "         'query': 'train query 21',\n",
    "         'relevant_docs': [{'id': '3', 'score': 1}]\n",
    "        }\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e76526-5b8a-4142-afa9-e92537f9f14f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def sample_data(\n",
    "    n_relevant: int,  # The number of relevant documents to sample.\n",
    "    n_irrelevant: int,  # The number of non-judged documents to sample. \n",
    ") -> PassageData:  \n",
    "\n",
    "    \"\"\"\n",
    "    Sample data from the passage ranking dataset.\n",
    "    \n",
    "    The final sample contains `n_relevant` train relevant documents, \n",
    "    `n_relevant` dev relevant documents and `n_irrelevant` random documents\n",
    "    sampled from the entire corpus. \n",
    "    \n",
    "    All the relevant sampled documents, both from train and dev sets, \n",
    "    are guaranteed to be on the `corpus_sample`, which will contain \n",
    "    2 * `n_relevant` + `n_irrelevant` documents.\n",
    "    \"\"\"\n",
    "    passage_corpus = ir_datasets.load(\"msmarco-passage\")\n",
    "    passage_train = ir_datasets.load(\"msmarco-passage/train/judged\")    \n",
    "    passage_dev = ir_datasets.load(\"msmarco-passage/dev/judged\")\n",
    "    \n",
    "    train_qrels_dict = passage_train.qrels_dict()\n",
    "    train_qrels_sample = sample_dict_items(d=train_qrels_dict, n=n_relevant)\n",
    "\n",
    "    dev_qrels_dict = passage_dev.qrels_dict() \n",
    "    dev_qrels_sample = sample_dict_items(d=dev_qrels_dict, n=n_relevant)\n",
    "    \n",
    "    train_queries = {k:v for k,v in passage_train.queries_iter()}\n",
    "    train_queries_sample = {k: train_queries[k] for k in train_qrels_sample.keys()}\n",
    "    dev_queries = {k:v for k,v in passage_dev.queries_iter()}\n",
    "    dev_queries_sample = {k: dev_queries[k] for k in dev_qrels_sample.keys()}\n",
    "\n",
    "    train_positive_doc_id_samples = [\n",
    "        doc_id[0]\n",
    "        for doc_id in [list(docs.keys()) for docs in train_qrels_sample.values()]\n",
    "    ]\n",
    "    dev_positive_doc_id_samples = [\n",
    "        doc_id[0]\n",
    "        for doc_id in [list(docs.keys()) for docs in dev_qrels_sample.values()]\n",
    "    ]\n",
    "\n",
    "    docs_id = [x[0] for x in passage_corpus.docs_iter()]    \n",
    "    docs_store = passage_corpus.docs_store()\n",
    "    \n",
    "    negative_doc_id_samples = random.sample(\n",
    "        k=n_irrelevant, population=sorted(docs_id)\n",
    "    )\n",
    "    doc_id_samples = list(\n",
    "        set(\n",
    "            train_positive_doc_id_samples\n",
    "            + dev_positive_doc_id_samples\n",
    "            + negative_doc_id_samples\n",
    "        )\n",
    "    )\n",
    "    corpus_sample = {k: docs_store.get(k)[1] for k in doc_id_samples}\n",
    "\n",
    "    return PassageData(\n",
    "        corpus=corpus_sample,\n",
    "        train_qrels=train_qrels_sample,\n",
    "        train_queries=train_queries_sample,\n",
    "        dev_qrels=dev_qrels_sample,\n",
    "        dev_queries=dev_queries_sample,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c159f1e1-8a55-415b-b050-efea499b9e42",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b02820-a7dd-4c8a-aaef-9370244ea3e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|eval:false\n",
    "#|notest\n",
    "sample = sample_data(n_relevant=1, n_irrelevant=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993b9804-e6df-4259-ba49-a395d6c88d2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "sample = PassageData.load(file_path=\"resources/passage/passage_sample.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39888c3-c454-4a2c-98aa-986728590128",
   "metadata": {},
   "source": [
    "The sampled corpus is a dict containing document id as key and the passage text as value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa955262-f5b7-4768-9336-9ede436a6505",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'890370': 'the map of europe gives you a clear view of the political boundaries that segregate the countries in the continent including germany uk france spain italy greece romania ukraine hungary austria sweden finland norway czech republic belgium luxembourg switzerland croatia and albaniahe map of europe gives you a clear view of the political boundaries that segregate the countries in the continent including germany uk france spain italy greece romania ukraine hungary austria sweden finland norway czech republic belgium luxembourg switzerland croatia and albania',\n",
       " '5060205': 'Setting custom HTTP headers with cURL can be done by using the CURLOPT_HTTPHEADER option, which can be set with the curl_setopt function. To add headers to your HTTP request you need to put them into a PHP Array, which you can then pass to the cul_setopt function, like demonstrated in the below example.',\n",
       " '6096573': \"The sugar in RNA is ribose, whereas the sugar in DNA is deoxyribose. The only difference between the two is that in deoxyribose, there is an oxygen missing from the 2' carbon …(there is a H there instead of an OH). This makes DNA more stable/less reactive than RNA. 1 person found this useful.\",\n",
       " '3092885': 'All three C-Ph bonds are typical of sp 3 - sp 2 carbon-carbon bonds with lengths of approximately 1.47 A, å while The-C o bond length is approximately.1 42. A å the presence of three adjacent phenyl groups confers special properties manifested in the reactivity of. the alcohol',\n",
       " '7275560': 'shortest phase of mitosis Anaphase is the shortest phase of mitosis. During anaphase the arranged chromosomes at the metaphase plate are migrate towards their respective poles. Before this migration started, chromosomes are divided into sister chromatids, by the separation of joined centromere of two sister chromatids of a chromosomes.'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample.corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430fb02a-d0be-4ef9-a11a-9affc8d267af",
   "metadata": {},
   "source": [
    "The size of the sampled document corpus is equal to 2 * `n_relevant` + `n_irrelevant`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6617b585-587e-4e8d-af8e-4345349142a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(sample.corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89963a8c-0bf4-4bae-98dc-635141fc51b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(len(sample.corpus), 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f3f64c-765a-431d-98af-7277fd9ee91c",
   "metadata": {},
   "source": [
    "Sampled queries are dict containing query id as key and query text as value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91122e87-b0c0-465c-96b4-3d95eb53feaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'899723': 'what sugar is found in rna'}\n",
      "{'994205': 'which is the shortest stage in duration'}\n"
     ]
    }
   ],
   "source": [
    "print(sample.train_queries)\n",
    "print(sample.dev_queries)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cf9f91b-8934-4b2c-b999-3aa9858c9177",
   "metadata": {},
   "source": [
    "Sampled qrels contains one relevant document for each query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b79271fe-96d3-433c-81da-4a590f1476cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'899723': {'6096573': 1}}\n",
      "{'994205': {'7275560': 1}}\n"
     ]
    }
   ],
   "source": [
    "print(sample.train_qrels)\n",
    "print(sample.dev_qrels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb5fe9a3-9192-4294-b914-e00397a10ccf",
   "metadata": {},
   "source": [
    "The following relevant documents are guaranteed to be included in the `corpus_sample`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90916e1b-2ac7-4c27-9409-51d1dea33976",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['6096573', '7275560']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#|echo:false\n",
    "relevant_doc_ids= [\n",
    "    x\n",
    "    for query_id in sample.train_qrels \n",
    "    for x in sample.train_qrels[query_id]\n",
    "] + [\n",
    "    x \n",
    "    for query_id in sample.dev_qrels \n",
    "    for x in sample.dev_qrels[query_id]\n",
    "]\n",
    "relevant_doc_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7b24157-4e32-4979-838b-a049ed105633",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "sample_corpus_ids = [x for x in sample.corpus]\n",
    "assert all([x in sample_corpus_ids for x in relevant_doc_ids])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53157953-e4d8-4cf5-a373-7904cd2d8698",
   "metadata": {},
   "source": [
    "## Basic search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bd01676-c87d-4048-a59c-47adb452c303",
   "metadata": {},
   "source": [
    "Code related to a basic search search engine for passage ranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32928903-3eaa-42b1-bd6b-dc6b070c9e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_basic_search_package(\n",
    "    name: str=\"PassageRanking\"  # Name of the application\n",
    ") -> ApplicationPackage: # pyvespa [ApplicationPackage](https://pyvespa.readthedocs.io/en/latest/reference-api.html#applicationpackage) instance.\n",
    "    \"\"\"\n",
    "    Create a basic Vespa application package for passage ranking.\n",
    "    \n",
    "    *Vespa fields*:\n",
    "        \n",
    "    The application contain two string fields: `doc_id` and `text`.\n",
    "    \n",
    "    *Vespa rank functions*:\n",
    "    \n",
    "    The application contain two rank profiles: \n",
    "    [bm25](https://docs.vespa.ai/en/reference/bm25.html) and\n",
    "    [nativeRank](https://docs.vespa.ai/en/reference/nativerank.html).\n",
    "    \n",
    "    \"\"\"\n",
    "    app_package = ApplicationPackage(name=name)\n",
    "    app_package.schema.add_fields(\n",
    "        Field(name=\"doc_id\", type=\"string\", indexing=[\"attribute\", \"summary\"]),\n",
    "        Field(name=\"text\", type=\"string\", indexing=[\"index\", \"summary\"], index=\"enable-bm25\"),\n",
    "    )\n",
    "    app_package.schema.add_field_set(\n",
    "        FieldSet(name=\"default\", fields=[\"text\"])\n",
    "    )\n",
    "    app_package.schema.add_rank_profile(\n",
    "        RankProfile(\n",
    "            name=\"bm25\",\n",
    "            first_phase=\"bm25(text)\",\n",
    "            summary_features=[\"bm25(text)\"]\n",
    "        )\n",
    "    )\n",
    "    app_package.schema.add_rank_profile(\n",
    "        RankProfile(\n",
    "            name=\"native_rank\", \n",
    "            first_phase=\"nativeRank(text)\"\n",
    "        )\n",
    "    )\n",
    "    app_package.query_profile.add_fields(\n",
    "        QueryField(name=\"maxHits\", value=10000)\n",
    "    )\n",
    "    return app_package"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd4e189-f2f4-4244-b321-de2529c43105",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19e9cde8-69e0-4975-802e-4a70a8461ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "app_package = create_basic_search_package(name=\"PassageModuleApp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979bfba8-211a-4d8a-a91a-eed88822fa1b",
   "metadata": {},
   "source": [
    "Check how the [Vespa schema definition](https://docs.vespa.ai/en/schemas.html) for this application looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a7825b-e05c-4d1a-b31b-0251f2815050",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "schema PassageModuleApp {\n",
      "    document PassageModuleApp {\n",
      "        field doc_id type string {\n",
      "            indexing: attribute | summary\n",
      "        }\n",
      "        field text type string {\n",
      "            indexing: index | summary\n",
      "            index: enable-bm25\n",
      "        }\n",
      "    }\n",
      "    fieldset default {\n",
      "        fields: text\n",
      "    }\n",
      "    rank-profile bm25 {\n",
      "        first-phase {\n",
      "            expression: bm25(text)\n",
      "        }\n",
      "        summary-features {\n",
      "            bm25(text)\n",
      "        }\n",
      "    }\n",
      "    rank-profile native_rank {\n",
      "        first-phase {\n",
      "            expression: nativeRank(text)\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(app_package.schema.schema_to_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc6a156c-3a09-4c1d-a238-90b95bb3d909",
   "metadata": {},
   "source": [
    "## Evaluate query models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84f9913d-60a8-4e27-bcfc-d33ee431f6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def evaluate_query_models(\n",
    "    app_package: ApplicationPackage, \n",
    "    query_models: List[QueryModel],\n",
    "    metrics: List[EvalMetric],\n",
    "    corpus_size: List[int],\n",
    "    output_file_path: str,\n",
    "    dev_query_percentage: float = 55578/8841823,\n",
    "    verbose: bool = True,\n",
    "    **kwargs\n",
    "):\n",
    "\n",
    "    print(\"*****\")\n",
    "    print(\"Deploy Vespa application:\")\n",
    "    print(\"*****\")\n",
    "    vespa_docker = VespaDocker(port=8183, cfgsrv_port=19173)\n",
    "    app = vespa_docker.deploy(application_package=app_package)\n",
    "    dfs = []\n",
    "    for idx, n in enumerate(corpus_size):\n",
    "        print(\"*****\")\n",
    "        print(f\"Corpus size:{n}\")\n",
    "        print(\"*****\")        \n",
    "        if verbose:\n",
    "            print(\"*****\")\n",
    "            print(\"Generate sample data:\")\n",
    "            print(\"*****\")\n",
    "        data = sample_data(\n",
    "            n_relevant=int(n * dev_query_percentage), \n",
    "            n_irrelevant=int(n*(1-2*dev_query_percentage))\n",
    "        )    \n",
    "        if verbose:\n",
    "            data.summary\n",
    "            print(\"*****\")\n",
    "            print(\"Feed sample data to Vespa app:\")\n",
    "            print(\"*****\")\n",
    "        responses = app.feed_df(df=data.get_corpus(), include_id=True, id_field=\"doc_id\", batch_size=10000)\n",
    "        if verbose:\n",
    "            print(\"*****\")\n",
    "            print(\"Evaluate query models\")\n",
    "            print(\"*****\")\n",
    "        labeled_data = data.get_labels(type=\"dev\")\n",
    "        assert len(labeled_data) > 0, \"Need at least one query.\"\n",
    "        evaluation_per_query = evaluate(\n",
    "            app=app,\n",
    "            labeled_data=labeled_data, \n",
    "            eval_metrics=metrics, \n",
    "            query_model=query_models, \n",
    "            id_field=\"doc_id\",\n",
    "            per_query=True,\n",
    "            **kwargs\n",
    "        )    \n",
    "        estimates = compute_evaluation_estimates(\n",
    "            df = evaluation_per_query\n",
    "        )    \n",
    "        estimates = estimates.assign(corpus_size=n, number_queries=len(labeled_data))\n",
    "        if idx == 0:\n",
    "            estimates.to_csv(output_file_path, index=False, mode=\"w\")\n",
    "        else:\n",
    "            estimates.to_csv(output_file_path, index=False, mode=\"a\", header=False)\n",
    "        dfs.append(estimates)\n",
    "        if verbose:\n",
    "            print(\"*****\")\n",
    "            print(\"Delete all documents\")\n",
    "            print(\"*****\")\n",
    "        app.delete_all_docs(content_cluster_name=\"PassageRanking_content\", schema=\"PassageRanking\")\n",
    "    if verbose:\n",
    "        print(\"*****\")\n",
    "        print(\"Stop and remove Docker container\")\n",
    "        print(\"*****\")\n",
    "    vespa_docker.container.stop(timeout=600)\n",
    "    vespa_docker.container.remove()\n",
    "    estimates = pd.concat(dfs)\n",
    "    return estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8d42641-18ea-4524-bce5-36686c9c55d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from learntorank.evaluation import (\n",
    "    MatchRatio,\n",
    "    Recall, \n",
    "    ReciprocalRank, \n",
    "    NormalizedDiscountedCumulativeGain\n",
    ")\n",
    "from learntorank.query import QueryModel, OR, Ranking\n",
    "\n",
    "corpus_size = [100, 200]\n",
    "app_package = create_basic_search_package(name=\"PassageEvaluationApp\")\n",
    "query_models = [\n",
    "    QueryModel(\n",
    "        name=\"bm25\", \n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking(name=\"bm25\")\n",
    "    ),\n",
    "    QueryModel(\n",
    "        name=\"native_rank\", \n",
    "        match_phase=OR(), \n",
    "        ranking=Ranking(name=\"native_rank\")\n",
    "    )\n",
    "]\n",
    "metrics = [\n",
    "    MatchRatio(),\n",
    "    Recall(at=100), \n",
    "    ReciprocalRank(at=10), \n",
    "    NormalizedDiscountedCumulativeGain(at=10)\n",
    "]\n",
    "output_file_path = \"test.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67772339-21c9-4ba5-bd47-0e6a34566c28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "*****\n",
      "Deploy Vespa application:\n",
      "*****\n",
      "Waiting for configuration server, 0/300 seconds...\n",
      "Waiting for configuration server, 5/300 seconds...\n",
      "Waiting for configuration server, 10/300 seconds...\n",
      "Waiting for application status, 0/300 seconds...\n",
      "Waiting for application status, 5/300 seconds...\n",
      "Waiting for application status, 10/300 seconds...\n",
      "Waiting for application status, 15/300 seconds...\n",
      "Waiting for application status, 20/300 seconds...\n",
      "Waiting for application status, 25/300 seconds...\n",
      "Waiting for application status, 30/300 seconds...\n",
      "Waiting for application status, 35/300 seconds...\n",
      "Waiting for application status, 40/300 seconds...\n",
      "Waiting for application status, 45/300 seconds...\n",
      "Waiting for application status, 50/300 seconds...\n",
      "Waiting for application status, 55/300 seconds...\n",
      "Waiting for application status, 60/300 seconds...\n",
      "Waiting for application status, 65/300 seconds...\n",
      "Waiting for application status, 70/300 seconds...\n",
      "Waiting for application status, 75/300 seconds...\n",
      "Waiting for application status, 80/300 seconds...\n"
     ]
    }
   ],
   "source": [
    "estimates = evaluate_query_models(\n",
    "    app_package=app_package,\n",
    "    query_models=query_models,\n",
    "    metrics=metrics,\n",
    "    corpus_size=corpus_size,\n",
    "    dev_query_percentage=0.5,\n",
    "    output_file_path=output_file_path, \n",
    "    verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b19d53ca-04b9-4293-aca7-63dc78161013",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "os.remove(\"test.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d77781-8f0d-46d0-92f6-22ef8b25719d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learntorank",
   "language": "python",
   "name": "learntorank"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}