002_module_evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a19ce1-bb8a-40c1-8ef3-ff78f60b71d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp evaluation"
   ]
  },
  {
   "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": [
    "# evaluation\n",
    "> Reference API related to evaluation function and metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "806e548c-5588-4a0a-b4d5-e744bbe1fc1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "import math\n",
    "from typing import Dict, List, Union\n",
    "from fastcore.utils import patch, patch_to\n",
    "from pandas import DataFrame\n",
    "from vespa.io import VespaQueryResponse\n",
    "from vespa.application import Vespa\n",
    "from learntorank.query import QueryModel, send_query, send_query_batch, _parse_labeled_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89292957-9a41-44bc-91bf-130e1a733340",
   "metadata": {},
   "source": [
    "## Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e4a35ee-d5ce-4676-9f28-7c63870e6a97",
   "metadata": {},
   "source": [
    "Abstract and concrete classes related to evaluation metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38100b48-65d7-4ed0-a204-9f3a53604fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class EvalMetric(object):\n",
    "    def __init__(self) -> None:\n",
    "        \"Abstract class for evaluation metric.\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46b14706-1799-49a9-ae4f-e66f9c8708f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: EvalMetric,\n",
    "    query_results,  # Raw query results returned by Vespa.\n",
    "    relevant_docs,  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field,  # The Vespa field representing the document id.\n",
    "    default_score,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Metric values.\n",
    "    \"Abstract method to be implemented by metrics inheriting from `EvalMetric` to evaluate query results.\"\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374401b0-0a00-40b1-96e6-04c1f9bd5d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class MatchRatio(EvalMetric):\n",
    "    def __init__(self) -> None:\n",
    "        \"Computes the ratio of documents retrieved by the match phase.\"\n",
    "        super().__init__()\n",
    "        self.name = \"match_ratio\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97032c92-d8c3-4e67-bbb0-629b5667f358",
   "metadata": {},
   "source": [
    "Instantiate the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c766081-2330-4743-ac1f-bc916bf81eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = MatchRatio()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcaeb7ae-aa50-4101-a1ba-438da46b3e5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: MatchRatio,\n",
    "    query_results: VespaQueryResponse,  # Raw query results returned by Vespa.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Returns the match ratio. In addition, if `detailed_metrics=False`, returns the number of retrieved docs `_retrieved_docs` and the number of docs available in the corpus `_docs_available`.\n",
    "    \"Evaluate query results according to match ratio metric.\"\n",
    "    \n",
    "    retrieved_docs = query_results.number_documents_retrieved\n",
    "    docs_available = query_results.number_documents_indexed\n",
    "    value = 0\n",
    "    if docs_available > 0:\n",
    "        value = retrieved_docs / docs_available\n",
    "    metrics = {\n",
    "        str(self.name): value,\n",
    "    }\n",
    "    if detailed_metrics:\n",
    "        metrics.update(\n",
    "            {\n",
    "                str(self.name) + \"_retrieved_docs\": retrieved_docs,\n",
    "                str(self.name) + \"_docs_available\": docs_available,\n",
    "            }\n",
    "        )\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b169be40-d518-4e92-a748-749f83e7a374",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query_results = VespaQueryResponse(\n",
    "    {\"root\": {\"fields\": {\"totalCount\": 1083},\n",
    "              \"coverage\": {\"documents\": 62529}}\n",
    "    }, \n",
    "    status_code=None, \n",
    "    url=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd62e5c-a644-4021-a444-67ef9c372e52",
   "metadata": {},
   "source": [
    "Compute match ratio:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88e2719e-108a-490f-84c8-7ec777b23468",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'match_ratio': 0.01731996353691887}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results, \n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b749f6-bc55-4d4b-b7cd-22a33b45dcf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation,\n",
    "    {\n",
    "        \"match_ratio\": 1083 / 62529,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6806d5d2-f84e-4186-b733-2322df4aceda",
   "metadata": {},
   "source": [
    "Return detailed metrics, in addition to match ratio:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec4a2a9-3606-491a-b3bc-5a681b125624",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'match_ratio': 0.01731996353691887,\n",
       " 'match_ratio_retrieved_docs': 1083,\n",
       " 'match_ratio_docs_available': 62529}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23c12158-ec99-4e57-9468-cf2dacad26ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"match_ratio_retrieved_docs\": 1083,\n",
    "        \"match_ratio_docs_available\": 62529,\n",
    "        \"match_ratio\": 1083 / 62529,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5c667a-00ef-4cef-bf36-df155766393a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without 'totalCount'\n",
    "query_results = VespaQueryResponse(\n",
    "    {\n",
    "        \"root\": {\n",
    "            \"coverage\": {\n",
    "                \"documents\": 62529,\n",
    "            },\n",
    "        }\n",
    "    }, \n",
    "    status_code=None, \n",
    "    url=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61e890b5-28ef-4207-8460-4e08f0d16ca3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without 'totalCount'\n",
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"match_ratio\": 0 / 62529,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aad19c90-dce1-4cdd-8e3e-909e6abb115f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without 'totalCount'\n",
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"match_ratio_retrieved_docs\": 0,\n",
    "        \"match_ratio_docs_available\": 62529,\n",
    "        \"match_ratio\": 0 / 62529,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f264a6-6036-49ca-89db-307cfe0abb3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without \"coverage\": {\"documents\": 62529}\n",
    "query_results=VespaQueryResponse({\n",
    "                \"root\": {\n",
    "                    \"id\": \"toplevel\",\n",
    "                    \"relevance\": 1.0,\n",
    "                    \"fields\": {\"totalCount\": 1083},\n",
    "                    \"coverage\": {\n",
    "                        \"coverage\": 100,\n",
    "                        \"full\": True,\n",
    "                        \"nodes\": 2,\n",
    "                        \"results\": 1,\n",
    "                        \"resultsFull\": 1,\n",
    "                    },\n",
    "                }\n",
    "            }, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29495b98-6f8e-46bd-9adb-1b46b8cf0d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without \"coverage\": {\"documents\": 62529}\n",
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"match_ratio\": 0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e535b59a-001e-4239-8fca-6bf201d27013",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# case without \"coverage\": {\"documents\": 62529}\n",
    "evaluation = metric.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=None,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"match_ratio_retrieved_docs\": 1083,\n",
    "        \"match_ratio_docs_available\": 0,\n",
    "        \"match_ratio\": 0,\n",
    "    },\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388bf07a-028c-4e51-ad09-5694939edbda",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class TimeQuery(EvalMetric):\n",
    "    def __init__(self) -> None:\n",
    "        \"Compute the time it takes for Vespa to execute the query..\"\n",
    "        super().__init__()\n",
    "        self.name = \"search_time\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4960958-942b-4215-8e17-b6411cf4451a",
   "metadata": {},
   "source": [
    "Instantiate the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb52b670-e768-457e-bf87-356445cad1c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_metric = TimeQuery()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f52cf202-89e9-4fe4-8800-0e9bcbf715f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export \n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: TimeQuery,\n",
    "    query_results: VespaQueryResponse,  # Raw query results returned by Vespa.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Returns the match ratio. In addition, if `detailed_metrics=False`, returns the number of retrieved docs `_retrieved_docs` and the number of docs available in the corpus `_docs_available`.\n",
    "    \"Evaluate query results according to query time metric.\"\n",
    "    try:\n",
    "        timing = query_results.json['timing']    \n",
    "    except KeyError:\n",
    "        timing = {\n",
    "            'querytime': 0,\n",
    "            'summaryfetchtime': 0,\n",
    "            'searchtime': 0\n",
    "        }\n",
    "    metrics = {\n",
    "        str(self.name): timing[\"searchtime\"],\n",
    "    }\n",
    "    if detailed_metrics:\n",
    "        metrics.update(\n",
    "            {\n",
    "                str(self.name) + \"_query_time\": timing[\"querytime\"],\n",
    "                str(self.name) + \"_summary_fetch_time\": timing[\"summaryfetchtime\"],\n",
    "            }\n",
    "        )\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "634cf6fe-e5f9-4fef-bc21-f00806cadc09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query_results = VespaQueryResponse({\n",
    "    \"timing\": {'querytime': 0.01,\n",
    "               'summaryfetchtime': 0.002,\n",
    "               'searchtime': 0.013\n",
    "              }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6abb22ce-0cde-4999-bf3e-1f390139035d",
   "metadata": {},
   "source": [
    "Compute the query time a client would observe (except network latency). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa92bcbf-2582-44ec-91ca-13a1299f35af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'search_time': 0.013}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time_metric.evaluate_query(\n",
    "    query_results=query_results, \n",
    "    relevant_docs=None, \n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7955c080-3b09-4494-9565-1748907683a8",
   "metadata": {},
   "source": [
    "Include detailed metrics. In addition to the `search_time` above, it returns the time to execute the first protocol phase/matching phase (`search_time_query_time`) and the time to execute the summary fill protocol phase for the globally ordered top-k hits (`search_time_summary_fetch_time`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4172f367-aabb-464f-acca-f65a6bd2729b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'search_time': 0.013,\n",
       " 'search_time_query_time': 0.01,\n",
       " 'search_time_summary_fetch_time': 0.002}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time_metric.evaluate_query(\n",
    "    query_results=query_results, \n",
    "    relevant_docs=None, \n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecebd54-96df-453a-9647-0d7810efe74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# in case timing is not available\n",
    "query_results = VespaQueryResponse({\n",
    "    \"something_else\": {\n",
    "        'querytime': 0.01,\n",
    "        'summaryfetchtime': 0.002,\n",
    "        'searchtime': 0.013\n",
    "    }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4934b807-91cf-43f5-9f04-fe9e73cb7215",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    time_metric.evaluate_query(\n",
    "        query_results=query_results, \n",
    "        relevant_docs=None, \n",
    "        id_field=\"vespa_id_field\",\n",
    "        default_score=0\n",
    "    ), \n",
    "    {'search_time': 0}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c10157-46ad-4ed6-b25b-6810e86c07b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    time_metric.evaluate_query(\n",
    "        query_results=query_results, \n",
    "        relevant_docs=None, \n",
    "        id_field=\"vespa_id_field\",\n",
    "        default_score=0,\n",
    "        detailed_metrics=True\n",
    "    ), \n",
    "    {\n",
    "        'search_time': 0,\n",
    "        'search_time_query_time': 0,\n",
    "        'search_time_summary_fetch_time': 0\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e502552-ba9b-4ab0-8125-cea8bc25310d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Recall(EvalMetric):\n",
    "    def __init__(\n",
    "        self, \n",
    "        at: int  # Maximum position on the resulting list to look for relevant docs.\n",
    "    ) -> None:\n",
    "        \"Compute the recall at position `at`.\"\n",
    "        super().__init__()\n",
    "        self.name = \"recall_\" + str(at)\n",
    "        self.at = at"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2088a8cb-dec5-47c0-a8bb-2e70263659ff",
   "metadata": {},
   "source": [
    "Instantiate the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5980f0f-5cd5-45bb-9472-dd95de9ba205",
   "metadata": {},
   "outputs": [],
   "source": [
    "recall_1 = Recall(at=1)\n",
    "recall_2 = Recall(at=2)\n",
    "recall_3 = Recall(at=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50fbf673-f3cd-4590-b89a-a1e0ca3c5007",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: Recall,\n",
    "    query_results: VespaQueryResponse,  # Raw query results returned by Vespa.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Returns the recall value.\n",
    "    \"\"\"\n",
    "    Evaluate query results according to recall metric.\n",
    "\n",
    "    There is an assumption that only documents with score > 0 are relevant. Recall is equal to zero in case no\n",
    "    relevant documents with score > 0 is provided.\n",
    "    \"\"\"\n",
    "\n",
    "    relevant_ids = {str(doc[\"id\"]) for doc in relevant_docs if doc.get(\"score\", default_score) > 0}\n",
    "    try:\n",
    "        retrieved_ids = {\n",
    "            str(hit[\"fields\"][id_field]) for hit in query_results.hits[: self.at]\n",
    "        }\n",
    "    except KeyError:\n",
    "        retrieved_ids = set()\n",
    "\n",
    "    return {str(self.name): len(relevant_ids & retrieved_ids) / len(relevant_ids) if len(relevant_ids) > 0 else 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28187738-c974-417a-b86d-27f08d56fba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "query_results = VespaQueryResponse({\n",
    "    \"root\": {\n",
    "        \"children\": [\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"ghi\",\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"def\",\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b13f00-439c-4a82-8a3f-b7608ca74ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "relevant_docs = [{\"id\": \"def\", \"score\": 1}, {\"id\": \"abc\", \"score\": 1}]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c8bf6d0-ad04-4462-ab2e-f25573eca8b9",
   "metadata": {},
   "source": [
    "Compute recall:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e805dba-f630-46d6-8fd7-a0e629557577",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'recall_2': 0.5}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = recall_2.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa4793b-d34b-45d3-b408-adc31fd42a09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_2\": 0.5,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f20914-078c-498d-98fa-c98b469b50cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# same data as above but with recall_1\n",
    "evaluation = recall_1.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_1\": 0.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c79ac0-dc99-4732-aa7b-ce72eaf04e75",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# There is an additional third hit when compared to 'query_results'\n",
    "query_results2 = VespaQueryResponse({\n",
    "    \"root\": {\n",
    "        \"children\": [\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"ghi\",\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"def\",\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"abc\",\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0469e96-e0a6-4eb4-95f7-418fb6588b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# different relevant scores, score != 1\n",
    "relevant_docs2 = [{\"id\": \"ghi\", \"score\": 1}, {\"id\": \"abc\", \"score\": 2}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05fdc954-4ada-4f7e-a6fb-c794f2d6146f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = recall_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs2,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_3\": 1,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3805f0b3-8a3a-4441-b165-590874832cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# id field is an integer\n",
    "query_results_int_id = VespaQueryResponse({\n",
    "    \"root\": {\n",
    "        \"children\": [\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": 1,\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": 2,\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": 3,\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5147947-b8e4-404c-8e6d-c891fed3af60",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "relevant_docs_int_id = [{\"id\": 1, \"score\": 1}, {\"id\": 3, \"score\": 2}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61dcc2b1-dddd-4a89-a425-5c2f27a9b772",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = recall_3.evaluate_query(\n",
    "    query_results=query_results_int_id,\n",
    "    relevant_docs=relevant_docs_int_id,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_3\": 1,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826428cb-5ede-4ed5-8758-61e2fd7a0e4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# relevant docs containing score = 0\n",
    "relevant_docs_with_zero_score = [{\"id\": \"ghi\", \"score\": 0}, {\"id\": \"abc\", \"score\": 2}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6098ff62-a759-4119-9091-be2512527461",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# test recall metric in the presence of score = 0 in the relevant docs\n",
    "evaluation = recall_1.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_1\": 0,\n",
    "    },\n",
    ")\n",
    "\n",
    "evaluation = recall_2.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_2\": 0,\n",
    "    },\n",
    ")\n",
    "\n",
    "evaluation = recall_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_3\": 1,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58787ca1-b5cb-4d50-a2e5-579eb6529208",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# test recall metric with relevant docs containing only score = 0\n",
    "evaluation = recall_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=[{\"id\": \"ghi\", \"score\": 0}, {\"id\": \"abc\", \"score\": 0}],\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_3\": 0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e38144-3e49-482a-9ea9-b39a9c4ae571",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "# missing 'fields'\n",
    "query_results_empty_field = VespaQueryResponse({\n",
    "    \"root\": {\n",
    "        \"children\": [\n",
    "            {\n",
    "                \"id\": \"ghi\"\n",
    "            },\n",
    "            {\n",
    "                \"fields\": {\n",
    "                    \"vespa_id_field\": \"def\",\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "}, status_code=None, url=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b065e84-9e59-4af7-934b-a2d79d6adaa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "relevant_docs = [{\"id\": \"def\", \"score\": 1}, {\"id\": \"abc\", \"score\": 1}]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c13425da-8397-4b1f-b93d-47cda6bcada6",
   "metadata": {},
   "source": [
    "Compute recall:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98cfac3-8c14-4516-94ee-6e4de2421e52",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = recall_2.evaluate_query(\n",
    "    query_results=query_results_empty_field,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"recall_2\": 0.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d58a53-c0d3-47df-8207-9e07328a515c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ReciprocalRank(EvalMetric):\n",
    "    def __init__(\n",
    "        self, \n",
    "        at: int  # Maximum position on the resulting list to look for relevant docs.\n",
    "    ):\n",
    "        \"Compute the reciprocal rank at position `at`\"\n",
    "        super().__init__()\n",
    "        self.name = \"reciprocal_rank_\" + str(at)\n",
    "        self.at = at"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932958f1-6cf2-430f-ab7a-b85a7fa0b01b",
   "metadata": {},
   "source": [
    "Instantiate the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b23272e5-5422-4ab8-b028-463a30edcd23",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr_1 = ReciprocalRank(at=1)\n",
    "rr_2 = ReciprocalRank(at=2)\n",
    "rr_3 = ReciprocalRank(at=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f7075e7-c17a-4c4d-b9dc-862219a6b3f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: ReciprocalRank,\n",
    "    query_results: VespaQueryResponse,  # Raw query results returned by Vespa.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Returns the reciprocal rank value.\n",
    "    \"\"\"\n",
    "    Evaluate query results according to reciprocal rank metric.\n",
    "\n",
    "    There is an assumption that only documents with score > 0 are relevant.\n",
    "    \"\"\"\n",
    "\n",
    "    relevant_ids = {str(doc[\"id\"]) for doc in relevant_docs if doc.get(\"score\", default_score) > 0}\n",
    "    rr = 0\n",
    "    hits = query_results.hits[: self.at]\n",
    "    for index, hit in enumerate(hits):\n",
    "        try:\n",
    "            if str(hit[\"fields\"][id_field]) in relevant_ids:\n",
    "                rr = 1 / (index + 1)\n",
    "                break\n",
    "        except KeyError:\n",
    "            rr = 0\n",
    "\n",
    "    return {str(self.name): rr}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3ca686-f093-4c4a-a462-e4948d8effca",
   "metadata": {},
   "source": [
    "Compute reciprocal rank:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa79e504-e158-44f0-8d23-a571c7dcc647",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'reciprocal_rank_2': 0.5}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = rr_2.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc74b315-d744-4c4e-8216-afbaeee1a377",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_2\": 0.5,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e4f0ff8-5bd0-455e-b481-f2520ffb7ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_1.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_1\": 0.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c1014a8-2478-4c84-be9e-5ef6d6ab8460",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs2,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_3\": 1.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b5630a-3824-4172-9eab-5b0396e49af9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_3.evaluate_query(\n",
    "    query_results=query_results_int_id,\n",
    "    relevant_docs=relevant_docs_int_id,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_3\": 1.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f09a49-7c5d-4efb-b760-e1f5493d03e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_1.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_1\": 0.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031818eb-ba1c-488e-a0cd-0508de478745",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "metric = ReciprocalRank(at=2)\n",
    "evaluation = rr_2.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_2\": 0.0,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea591cbd-efda-4254-b918-d79bc387376f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs_with_zero_score,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_3\": 1 / 3,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "266447f5-ec4a-4f34-99fc-1ca7c13fd7e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = rr_2.evaluate_query(\n",
    "    query_results=query_results_empty_field,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"reciprocal_rank_2\": 0.5,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9ecef1-372c-4034-96d6-7fc5bea574e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class NormalizedDiscountedCumulativeGain(EvalMetric):\n",
    "    def __init__(\n",
    "        self, \n",
    "        at: int  # Maximum position on the resulting list to look for relevant docs.\n",
    "    ):\n",
    "        \"Compute the normalized discounted cumulative gain at position `at`.\"\n",
    "        super().__init__()\n",
    "        self.name = \"ndcg_\" + str(at)\n",
    "        self.at = at\n",
    "\n",
    "    @staticmethod    \n",
    "    def _compute_dcg(scores: List[int]) -> float:\n",
    "        return sum([score / math.log2(idx + 2) for idx, score in enumerate(scores)])        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dae8ee5b-8e23-4c88-88de-e245ce165bd1",
   "metadata": {},
   "source": [
    "Instantiate the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a98e4d10-cebc-4cbd-8dcc-4cf1a05c670b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ndcg_1 = NormalizedDiscountedCumulativeGain(at=1)\n",
    "ndcg_2 = NormalizedDiscountedCumulativeGain(at=2)\n",
    "ndcg_3 = NormalizedDiscountedCumulativeGain(at=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89e2a87-e1a0-48d3-82a6-a95b5862c2d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@patch\n",
    "def evaluate_query(\n",
    "    self: NormalizedDiscountedCumulativeGain,\n",
    "    query_results: VespaQueryResponse,  # Raw query results returned by Vespa.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int,  # Score to assign to the additional documents that are not relevant. Default to 0.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    ") -> Dict:  # Returns the normalized discounted cumulative gain. In addition, if `detailed_metrics=False`, returns the ideal discounted cumulative gain `_ideal_dcg`, the discounted cumulative gain `_dcg`.\n",
    "    \"\"\"\n",
    "    Evaluate query results according to normalized discounted cumulative gain.\n",
    "\n",
    "    There is an assumption that documents returned by the query that are not included in the set of relevant\n",
    "    documents have score equal to zero. Similarly, if the query returns a number `N < at` documents, we will\n",
    "    assume that those `N - at` missing scores are equal to zero.\n",
    "    \"\"\"\n",
    "\n",
    "    at = self.at\n",
    "    relevant_scores = {str(doc[\"id\"]): doc[\"score\"] for doc in relevant_docs}\n",
    "    assert default_score == 0, \"NDCG default score should be zero.\"\n",
    "    search_scores = [default_score] * at\n",
    "    ideal_scores = [default_score] * at\n",
    "\n",
    "    hits = query_results.hits[:at]\n",
    "    for idx, hit in enumerate(hits):\n",
    "        try:\n",
    "            doc_id = str(hit[\"fields\"][id_field])\n",
    "            search_scores[idx] = relevant_scores.get(\n",
    "                doc_id, default_score\n",
    "            )\n",
    "        except KeyError:\n",
    "            search_scores[idx] = default_score\n",
    "\n",
    "    sorted_score_list = sorted(list(relevant_scores.values()), reverse=True)[:at]\n",
    "    for idx, score in enumerate(sorted_score_list):\n",
    "        ideal_scores[idx] = score\n",
    "\n",
    "    ideal_dcg = self._compute_dcg(ideal_scores)\n",
    "    dcg = self._compute_dcg(search_scores)\n",
    "\n",
    "    ndcg = 0\n",
    "    if ideal_dcg > 0:\n",
    "        ndcg = dcg / ideal_dcg\n",
    "\n",
    "    metrics = {\n",
    "        str(self.name): ndcg,\n",
    "    }\n",
    "    if detailed_metrics:\n",
    "        metrics.update(\n",
    "            {\n",
    "                str(self.name) + \"_ideal_dcg\": ideal_dcg,\n",
    "                str(self.name) + \"_dcg\": dcg,\n",
    "            }\n",
    "        )\n",
    "    return metrics\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e96f8a5-2724-454c-8237-db578836fd48",
   "metadata": {},
   "source": [
    "Compute NDCG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98837e94-4266-4fbd-ab52-2cc652804207",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ndcg_2': 0.38685280723454163}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metric = NormalizedDiscountedCumulativeGain(at=2)\n",
    "evaluation = ndcg_2.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37e889e2-2be7-4fcf-a636-7cdc773649b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "expected_dcg = 0 / math.log2(2) + 1 / math.log2(3)\n",
    "expected_ideal_dcg = 1 / math.log2(2) + 1 / math.log2(3)\n",
    "expected_ndcg = expected_dcg / expected_ideal_dcg\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_2\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184bab64-cc4a-4777-8aaf-d5759c569b01",
   "metadata": {},
   "source": [
    "Return detailed metrics, in addition to NDCG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b3e2563-9086-44ac-8ad9-0da6c84e3067",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ndcg_2': 0.38685280723454163,\n",
       " 'ndcg_2_ideal_dcg': 1.6309297535714575,\n",
       " 'ndcg_2_dcg': 0.6309297535714575}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = ndcg_2.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ae015c-752b-47d0-ab50-485995adc392",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_2_ideal_dcg\": expected_ideal_dcg,\n",
    "        \"ndcg_2_dcg\": expected_dcg,\n",
    "        \"ndcg_2\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c72c312-1f71-4347-ade2-a88cce314c48",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_1.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "expected_dcg = 0 / math.log2(2)\n",
    "expected_ideal_dcg = 1 / math.log2(2)\n",
    "expected_ndcg = 0\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_1\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72a161f7-1ed9-4a63-8af5-709836799726",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide \n",
    "evaluation = ndcg_1.evaluate_query(\n",
    "    query_results=query_results,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_1_ideal_dcg\": expected_ideal_dcg,\n",
    "        \"ndcg_1_dcg\": expected_dcg,\n",
    "        \"ndcg_1\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe237f38-72b1-42c1-892c-b37a1e5e5463",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs2,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "expected_dcg = 1 / math.log2(2) + 0 / math.log2(3) + 2 / math.log2(4)\n",
    "expected_ideal_dcg = 2 / math.log2(2) + 1 / math.log2(3) + 0 / math.log2(4)\n",
    "expected_ndcg = expected_dcg / expected_ideal_dcg\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_3\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c980bb8-bd37-4699-88f4-ce6b20d13f39",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_3.evaluate_query(\n",
    "    query_results=query_results2,\n",
    "    relevant_docs=relevant_docs2,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_3_ideal_dcg\": expected_ideal_dcg,\n",
    "        \"ndcg_3_dcg\": expected_dcg,\n",
    "        \"ndcg_3\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7309423-dec3-4ce4-975f-046ebf5669d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_3.evaluate_query(\n",
    "    query_results=query_results_int_id,\n",
    "    relevant_docs=relevant_docs_int_id,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    ")\n",
    "expected_dcg = 1 / math.log2(2) + 0 / math.log2(3) + 2 / math.log2(4)\n",
    "expected_ideal_dcg = 2 / math.log2(2) + 1 / math.log2(3) + 0 / math.log2(4)\n",
    "expected_ndcg = expected_dcg / expected_ideal_dcg\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_3\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f83f404f-f092-4b34-9ded-93df54380411",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_3.evaluate_query(\n",
    "    query_results=query_results_int_id,\n",
    "    relevant_docs=relevant_docs_int_id,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_3_ideal_dcg\": expected_ideal_dcg,\n",
    "        \"ndcg_3_dcg\": expected_dcg,\n",
    "        \"ndcg_3\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d3041d8-c501-4fa9-8a93-bcd9f022e49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = ndcg_2.evaluate_query(\n",
    "    query_results=query_results_empty_field,\n",
    "    relevant_docs=relevant_docs,\n",
    "    id_field=\"vespa_id_field\",\n",
    "    default_score=0,\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "expected_dcg = 0 / math.log2(2) + 1 / math.log2(3)\n",
    "expected_ideal_dcg = 1 / math.log2(2) + 1 / math.log2(3)\n",
    "expected_ndcg = expected_dcg / expected_ideal_dcg\n",
    "\n",
    "test_eq(\n",
    "    evaluation,\n",
    "    {\n",
    "        \"ndcg_2_ideal_dcg\": expected_ideal_dcg,\n",
    "        \"ndcg_2_dcg\": expected_dcg,\n",
    "        \"ndcg_2\": expected_ndcg,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d722b46f-35d5-4972-affc-647aadf39058",
   "metadata": {},
   "source": [
    "## Evaluation queries in batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "623f3181-3ee7-4a8c-afde-818e64705fed",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def _evaluate_query_retry(app, flat_labeled_data, model, timeout, **kwargs):\n",
    "    query_responses = send_query_batch(\n",
    "        app=app,\n",
    "        query_batch = [x[0] for x in flat_labeled_data], \n",
    "        query_model = model, \n",
    "        **{\"ranking.softtimeout.enable\": \"false\",\n",
    "          \"timeout\": timeout},\n",
    "        **kwargs)\n",
    "    failed_queries = [idx for idx, x in enumerate(query_responses) if x.status_code != 200]\n",
    "    count = 0\n",
    "    while len(failed_queries) > 0:\n",
    "        query_batch = [flat_labeled_data[idx][0] for idx in failed_queries]\n",
    "        retry_query_responses = send_query_batch(\n",
    "            app=app,\n",
    "            query_batch=query_batch,\n",
    "            query_model = model, \n",
    "            **{\"ranking.softtimeout.enable\": \"false\",\n",
    "              \"timeout\": timeout},           \n",
    "            **kwargs\n",
    "        )\n",
    "        for idx, query_response_idx in enumerate(failed_queries):\n",
    "            query_responses[query_response_idx] = retry_query_responses[idx]\n",
    "        failed_queries = [idx for idx, x in enumerate(query_responses) if x.status_code != 200]\n",
    "        count+=1\n",
    "        if count>=3: break   \n",
    "    return query_responses\n",
    "    \n",
    "def evaluate(\n",
    "    app: Vespa,  # Connection to a Vespa application.\n",
    "    labeled_data: Union[List[Dict], DataFrame],  # Data containing query, query_id and relevant docs. See examples below for format.\n",
    "    eval_metrics: List[EvalMetric],  # Evaluation metrics\n",
    "    query_model: Union[QueryModel, List[QueryModel]],  # Query models to be evaluated\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    default_score: int = 0,  # Score to assign to the additional documents that are not relevant.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available. \n",
    "    per_query=False,  # Set to True to return evaluation metrics per query.\n",
    "    aggregators=None,  # Used only if `per_query=False`. List of pandas friendly aggregators to summarize per model metrics. We use [\"mean\", \"median\", \"std\"] by default.\n",
    "    timeout=1000,  # Vespa query timeout in ms.\n",
    "    **kwargs,  # Extra keyword arguments to be included in the Vespa Query.\n",
    ") -> DataFrame:  # Returns query_id and metrics according to the selected evaluation metrics.\n",
    "    \"Evaluate a `QueryModel` according to a list of `EvalMetric`.\"\n",
    "    \n",
    "    if isinstance(labeled_data, DataFrame):\n",
    "        labeled_data = _parse_labeled_data(df=labeled_data)\n",
    "\n",
    "    if isinstance(query_model, QueryModel):\n",
    "        query_model = [query_model]\n",
    "\n",
    "    model_names = [model.name for model in query_model]\n",
    "    assert len(model_names) == len(\n",
    "        set(model_names)\n",
    "    ), \"Duplicate model names. Choose unique model names.\"\n",
    "\n",
    "    evaluation = []\n",
    "\n",
    "    for model in query_model:\n",
    "        flat_labeled_data = [(x[\"query\"], x[\"query_id\"], x[\"relevant_docs\"]) for x in labeled_data]\n",
    "        query_responses = _evaluate_query_retry(app, flat_labeled_data, model, timeout, **kwargs)\n",
    "        failed_queries = [idx for idx, x in enumerate(query_responses) if x.status_code != 200]\n",
    "        if len(failed_queries) > 0:\n",
    "            print(f\"Failed queries for query model {model.name}: {len(failed_queries)}/{len(query_responses)}\")\n",
    "        timedout_queries = [idx for idx, x in enumerate(query_responses) if x.json.get(\"root\", {}).get(\"errors\", None) is not None]\n",
    "        if len(timedout_queries) > 0:\n",
    "            print(f\"Timeout queries for query model {model.name}: {len(timedout_queries)}/{len(query_responses)}\")\n",
    "            \n",
    "        for ((query, query_id, relevant_docs), query_response) in zip(flat_labeled_data, query_responses):\n",
    "            evaluation_query = {\"model\": model.name, \"query_id\": query_id}\n",
    "            for evaluator in eval_metrics:\n",
    "                evaluation_query.update(\n",
    "                    evaluator.evaluate_query(\n",
    "                        query_response,\n",
    "                        relevant_docs,\n",
    "                        id_field,\n",
    "                        default_score,\n",
    "                        detailed_metrics,\n",
    "                    )\n",
    "                )\n",
    "            evaluation.append(evaluation_query)\n",
    "    evaluation = DataFrame.from_records(evaluation)\n",
    "    if not per_query:\n",
    "        if not aggregators:\n",
    "            aggregators = [\"mean\", \"median\", \"std\"]\n",
    "        evaluation = (\n",
    "            evaluation[[x for x in evaluation.columns if x != \"query_id\"]]\n",
    "            .groupby(by=\"model\")\n",
    "            .agg(aggregators)\n",
    "            .T\n",
    "        )\n",
    "    return evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72c17e6a-6f2b-468c-9d44-a8dbe86393a8",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d5d2c09-4725-48e8-bf68-530c6444f8a5",
   "metadata": {},
   "source": [
    "Setup and feed a Vespa application:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c409876-8be0-4695-a5f1-0e8cc3128d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "from learntorank.passage import create_basic_search_package\n",
    "from learntorank.passage import PassageData\n",
    "from vespa.deployment import VespaDocker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f0d3356-fbe1-49b1-869a-f459fb1eda82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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 configuration server, 15/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",
      "Finished deployment.\n",
      "Successful documents fed: 1000/1000.\n",
      "Batch progress: 1/1.\n"
     ]
    }
   ],
   "source": [
    "#|output: false\n",
    "app_package = create_basic_search_package(name=\"EvaluationApp\")\n",
    "vespa_docker = VespaDocker(port=8082, cfgsrv_port=19072)\n",
    "app = vespa_docker.deploy(application_package=app_package)\n",
    "data = PassageData.load()\n",
    "responses = app.feed_df(\n",
    "    df=data.get_corpus(), \n",
    "    include_id=True, \n",
    "    id_field=\"doc_id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42fae8bb-aca6-466a-ac64-1e2fea592d47",
   "metadata": {},
   "source": [
    "Define query models to be evaluated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4acbdeae-237c-4856-b60f-9d337bf8ad01",
   "metadata": {},
   "outputs": [],
   "source": [
    "from learntorank.query import OR, Ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305348d8-3c91-413c-b3a8-30e0b50d3250",
   "metadata": {},
   "outputs": [],
   "source": [
    "bm25_query_model = QueryModel(\n",
    "    name=\"bm25\", \n",
    "    match_phase=OR(), \n",
    "    ranking=Ranking(name=\"bm25\")\n",
    ")\n",
    "native_query_model = QueryModel(\n",
    "    name=\"native_rank\", \n",
    "    match_phase=OR(), \n",
    "    ranking=Ranking(name=\"native_rank\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bab6315-d908-4ccd-a727-485a6fc72bf0",
   "metadata": {},
   "source": [
    "Define metrics to compute during evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6453cdcd-e5d0-41ce-b686-39ef7b3fde3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = [\n",
    "    Recall(at=10), \n",
    "    ReciprocalRank(at=3), \n",
    "    NormalizedDiscountedCumulativeGain(at=3)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "882cdde7-e26c-4f4e-a378-f447482f4592",
   "metadata": {},
   "source": [
    "Get labeled data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb7e842b-5340-49af-9e39-3deb782bcf0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'query_id': '1101971',\n",
       "  'query': 'why say the sky is the limit',\n",
       "  'relevant_docs': [{'id': '7407715', 'score': 1}]},\n",
       " {'query_id': '712898',\n",
       "  'query': 'what is an cvc in radiology',\n",
       "  'relevant_docs': [{'id': '7661336', 'score': 1}]}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labeled_data = data.get_labels(type=\"dev\")\n",
    "labeled_data[0:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef506c06-79ee-41c3-b5d8-282b8a27b4f4",
   "metadata": {},
   "source": [
    "Evaluate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe1f373-924e-4fc2-a965-136a89d9a40f",
   "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>model</th>\n",
       "      <th>bm25</th>\n",
       "      <th>native_rank</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">recall_10</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935833</td>\n",
       "      <td>0.845833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.215444</td>\n",
       "      <td>0.342749</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">reciprocal_rank_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935000</td>\n",
       "      <td>0.746667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.231977</td>\n",
       "      <td>0.399551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">ndcg_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.912839</td>\n",
       "      <td>0.740814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.242272</td>\n",
       "      <td>0.387611</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "model                         bm25  native_rank\n",
       "recall_10         mean    0.935833     0.845833\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.215444     0.342749\n",
       "reciprocal_rank_3 mean    0.935000     0.746667\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.231977     0.399551\n",
       "ndcg_3            mean    0.912839     0.740814\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.242272     0.387611"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data, \n",
    "    eval_metrics=metrics, \n",
    "    query_model=[native_query_model, bm25_query_model], \n",
    "    id_field=\"doc_id\",\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33294e2e-8f4c-4832-9afa-5b0b42c8cbca",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation.shape, (9,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e1c7204-b2d1-4695-b4fb-aabf5dba27d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "labeled_df = DataFrame.from_records(\n",
    "    [\n",
    "        {\n",
    "            \"qid\": str(q[\"query_id\"]), \n",
    "            \"query\": str(q[\"query\"]), \n",
    "            \"doc_id\": str(d[\"id\"]), \n",
    "            \"relevance\": int(d[\"score\"])\n",
    "        } for q in labeled_data for d in q[\"relevant_docs\"]\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75425a1e-b0da-43e6-adf7-80c721baa41b",
   "metadata": {},
   "source": [
    "The evaluate function also accepts labeled data as a data frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14766c87-3ae9-40b7-8f39-d7ace8068e9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qid</th>\n",
       "      <th>query</th>\n",
       "      <th>doc_id</th>\n",
       "      <th>relevance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1101971</td>\n",
       "      <td>why say the sky is the limit</td>\n",
       "      <td>7407715</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>712898</td>\n",
       "      <td>what is an cvc in radiology</td>\n",
       "      <td>7661336</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>154469</td>\n",
       "      <td>dmv california how long does it take to get id</td>\n",
       "      <td>7914544</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>930015</td>\n",
       "      <td>what's an epigraph</td>\n",
       "      <td>7928705</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>860085</td>\n",
       "      <td>what is va tax</td>\n",
       "      <td>2915383</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       qid                                           query   doc_id  relevance\n",
       "0  1101971                    why say the sky is the limit  7407715          1\n",
       "1   712898                     what is an cvc in radiology  7661336          1\n",
       "2   154469  dmv california how long does it take to get id  7914544          1\n",
       "3   930015                              what's an epigraph  7928705          1\n",
       "4   860085                                  what is va tax  2915383          1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labeled_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b0b05a0-f10f-46d5-98f4-1a1023ceb46a",
   "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>model</th>\n",
       "      <th>bm25</th>\n",
       "      <th>native_rank</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">recall_10</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935833</td>\n",
       "      <td>0.845833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.215444</td>\n",
       "      <td>0.342749</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">reciprocal_rank_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935000</td>\n",
       "      <td>0.746667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.231977</td>\n",
       "      <td>0.399551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"3\" valign=\"top\">ndcg_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.912839</td>\n",
       "      <td>0.740814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>median</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.242272</td>\n",
       "      <td>0.387611</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "model                         bm25  native_rank\n",
       "recall_10         mean    0.935833     0.845833\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.215444     0.342749\n",
       "reciprocal_rank_3 mean    0.935000     0.746667\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.231977     0.399551\n",
       "ndcg_3            mean    0.912839     0.740814\n",
       "                  median  1.000000     1.000000\n",
       "                  std     0.242272     0.387611"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation_df = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_df, \n",
    "    eval_metrics=metrics, \n",
    "    query_model=[native_query_model, bm25_query_model], \n",
    "    id_field=\"doc_id\",\n",
    ")\n",
    "evaluation_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71dceafa-c04b-4f35-92e6-c3bc44ba7016",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation, evaluation_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "145c7c80-4d06-4ca1-a62b-27cd477e22ee",
   "metadata": {},
   "source": [
    "Control which aggregators are computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dbe598d-f5b2-494b-90d8-57500412465c",
   "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>model</th>\n",
       "      <th>bm25</th>\n",
       "      <th>native_rank</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">recall_10</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935833</td>\n",
       "      <td>0.845833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.215444</td>\n",
       "      <td>0.342749</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">reciprocal_rank_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935000</td>\n",
       "      <td>0.746667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.231977</td>\n",
       "      <td>0.399551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ndcg_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.912839</td>\n",
       "      <td>0.740814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.242272</td>\n",
       "      <td>0.387611</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "model                       bm25  native_rank\n",
       "recall_10         mean  0.935833     0.845833\n",
       "                  std   0.215444     0.342749\n",
       "reciprocal_rank_3 mean  0.935000     0.746667\n",
       "                  std   0.231977     0.399551\n",
       "ndcg_3            mean  0.912839     0.740814\n",
       "                  std   0.242272     0.387611"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data, \n",
    "    eval_metrics=metrics, \n",
    "    query_model=[native_query_model, bm25_query_model], \n",
    "    id_field=\"doc_id\",\n",
    "    aggregators=[\"mean\", \"std\"]\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b50d7523-042c-4302-b5d1-4d24bb7abe3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation.shape, (6,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e12b16-15f6-4e8e-a845-2c1975945ee7",
   "metadata": {},
   "source": [
    "Include detailed metrics when available, this includes intermediate steps that are available for some of the metrics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4949070-1c44-47f6-bd16-4c5fe3bb4fa1",
   "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>model</th>\n",
       "      <th>bm25</th>\n",
       "      <th>native_rank</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">recall_10</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935833</td>\n",
       "      <td>0.845833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.215444</td>\n",
       "      <td>0.342749</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">reciprocal_rank_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.935000</td>\n",
       "      <td>0.746667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.231977</td>\n",
       "      <td>0.399551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ndcg_3</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.912839</td>\n",
       "      <td>0.740814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.242272</td>\n",
       "      <td>0.387611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ndcg_3_ideal_dcg</th>\n",
       "      <th>mean</th>\n",
       "      <td>1.054165</td>\n",
       "      <td>1.054165</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.207315</td>\n",
       "      <td>0.207315</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">ndcg_3_dcg</th>\n",
       "      <th>mean</th>\n",
       "      <td>0.938928</td>\n",
       "      <td>0.765474</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.225533</td>\n",
       "      <td>0.387161</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "model                       bm25  native_rank\n",
       "recall_10         mean  0.935833     0.845833\n",
       "                  std   0.215444     0.342749\n",
       "reciprocal_rank_3 mean  0.935000     0.746667\n",
       "                  std   0.231977     0.399551\n",
       "ndcg_3            mean  0.912839     0.740814\n",
       "                  std   0.242272     0.387611\n",
       "ndcg_3_ideal_dcg  mean  1.054165     1.054165\n",
       "                  std   0.207315     0.207315\n",
       "ndcg_3_dcg        mean  0.938928     0.765474\n",
       "                  std   0.225533     0.387161"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data, \n",
    "    eval_metrics=metrics, \n",
    "    query_model=[native_query_model, bm25_query_model], \n",
    "    id_field=\"doc_id\",\n",
    "    aggregators=[\"mean\", \"std\"],\n",
    "    detailed_metrics=True\n",
    ")\n",
    "evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24a9d247-ca7c-48c9-af9f-4563290cdc82",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation.shape, (10,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a5d60a7-cb72-4dde-918e-4264f4020293",
   "metadata": {},
   "source": [
    "Generate results per query:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14516da7-dea2-4aea-b6fb-2d74b53d8911",
   "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>model</th>\n",
       "      <th>query_id</th>\n",
       "      <th>recall_10</th>\n",
       "      <th>reciprocal_rank_3</th>\n",
       "      <th>ndcg_3</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>native_rank</td>\n",
       "      <td>1101971</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>native_rank</td>\n",
       "      <td>712898</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>native_rank</td>\n",
       "      <td>154469</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>native_rank</td>\n",
       "      <td>930015</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>native_rank</td>\n",
       "      <td>860085</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         model query_id  recall_10  reciprocal_rank_3  ndcg_3\n",
       "0  native_rank  1101971        1.0                1.0     1.0\n",
       "1  native_rank   712898        0.0                0.0     0.0\n",
       "2  native_rank   154469        1.0                0.0     0.0\n",
       "3  native_rank   930015        1.0                0.0     0.0\n",
       "4  native_rank   860085        0.0                0.0     0.0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data, \n",
    "    eval_metrics=metrics, \n",
    "    query_model=[native_query_model, bm25_query_model], \n",
    "    id_field=\"doc_id\",\n",
    "    per_query=True\n",
    ")\n",
    "evaluation.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3892ed19-5b54-4c35-a967-fcc8e982ba08",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(evaluation.shape, (200,5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65ffba0c-7d3d-441e-924a-7915d19215a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "vespa_docker.container.stop(timeout=600)\n",
    "vespa_docker.container.remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4b5d6e-3858-4ac2-8774-4c49e3c92055",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from learntorank.query import WeakAnd, ANN, Union, Ranking, QueryRankingFeature\n",
    "from random import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35b489b3-eaed-44c3-ac77-c59c2f1f17d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "#\n",
    "# Connect to a running Vespa Application\n",
    "#\n",
    "app = Vespa(url=\"https://api.cord19.vespa.ai\")\n",
    "#\n",
    "# Define a query model\n",
    "#\n",
    "match_phase = Union(\n",
    "    WeakAnd(hits=10),\n",
    "    ANN(\n",
    "        doc_vector=\"title_embedding\",\n",
    "        query_vector=\"title_vector\",\n",
    "        hits=10,\n",
    "        label=\"title\",\n",
    "    ),\n",
    ")\n",
    "ranking = Ranking(name=\"bm25\", list_features=True)\n",
    "query_model = QueryModel(\n",
    "    name=\"ANN_bm25\",\n",
    "    query_properties=[\n",
    "        QueryRankingFeature(\n",
    "            name=\"title_vector\",\n",
    "            mapping=lambda x: [random() for x in range(768)],\n",
    "        )\n",
    "    ],\n",
    "    match_phase=match_phase,\n",
    "    ranking=ranking,\n",
    ")\n",
    "#\n",
    "# Define labelled data\n",
    "#\n",
    "labeled_data = [\n",
    "    {\n",
    "        \"query_id\": 0,\n",
    "        \"query\": \"Intrauterine virus infections and congenital heart disease\",\n",
    "        \"relevant_docs\": [{\"id\": 0, \"score\": 1}, {\"id\": 3, \"score\": 1}],\n",
    "    },\n",
    "    {\n",
    "        \"query_id\": 1,\n",
    "        \"query\": \"Clinical and immunologic studies in identical twins discordant for systemic lupus erythematosus\",\n",
    "        \"relevant_docs\": [{\"id\": 1, \"score\": 1}, {\"id\": 5, \"score\": 1}],\n",
    "    },\n",
    "]\n",
    "# equivalent data in df format\n",
    "labeled_data_df = DataFrame(\n",
    "    data={\n",
    "        \"qid\": [0, 0, 1, 1],\n",
    "        \"query\": [\"Intrauterine virus infections and congenital heart disease\"]\n",
    "        * 2\n",
    "        + [\n",
    "            \"Clinical and immunologic studies in identical twins discordant for systemic lupus erythematosus\"\n",
    "        ]\n",
    "        * 2,\n",
    "        \"doc_id\": [0, 3, 1, 5],\n",
    "        \"relevance\": [1, 1, 1, 1],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dd548f6-a4b4-4ba4-945f-fa1238c27388",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "eval_metrics = [\n",
    "    MatchRatio(), \n",
    "    Recall(at=10), \n",
    "    ReciprocalRank(at=10)\n",
    "]\n",
    "ltr_evaluation = evaluate(\n",
    "    app=app, \n",
    "    labeled_data=labeled_data, \n",
    "    eval_metrics=eval_metrics, \n",
    "    query_model=query_model, \n",
    "    id_field=\"id\"\n",
    ")\n",
    "test_eq(ltr_evaluation.shape, (9, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a02af19c-9344-41db-b289-17ef5677dc77",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_fail(\n",
    "    evaluate, kwargs={\n",
    "        \"app\": app, \n",
    "        \"labeled_data\": labeled_data,\n",
    "        \"eval_metrics\": eval_metrics,\n",
    "        \"query_model\": [QueryModel(), QueryModel(), query_model],\n",
    "        \"id_field\": \"id\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "701997b8-bd3e-40cc-b340-64d66a3ae022",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data,\n",
    "    eval_metrics=eval_metrics,\n",
    "    query_model=[QueryModel(), query_model],\n",
    "    id_field=\"id\",\n",
    ")\n",
    "test_eq(evaluation.shape, (9, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c3e42d3-fe89-4de4-8e85-781cc47a91c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data_df,\n",
    "    eval_metrics=eval_metrics,\n",
    "    query_model=query_model,\n",
    "    id_field=\"id\",\n",
    "    detailed_metrics=True,\n",
    ")\n",
    "test_eq(evaluation.shape, (15, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4501eda-4b49-4c49-aebc-9d4b8a7df1a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "evaluation = evaluate(\n",
    "    app=app,\n",
    "    labeled_data=labeled_data_df,\n",
    "    eval_metrics=eval_metrics,\n",
    "    query_model=query_model,\n",
    "    id_field=\"id\",\n",
    "    detailed_metrics=True,\n",
    "    per_query=True,\n",
    ")\n",
    "test_eq(evaluation.shape, (2, 7))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98ede4e8-146a-42d8-bfc5-ef2e3831790d",
   "metadata": {},
   "source": [
    "## Evaluate specific query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d224af52-3317-471e-9c5d-2497ca3a8226",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def evaluate_query(\n",
    "    app: Vespa,  # Connection to a Vespa application.\n",
    "    eval_metrics: List[EvalMetric],  # Evaluation metrics\n",
    "    query_model: QueryModel,  # Query model to be evaluated  \n",
    "    query_id: str,  # Query id represented as str.\n",
    "    query: str,  # Query string.\n",
    "    id_field: str,  # The Vespa field representing the document id.\n",
    "    relevant_docs: List[Dict],  # Each dict contains a doc id a optionally a doc score.\n",
    "    default_score: int = 0,  # Score to assign to the additional documents that are not relevant.\n",
    "    detailed_metrics=False,  # Return intermediate computations if available.\n",
    "    **kwargs,  # Extra keyword arguments to be included in the Vespa Query.\n",
    ") -> Dict:  # Contains query_id and metrics according to the selected evaluation metrics.\n",
    "    \"Evaluate a single query according to evaluation metrics\"\n",
    "    \n",
    "    query_results = send_query(\n",
    "        app=app, \n",
    "        query=query, \n",
    "        query_model=query_model, \n",
    "        **kwargs\n",
    "    )\n",
    "    evaluation = {\"model\": query_model.name, \"query_id\": query_id}\n",
    "    for evaluator in eval_metrics:\n",
    "        evaluation.update(\n",
    "            evaluator.evaluate_query(\n",
    "                query_results,\n",
    "                relevant_docs,\n",
    "                id_field,\n",
    "                default_score,\n",
    "                detailed_metrics,\n",
    "            )\n",
    "        )\n",
    "    return evaluation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe31a67d-2b9d-40e2-89a3-bb42bb58428a",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e84ca85e-ab20-4b7d-a93e-e76f7c13fe6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = Vespa(url = \"https://api.cord19.vespa.ai\")\n",
    "query_model = QueryModel(\n",
    "    match_phase = OR(),\n",
    "    ranking = Ranking(name=\"bm25\", list_features=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc810af-db2e-4fbb-b94c-5321e890d948",
   "metadata": {},
   "source": [
    "Evaluate a single query:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d45c695-79e5-452d-84b0-89c6d233dbb5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model': 'bm25',\n",
       " 'query_id': '0',\n",
       " 'match_ratio': 0.814424921006077,\n",
       " 'recall_10': 0.0,\n",
       " 'reciprocal_rank_10': 0}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_evaluation = evaluate_query(\n",
    "    app=app,\n",
    "    eval_metrics = eval_metrics, \n",
    "    query_model = bm25_query_model, \n",
    "    query_id = \"0\", \n",
    "    query = \"Intrauterine virus infections and congenital heart disease\", \n",
    "    id_field = \"id\",\n",
    "    relevant_docs = [{\"id\": 0, \"score\": 1}, {\"id\": 3, \"score\": 1}],\n",
    "    default_score = 0\n",
    ")\n",
    "query_evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb89cf1-b590-41d5-a527-949be3ca2af0",
   "metadata": {},
   "source": [
    "## Evaluate query under specific document ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e89d4cf-40f6-4c0e-b64f-1a7b6fb02aab",
   "metadata": {},
   "source": [
    "Use `recall` to specify which documents should be included in the evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d46443c-c46c-4b02-8a78-9c52fd7b448f",
   "metadata": {},
   "source": [
    "In the example below, we include documents with id equal to 0, 1 and 2. Since the relevant documents for this query are the documents with id 0 and 3, we should get recall equal to 0.5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "202d9e2e-be04-4841-903b-1bbf575c39b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model': 'default_name',\n",
       " 'query_id': 0,\n",
       " 'match_ratio': 9.70242657688688e-06,\n",
       " 'recall_10': 0.5,\n",
       " 'reciprocal_rank_10': 1.0}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_evaluation = evaluate_query(\n",
    "    app=app,\n",
    "    eval_metrics = eval_metrics, \n",
    "    query_model = query_model, \n",
    "    query_id = 0, \n",
    "    query = \"Intrauterine virus infections and congenital heart disease\", \n",
    "    id_field = \"id\",\n",
    "    relevant_docs = [{\"id\": 0, \"score\": 1}, {\"id\": 3, \"score\": 1}],\n",
    "    default_score = 0,\n",
    "    recall = (\"id\", [0, 1, 2])\n",
    ")\n",
    "query_evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3a243c-ce56-4691-95a9-bcc411edb2ae",
   "metadata": {},
   "source": [
    "We now include documents with id equal to 0, 1, 2 and 3. This should give a recall equal to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c3e5de2-7f31-4828-8aa3-b835ec37d3e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model': 'default_name',\n",
       " 'query_id': 0,\n",
       " 'match_ratio': 1.2936568769182506e-05,\n",
       " 'recall_10': 1.0,\n",
       " 'reciprocal_rank_10': 1.0}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_evaluation = evaluate_query(\n",
    "    app=app,\n",
    "    eval_metrics = eval_metrics, \n",
    "    query_model = query_model, \n",
    "    query_id = 0, \n",
    "    query = \"Intrauterine virus infections and congenital heart disease\", \n",
    "    id_field = \"id\",\n",
    "    relevant_docs = [{\"id\": 0, \"score\": 1}, {\"id\": 3, \"score\": 1}],\n",
    "    default_score = 0,\n",
    "    recall = (\"id\", [0, 1, 2, 3])\n",
    ")\n",
    "query_evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5171e9fc-a09a-404d-837f-d8f2844c6749",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "test_eq(query_evaluation[\"recall_10\"], 1.0)"
   ]
  },
  {
   "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
}