{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "59b73671",
   "metadata": {},
   "source": [
    "# Information Retrieval — Class 1: Foundations & Evaluation Metrics\n",
    "\n",
    "**Date:** March 13, 2026\n",
    "\n",
    "**Goal:** Build your first document retrieval system on real web-search data (MS MARCO), then evaluate it with proper IR metrics.\n",
    "\n",
    "By the end of this session you will have:\n",
    "- Loaded and explored a real IR benchmark dataset (1,000 queries, 92k documents)\n",
    "- Built two keyword-based retrieval systems (boolean + TF-IDF)\n",
    "- Implemented five IR evaluation metrics **from scratch**: P@k, R@k, MRR, MAP, nDCG\n",
    "- Diagnosed where simple retrieval fails — setting the stage for the rest of the course\n",
    "\n",
    "**Instructions:**\n",
    "1. Run the cells to see the results.\n",
    "2. Look for **🚧 TODO:** markers. These are exercises for you to complete.\n",
    "3. For \"Explain\" questions, write your answer in the markdown cell provided.\n",
    "4. You have **internet access** and may use **Copilot** — but make sure you understand every line you write.\n",
    "\n",
    "This notebook is designed to be completed in class (1h30)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99cf2f2e",
   "metadata": {},
   "source": [
    "<span style=\"color:magenta\">Group Members:</span>\n",
    "\n",
    "* Name 1\n",
    "* Name 2\n",
    "* Name 3\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc0f183a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "import urllib.request\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer\n",
    "from tqdm import tqdm\n",
    "\n",
    "tqdm.pandas()\n",
    "sns.set_style(\"darkgrid\")\n",
    "sns.set_context(\"notebook\")\n",
    "pd.set_option(\"display.max_colwidth\", 120)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ace7ea7",
   "metadata": {},
   "source": [
    "## Download the MS MARCO Subset\n",
    "\n",
    "[MS MARCO](https://microsoft.github.io/msmarco/) (Microsoft MAchine Reading COmprehension) is a large-scale information retrieval benchmark built from real Bing search queries. We use a curated subset of ~1,000 queries and ~92,000 documents, small enough to run on any laptop.\n",
    "\n",
    "The dataset contains three files:\n",
    "- **queries:** a query ID and the actual search query text\n",
    "- **docs:** a document ID, title, and body text\n",
    "- **gold:** for each query, the top-100 documents ranked by a strong baseline system (Indri), which we use as relevance judgments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d2071e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and extract MS MARCO data if not already present\n",
    "QUERIES_FILE = Path(\"our.msmarco.queries.tsv\")\n",
    "\n",
    "if not QUERIES_FILE.exists():\n",
    "    print(\"Downloading MS MARCO subset...\")\n",
    "    url = \"https://gerdes.fr/saclay/inforet/our_msmarco.zip\"\n",
    "    zip_path = \"our_msmarco.zip\"\n",
    "    urllib.request.urlretrieve(url, zip_path)\n",
    "    print(\"Extracting...\")\n",
    "    with zipfile.ZipFile(zip_path, 'r') as z:\n",
    "        z.extractall(\".\")\n",
    "    os.remove(zip_path)\n",
    "    print(\"Done!\")\n",
    "else:\n",
    "    print(\"Data files already present.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941bf723",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1: Exploring MS MARCO (~20 min)\n",
    "\n",
    "Before building any retrieval system, we must understand the data.\n",
    "\n",
    "### 1.1 Load and Explore the Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7929d2d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = pd.read_csv(\"our.msmarco.queries.tsv\", sep=\"\\t\", usecols=[1, 2])\n",
    "print(f\"{len(queries)} queries loaded\")\n",
    "queries.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8a5e80",
   "metadata": {},
   "source": [
    "**🚧 TODO 1.1:** Compute basic statistics about the queries.\n",
    "1. What is the average number of **words** per query? (split on whitespace)\n",
    "2. What is the shortest query? The longest?\n",
    "3. Plot a histogram of query lengths (in words).\n",
    "4. What do you notice about the distribution? How does this compare to what you'd expect from a web search engine?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fcd1287",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 1.1: Query statistics\n",
    "\n",
    "queries[\"n_words\"] = ... your code here ...\n",
    "\n",
    "print(f\"Average words per query: ...\")\n",
    "print(f\"Shortest query: ...\")\n",
    "print(f\"Longest query: ...\")\n",
    "\n",
    "plt.figure(figsize=(10, 4))\n",
    "... your code here ...\n",
    "plt.xlabel(\"Number of words\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.title(\"Distribution of Query Lengths (MS MARCO)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8fca869",
   "metadata": {},
   "source": [
    "**Answer 1.1:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### 1.2 Load and Explore the Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b000170",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = pd.read_csv(\"our.msmarco.docs.tsv\", sep=\"\\t\", usecols=[1, 2, 3])\n",
    "print(f\"{len(docs)} documents loaded\")\n",
    "\n",
    "# Combine title and body into a single text field\n",
    "docs[\"text\"] = docs[\"title\"].fillna(\"\") + \" \" + docs[\"body\"].fillna(\"\")\n",
    "docs.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ccf5a57",
   "metadata": {},
   "source": [
    "**🚧 TODO 1.2:** Explore the documents.\n",
    "1. What is the average document length (in words)?\n",
    "2. Plot the distribution of document lengths. Use `np.clip` to cap at 2000 words for readability.\n",
    "3. Print 3 random documents (first 300 characters each). What types of web content do you see?\n",
    "4. **Hard question:** Why do you think document lengths vary so much more than query lengths? What consequences does this have for retrieval?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe18fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 1.2: Document statistics\n",
    "\n",
    "docs[\"n_words\"] = ... your code here ...\n",
    "\n",
    "print(f\"Average words per document: ...\")\n",
    "print(f\"Median: ...\")\n",
    "\n",
    "plt.figure(figsize=(10, 4))\n",
    "... your code here ...\n",
    "plt.title(\"Distribution of Document Lengths (MS MARCO)\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\n--- 3 Random Documents (first 300 chars) ---\")\n",
    "... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79073f36",
   "metadata": {},
   "source": [
    "**Answer 1.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### 1.3 Load and Explore the Gold Judgments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9e65c2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gold = pd.read_csv(\"our.msmarco.gold.tsv\", sep=\"\\t\", usecols=[1, 3])\n",
    "print(f\"{len(gold)} gold judgments loaded\")\n",
    "gold.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e4ad9f5",
   "metadata": {},
   "source": [
    "**🚧 TODO 1.3:** Analyze the gold judgments.\n",
    "1. How many unique queries have gold judgments?\n",
    "2. How many gold documents does each query have on average? What is the distribution?\n",
    "3. Create a dictionary `gold_dict` that maps each `qid` to the **set** of its relevant `docid`s. This will be essential for evaluation later.\n",
    "4. Split the queries into a **training** set (first 500) and a **test** set (remaining). We'll use the training set for development and the test set for final evaluation. How many queries does each split have?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b352483",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 1.3: Gold statistics and data preparation\n",
    "\n",
    "# 1. Unique queries with gold judgments\n",
    "print(f\"Queries with gold judgments: ...\")\n",
    "\n",
    "# 2. Gold docs per query\n",
    "... your code here ...\n",
    "\n",
    "# 3. Gold dictionary for fast lookup\n",
    "gold_dict = ... your code here ...\n",
    "\n",
    "# 4. Train/test split\n",
    "train_queries = ... your code here ...\n",
    "test_queries = ... your code here ...\n",
    "print(f\"Training queries: {len(train_queries)}\")\n",
    "print(f\"Test queries: {len(test_queries)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0604601e",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 2: Building a Retrieval System (~25 min)\n",
    "\n",
    "In classification (HoNLP), you assigned labels to documents. In retrieval, you **rank** all documents by relevance to a query. The simplest approach: count how many query terms appear in each document.\n",
    "\n",
    "### 2.1 Building a Boolean Retrieval Index\n",
    "\n",
    "A **boolean retrieval** model asks a simple binary question for each term: does it appear in the document (1) or not (0)? The score is the number of matching query terms.\n",
    "\n",
    "$$\\text{score}(q, d) = \\sum_{t \\in q} \\mathbf{1}[t \\in d]$$\n",
    "\n",
    "Scikit-learn's `CountVectorizer(binary=True)` builds exactly this: a sparse term-document matrix with 1s and 0s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c9ffbf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the boolean document index (this takes ~15-30 seconds on 92k docs)\n",
    "print(\"Building boolean index...\")\n",
    "bool_vectorizer = CountVectorizer(binary=True)\n",
    "bool_matrix = bool_vectorizer.fit_transform(docs[\"text\"])\n",
    "print(f\"Index shape: {bool_matrix.shape} (documents × vocabulary)\")\n",
    "print(f\"Vocabulary size: {len(bool_vectorizer.vocabulary_)}\")\n",
    "print(f\"Sparsity: {1 - bool_matrix.nnz / (bool_matrix.shape[0] * bool_matrix.shape[1]):.4%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad7aed6",
   "metadata": {},
   "source": [
    "**🚧 TODO 2.2:** Write a function `retrieve(query_text, vectorizer, doc_matrix, k=10)` that:\n",
    "1. Transforms the query into the same vector space as the documents\n",
    "2. Computes the dot product between the query vector and all document vectors\n",
    "3. Returns the **indices** of the top-k highest-scoring documents\n",
    "\n",
    "Use sparse matrix multiplication (`doc_matrix @ query_vec.T`) for efficiency — never `.toarray()` on the full matrix!\n",
    "\n",
    "**Hint:** `np.argpartition` is faster than `np.argsort` when you only need the top-k.\n",
    "\n",
    "Then write a helper function `show_results(query_text, vectorizer, doc_matrix, k=10)` that prints the query and the top-k retrieved documents (first 200 characters each), along with their scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e44a8371",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 2.2: Retrieval function\n",
    "\n",
    "def retrieve(query_text, vectorizer, doc_matrix, k=10):\n",
    "    \"\"\"Retrieve top-k document indices for a query.\"\"\"\n",
    "    query_vec = vectorizer.transform([query_text])\n",
    "    scores = ... your code here ...\n",
    "    top_k_idx = ... your code here ...\n",
    "    return top_k_idx, scores[top_k_idx]\n",
    "\n",
    "\n",
    "def show_results(query_text, vectorizer, doc_matrix, k=10):\n",
    "    \"\"\"Retrieve and display top-k results for a query.\"\"\"\n",
    "    top_k_idx, top_k_scores = retrieve(query_text, vectorizer, doc_matrix, k)\n",
    "    print(f\"Query: {query_text!r}\")\n",
    "    print(f\"{'='*80}\")\n",
    "    for rank, (idx, score) in enumerate(zip(top_k_idx, top_k_scores), 1):\n",
    "        docid = docs.iloc[idx][\"docid\"]\n",
    "        text_preview = docs.iloc[idx][\"text\"][:200].replace(\"\\n\", \" \")\n",
    "        print(f\"  #{rank} [score={score:.2f}] {docid}: {text_preview}...\")\n",
    "    print()\n",
    "\n",
    "\n",
    "# Test it!\n",
    "show_results(\"what is a po box\", bool_vectorizer, bool_matrix, k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f0f6e96",
   "metadata": {},
   "source": [
    "**🚧 TODO 2.3:** Test your retrieval system on these queries and comment on the quality:\n",
    "1. `\"what is a po box\"`\n",
    "2. `\"how long does it take to get to mars\"`\n",
    "3. `\"python programming language\"`\n",
    "4. `\"heart attack symptoms\"`\n",
    "\n",
    "For each query, look at the top-5 results. Are they relevant? For the **bad** results, can you explain why the system failed?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "418bbac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 2.3: Test and comment\n",
    "\n",
    "test_queries_list = [\n",
    "    \"what is a po box\",\n",
    "    \"how long does it take to get to mars\",\n",
    "    \"python programming language\",\n",
    "    \"heart attack symptoms\",\n",
    "]\n",
    "\n",
    "for q in test_queries_list:\n",
    "    show_results(q, bool_vectorizer, bool_matrix, k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "486bb6b5",
   "metadata": {},
   "source": [
    "**Answer 2.3:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### 2.4 Upgrade: TF-IDF Retrieval\n",
    "\n",
    "**🚧 TODO 2.4:** Build a TF-IDF retrieval system using `TfidfVectorizer` and compare it to boolean retrieval on the same queries.\n",
    "\n",
    "Use `sublinear_tf=True` (replaces raw TF with $1 + \\log(\\text{tf})$) and `strip_accents='unicode'` for better handling of accented characters.\n",
    "\n",
    "**Questions:**\n",
    "1. Why should `sublinear_tf=True` help? What would go wrong with raw term frequencies?\n",
    "2. Where exactly in the TF-IDF formula does the \"common words\" problem get solved?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20181468",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 2.4: TF-IDF retrieval\n",
    "\n",
    "print(\"Building TF-IDF index...\")\n",
    "tfidf_vectorizer = ... your code here ...\n",
    "tfidf_matrix = ... your code here ...\n",
    "print(f\"TF-IDF matrix shape: {tfidf_matrix.shape}\")\n",
    "\n",
    "# Compare on the same queries\n",
    "for q in test_queries_list:\n",
    "    show_results(q, tfidf_vectorizer, tfidf_matrix, k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f5df28b",
   "metadata": {},
   "source": [
    "**Answer 2.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "## Part 3: Evaluation Metrics — From Scratch (~30 min)\n",
    "\n",
    "How do you measure whether a ranked list is \"good\"? Unlike classification where you check if the label matches, in retrieval you must evaluate the **quality of a ranking**. We'll implement five metrics from scratch.\n",
    "\n",
    "### Notation\n",
    "\n",
    "- A **query** $q$ has a set of **relevant documents** $\\text{Rel}_q$ (from our gold file)\n",
    "- Our system returns a **ranked list** of document IDs: $[d_1, d_2, ..., d_k]$\n",
    "- $\\text{rel}(i) = 1$ if $d_i \\in \\text{Rel}_q$, else 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6dee1d0",
   "metadata": {},
   "source": [
    "### 3.1 Precision@k and Recall@k\n",
    "\n",
    "**Precision@k** = fraction of top-k results that are relevant:\n",
    "\n",
    "$$P@k = \\frac{|\\{d_i : i \\leq k, d_i \\in \\text{Rel}_q\\}|}{k}$$\n",
    "\n",
    "**Recall@k** = fraction of all relevant documents found in the top-k:\n",
    "\n",
    "$$R@k = \\frac{|\\{d_i : i \\leq k, d_i \\in \\text{Rel}_q\\}|}{|\\text{Rel}_q|}$$\n",
    "\n",
    "**🚧 TODO 3.1:** Implement `precision_at_k` and `recall_at_k`.\n",
    "\n",
    "Inputs:\n",
    "- `retrieved`: list of retrieved document IDs (in ranked order)\n",
    "- `relevant`: set of relevant document IDs\n",
    "- `k`: cutoff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e08f3f30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.1: Precision@k and Recall@k\n",
    "\n",
    "def precision_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Precision at rank k.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "def recall_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Recall at rank k.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "# Sanity check with a toy example\n",
    "toy_retrieved = [\"d1\", \"d3\", \"d2\", \"d5\", \"d4\"]  # d1 and d2 are relevant\n",
    "toy_relevant = {\"d1\", \"d2\", \"d7\"}\n",
    "\n",
    "print(\"Toy example: retrieved =\", toy_retrieved, \", relevant =\", toy_relevant)\n",
    "for k in [1, 2, 3, 5]:\n",
    "    print(f\"  P@{k} = {precision_at_k(toy_retrieved, toy_relevant, k):.3f}, \"\n",
    "          f\"R@{k} = {recall_at_k(toy_retrieved, toy_relevant, k):.3f}\")\n",
    "# Expected: P@1=1.0, P@2=0.5, P@3=0.667, P@5=0.4\n",
    "#           R@1=0.333, R@2=0.333, R@3=0.667, R@5=0.667"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15beb69f",
   "metadata": {},
   "source": [
    "### 3.2 Mean Reciprocal Rank (MRR)\n",
    "\n",
    "**Reciprocal Rank** for a single query = $\\frac{1}{\\text{rank of first relevant document}}$\n",
    "\n",
    "$$\\text{MRR} = \\frac{1}{|Q|} \\sum_{q \\in Q} \\frac{1}{\\text{rank}_q}$$\n",
    "\n",
    "MRR focuses on **where the first relevant result appears**. It's the standard metric for question answering and navigational queries.\n",
    "\n",
    "**🚧 TODO 3.2:** Implement `reciprocal_rank` (for a single query) and `mean_reciprocal_rank` (averaged over a list of queries).\n",
    "\n",
    "Edge case: if no relevant document is retrieved, RR = 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcce253",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.2: MRR\n",
    "\n",
    "def reciprocal_rank(retrieved, relevant):\n",
    "    \"\"\"Reciprocal rank for a single query.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "def mean_reciprocal_rank(all_retrieved, all_relevant):\n",
    "    \"\"\"MRR over a list of queries.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "# Sanity check\n",
    "print(\"RR for toy example:\", reciprocal_rank(toy_retrieved, toy_relevant))\n",
    "# d1 is relevant and at rank 1 → RR = 1.0\n",
    "\n",
    "toy_ret2 = [\"d3\", \"d5\", \"d1\", \"d2\", \"d4\"]\n",
    "print(\"RR when first relevant at rank 3:\", reciprocal_rank(toy_ret2, toy_relevant))\n",
    "# d1 at rank 3 → RR = 1/3 ≈ 0.333"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd1f932",
   "metadata": {},
   "source": [
    "### 3.3 Average Precision (AP) and Mean Average Precision (MAP)\n",
    "\n",
    "**Average Precision** rewards systems that place relevant documents **early** in the ranking:\n",
    "\n",
    "$$\\text{AP}(q) = \\frac{1}{|\\text{Rel}_q|} \\sum_{k=1}^{n} P@k \\times \\text{rel}(k)$$\n",
    "\n",
    "Only positions where a relevant document appears contribute to the sum. AP is sensitive to the order of relevant documents.\n",
    "\n",
    "**MAP** = mean of AP across all queries.\n",
    "\n",
    "**🚧 TODO 3.3:** Implement `average_precision` and `mean_average_precision`.\n",
    "\n",
    "Then verify on this example:\n",
    "- retrieved = [✓, ✗, ✓, ✗, ✓], relevant = {d1, d3, d5} (3 total relevant)\n",
    "- AP = (1/3) × (P@1 × 1 + P@3 × 1 + P@5 × 1) = (1/3) × (1/1 + 2/3 + 3/5) = (1/3)(1 + 0.667 + 0.6) = **0.756**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48450815",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.3: MAP\n",
    "\n",
    "def average_precision(retrieved, relevant):\n",
    "    \"\"\"Average Precision for a single query.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "def mean_average_precision(all_retrieved, all_relevant):\n",
    "    \"\"\"MAP over a list of queries.\"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "# Sanity check: retrieved = [rel, non, rel, non, rel], 3 relevant total\n",
    "toy_ret3 = [\"d1\", \"d99\", \"d2\", \"d98\", \"d7\"]\n",
    "toy_rel3 = {\"d1\", \"d2\", \"d7\"}\n",
    "ap = average_precision(toy_ret3, toy_rel3)\n",
    "print(f\"AP = {ap:.4f}\")  # Expected: 0.7556\n",
    "assert abs(ap - 0.7556) < 0.001, f\"Expected ~0.756, got {ap}\"\n",
    "print(\"✓ AP check passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd0c278c",
   "metadata": {},
   "source": [
    "### 3.4 (Normalized) Discounted Cumulative Gain — nDCG@k\n",
    "\n",
    "nDCG handles **graded relevance** (not just binary) and gives a **logarithmic discount** for lower-ranked results.\n",
    "\n",
    "**DCG@k:**\n",
    "$$\\text{DCG}@k = \\sum_{i=1}^{k} \\frac{2^{\\text{rel}_i} - 1}{\\log_2(i + 1)}$$\n",
    "\n",
    "**Ideal DCG (IDCG@k):** DCG of the perfect ranking (sort documents by relevance, descending).\n",
    "\n",
    "**nDCG@k:**\n",
    "$$\\text{nDCG}@k = \\frac{\\text{DCG}@k}{\\text{IDCG}@k}$$\n",
    "\n",
    "nDCG is normalized to [0, 1]. A value of 1.0 means perfect ranking.\n",
    "\n",
    "**🚧 TODO 3.4:** Implement `dcg_at_k` and `ndcg_at_k`.\n",
    "\n",
    "For our MS MARCO gold file, we use **binary relevance** (relevant = 1, not relevant = 0). But your implementation should work with any integer relevance values.\n",
    "\n",
    "Verify on this toy example:\n",
    "- relevances = [3, 2, 0, 1] → DCG@4 = 7/1 + 3/1.585 + 0/2 + 1/2.322 = 7 + 1.893 + 0 + 0.431 = **9.324**\n",
    "- ideal = [3, 2, 1, 0] → IDCG@4 = 7/1 + 3/1.585 + 1/2 + 0/2.322 = 7 + 1.893 + 0.5 + 0 = **9.393**\n",
    "- nDCG@4 = 9.324 / 9.393 = **0.993**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9feedbca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.4: nDCG\n",
    "\n",
    "def dcg_at_k(relevances, k):\n",
    "    \"\"\"Discounted Cumulative Gain at rank k.\n",
    "\n",
    "    relevances: list of relevance scores in ranked order\n",
    "    \"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "def ndcg_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Normalized DCG at rank k.\n",
    "\n",
    "    retrieved: list of retrieved document IDs (in ranked order)\n",
    "    relevant: set of relevant document IDs (binary relevance)\n",
    "    \"\"\"\n",
    "    ... your code here ...\n",
    "\n",
    "\n",
    "# Sanity check with graded relevance\n",
    "graded_rels = [3, 2, 0, 1]\n",
    "dcg_val = dcg_at_k(graded_rels, 4)\n",
    "print(f\"DCG@4 = {dcg_val:.3f}\")  # Expected ~9.324\n",
    "\n",
    "ideal_rels = [3, 2, 1, 0]\n",
    "idcg_val = dcg_at_k(ideal_rels, 4)\n",
    "print(f\"IDCG@4 = {idcg_val:.3f}\")  # Expected ~9.393\n",
    "\n",
    "ndcg_val = dcg_val / idcg_val\n",
    "print(f\"nDCG@4 = {ndcg_val:.3f}\")  # Expected ~0.993\n",
    "\n",
    "# Binary relevance check\n",
    "print(f\"\\nnDCG for toy example: {ndcg_at_k(toy_retrieved, toy_relevant, 5):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04b59376",
   "metadata": {},
   "source": [
    "### 3.5 Conceptual Questions\n",
    "\n",
    "**🚧 TODO 3.5:** Answer these without running code. Think carefully.\n",
    "\n",
    "1. A system has **MRR = 1.0** and **MAP = 0.25**. Is this possible? What kind of system would produce this? Give a concrete example with 4 retrieved documents.\n",
    "\n",
    "2. System A has **P@10 = 0.3** and system B has **P@10 = 0.3**. Can their **nDCG@10** be different? Why?\n",
    "\n",
    "3. Your gold set has 100 relevant documents per query, but you only retrieve 10. What is the **maximum possible Recall@10**? What does this tell you about using Recall@10 as an evaluation metric here?\n",
    "\n",
    "4. You're building a search engine for a hospital's medical knowledge base. Doctors need to find the right treatment protocol for a patient. Which metric matters most: P@1, MRR, MAP, or nDCG? Justify your answer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeb81e47",
   "metadata": {},
   "source": [
    "**Answers 3.5:**\n",
    "\n",
    "1. ... your answer here ...\n",
    "\n",
    "2. ... your answer here ...\n",
    "\n",
    "3. ... your answer here ...\n",
    "\n",
    "4. ... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "798e56a6",
   "metadata": {},
   "source": [
    "## Part 4: Evaluating Our Baseline (~15 min)\n",
    "\n",
    "### 4.1 Evaluate Boolean Retrieval on All Training Queries\n",
    "\n",
    "**🚧 TODO 4.1:** Write a function `evaluate_system(vectorizer, doc_matrix, queries_df, gold_dict, k=10)` that:\n",
    "1. For each query, retrieves the top-k documents\n",
    "2. Computes P@k, MRR, MAP, and nDCG@k\n",
    "3. Returns a dictionary with the **mean** of each metric over all queries\n",
    "\n",
    "Apply it to the **boolean** retrieval system on the **training** queries.\n",
    "\n",
    "**Hint:** To go from document indices to document IDs, use `docs.iloc[idx][\"docid\"]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac4b982",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.1: Systematic evaluation\n",
    "\n",
    "def evaluate_system(vectorizer, doc_matrix, queries_df, gold_dict, k=10):\n",
    "    \"\"\"Evaluate a retrieval system on a set of queries.\"\"\"\n",
    "    ... your code here ...\n",
    "    return {\n",
    "        f\"P@{k}\": ...,\n",
    "        \"MRR\": ...,\n",
    "        \"MAP\": ...,\n",
    "        f\"nDCG@{k}\": ...,\n",
    "    }\n",
    "\n",
    "\n",
    "print(\"Evaluating Boolean retrieval on training queries...\")\n",
    "bool_results = evaluate_system(bool_vectorizer, bool_matrix, train_queries, gold_dict, k=10)\n",
    "print(\"\\nBoolean retrieval results:\")\n",
    "for metric, value in bool_results.items():\n",
    "    print(f\"  {metric}: {value:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3a2376",
   "metadata": {},
   "source": [
    "**🚧 TODO 4.2:** Now evaluate the **TF-IDF** system on the same training queries. Present both results side by side in a bar chart.\n",
    "\n",
    "Which metric improved the most? Which improved the least? Why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b0c5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.2: TF-IDF evaluation + comparison\n",
    "\n",
    "print(\"Evaluating TF-IDF retrieval on training queries...\")\n",
    "tfidf_results = evaluate_system(tfidf_vectorizer, tfidf_matrix, train_queries, gold_dict, k=10)\n",
    "print(\"\\nTF-IDF retrieval results:\")\n",
    "for metric, value in tfidf_results.items():\n",
    "    print(f\"  {metric}: {value:.4f}\")\n",
    "\n",
    "# Side-by-side bar chart\n",
    "... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bfb93a6",
   "metadata": {},
   "source": [
    "**Answer 4.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### 4.3 Error Mining\n",
    "\n",
    "**🚧 TODO 4.3:** Let's find out where TF-IDF fails.\n",
    "1. For each training query, compute its nDCG@10. Store the per-query scores.\n",
    "2. Find the **10 worst-scoring queries** (lowest nDCG@10).\n",
    "3. For the single worst query: print the query text, the top-5 retrieved documents, and 3 gold documents. What went wrong?\n",
    "4. Can you categorize the failure modes? (synonym mismatch, ambiguity, length bias, ...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86a48f90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.3: Error mining\n",
    "\n",
    "# 1. Per-query nDCG@10\n",
    "per_query_scores = []\n",
    "... your code here ...\n",
    "\n",
    "per_query_df = pd.DataFrame(per_query_scores).sort_values(\"ndcg10\")\n",
    "\n",
    "# 2. Worst 10 queries\n",
    "print(\"=== 10 Worst Queries by nDCG@10 ===\")\n",
    "print(per_query_df.head(10).to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca7e8d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Deep-dive into the worst query\n",
    "worst_qid = per_query_df.iloc[0][\"qid\"]\n",
    "worst_query = per_query_df.iloc[0][\"query\"]\n",
    "print(f\"Worst query: qid={worst_qid}, text={worst_query!r}\\n\")\n",
    "\n",
    "# What we retrieved\n",
    "... your code here ...\n",
    "\n",
    "# What the gold says\n",
    "... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc9c4027",
   "metadata": {},
   "source": [
    "**Answer 4.3:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6823efa9",
   "metadata": {},
   "source": [
    "## Part 5: Looking Ahead\n",
    "\n",
    "### 5.1 Final Evaluation on Test Queries\n",
    "\n",
    "**🚧 TODO 5.1:** Evaluate your TF-IDF system on the **test** queries (the ones we haven't used yet). Is performance similar to the training set? What does this tell you about the generalizability of keyword-based retrieval?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfccc334",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 5.1: Test set evaluation\n",
    "\n",
    "print(\"Evaluating TF-IDF on test queries...\")\n",
    "test_results = evaluate_system(tfidf_vectorizer, tfidf_matrix, test_queries, gold_dict, k=10)\n",
    "print(\"\\nTest set results:\")\n",
    "for metric, value in test_results.items():\n",
    "    print(f\"  {metric}: {value:.4f}\")\n",
    "\n",
    "print(\"\\nComparison (TF-IDF):\")\n",
    "... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e994b7ef",
   "metadata": {},
   "source": [
    "### 5.2 Discussion: What Comes Next?\n",
    "\n",
    "**🚧 TODO 5.2:** Based on your error analysis, propose **three concrete ideas** to improve retrieval quality. For each idea, explain:\n",
    "- What failure mode it addresses\n",
    "- How you would implement it (briefly)\n",
    "- Any potential downsides\n",
    "\n",
    "**Answer:**\n",
    "\n",
    "... your answer here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25fb6b65",
   "metadata": {},
   "source": [
    "### 5.3 The Codabench Challenge (Preview)\n",
    "\n",
    "Starting in **Class 3**, you will participate in a retrieval challenge on [codabench.org](https://www.codabench.org/):\n",
    "\n",
    "- **Task:** Scientific article retrieval (a different domain from MS MARCO)\n",
    "- **Goal:** Given queries about scientific topics, rank the most relevant articles\n",
    "- **Primary metric:** nDCG@10 — the same metric you just implemented!\n",
    "- **Groups:** 2–3 students\n",
    "- **Submission format:** TSV with columns `qid`, `docid`, `rank`, `score`\n",
    "- We use MS MARCO in the notebooks to learn techniques; the challenge tests if you can **transfer** them\n",
    "\n",
    "The techniques that improve your MS MARCO results will likely improve your challenge results too. Every week you'll learn tools that directly apply.\n",
    "\n",
    "---\n",
    "\n",
    "*Next week: **BM25** — the most successful retrieval algorithm ever built, used by Elasticsearch, Solr, and Lucene. You'll have a prep notebook to complete before class. We'll also learn about stemming, stopwords, and how BM25 tunes its behavior via two parameters.*\n",
    "\n",
    "**Make sure to submit this notebook before leaving!** 🎉"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}