{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b082b2f6",
   "metadata": {},
   "source": [
    "# Information Retrieval — Class 2 Prep: BM25 & Sparse Retrieval\n",
    "\n",
    "**Due:** March 20, 2026 (before class)\n",
    "\n",
    "**Goal:** Understand BM25 in depth — theory, implementation from scratch, parameter tuning, and preprocessing for retrieval. This notebook will prepare you for the in-class session where you'll implement BM25 variants and build the strongest possible sparse baseline.\n",
    "\n",
    "## Notebook Outline\n",
    "\n",
    "| Part | Topic | Description |\n",
    "|------|-------|-------------|\n",
    "| **1** | Setup & Data | Load MS MARCO + evaluation toolkit from Class 1 |\n",
    "| **2** | From TF-IDF to BM25: Theory | Mathematical derivation, parameter intuitions |\n",
    "| **3** | BM25 from Scratch | Implement BM25 scoring in NumPy — no libraries |\n",
    "| **4** | Using `rank_bm25` | Compare your implementation with the standard library |\n",
    "| **5** | Preprocessing for Retrieval | Stemming, stopwords, lowercasing — what helps? |\n",
    "| **6** | Parameter Tuning | Grid search over $k_1$, $b$; heatmap visualization |\n",
    "| **7** | Error Analysis & Discussion | Where BM25 fails, and why |\n",
    "\n",
    "**Key Concepts:**\n",
    "- BM25 algorithm (the world's most deployed retrieval model)\n",
    "- TF saturation and document length normalization\n",
    "- Parameter sensitivity ($k_1$ and $b$)\n",
    "- Preprocessing trade-offs for sparse retrieval\n",
    "- Systematic hyperparameter search with visualization\n",
    "\n",
    "**Required Libraries:** `rank_bm25`, `nltk`, `scikit-learn`, `numpy`, `pandas`, `matplotlib`, `seaborn`\n",
    "\n",
    "**Essential Reading & Resources:**\n",
    "\n",
    "| Resource | Link |\n",
    "|----------|------|\n",
    "| Robertson & Zaragoza (2009), *\"The Probabilistic Relevance Framework: BM25 and Beyond\"* | [PDF](https://www.staff.city.ac.uk/~sbrp622/papers/foundations_bm25_review.pdf) |\n",
    "| Manning et al., *Introduction to Information Retrieval*, Ch. 11 (Probabilistic IR) | [Free online](https://nlp.stanford.edu/IR-book/html/htmledition/probabilistic-information-retrieval-1.html) |\n",
    "| Wikipedia: Okapi BM25 | [Link](https://en.wikipedia.org/wiki/Okapi_BM25) |\n",
    "| `rank_bm25` Python library | [GitHub](https://github.com/dorianbrown/rank_bm25) |\n",
    "| scikit-learn `TfidfVectorizer` documentation | [Docs](https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.text.TfidfVectorizer.html) |\n",
    "| NLTK Stemming HOWTO | [Docs](https://www.nltk.org/howto/stem.html) |\n",
    "| Video: BM25 explained in 10 minutes (Pinecone) | [YouTube](https://www.youtube.com/watch?v=rOOJCih4Z-0) |\n",
    "| Blog: The BM25 Algorithm and its Variables (kmwllc) | [Link](https://kmwllc.com/index.php/2020/03/20/understanding-tf-idf-and-bm-25/) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa1ff16",
   "metadata": {},
   "source": [
    "<span style=\"color:magenta\">Student Name:</span>\n",
    "\n",
    "* Name\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e43d309",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "import urllib.request\n",
    "import math\n",
    "import pickle\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": "be900117",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1: Data Loading & Evaluation Toolkit\n",
    "\n",
    "We reuse the MS MARCO dataset and the evaluation functions from Class 1. Make sure your Class 1 notebook is complete and working — we'll need those metrics.\n",
    "\n",
    "> 📌 **Tip:** If you haven't already, install required libraries:\n",
    "> ```bash\n",
    "> pip install rank_bm25 nltk\n",
    "> ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8d54eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download MS MARCO data\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": "code",
   "execution_count": null,
   "id": "2d7cbaaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all three data files\n",
    "queries = pd.read_csv(\"our.msmarco.queries.tsv\", sep=\"\\t\", usecols=[1, 2])\n",
    "docs_full = pd.read_csv(\"our.msmarco.docs.tsv\", sep=\"\\t\", usecols=[1, 2, 3])\n",
    "gold = pd.read_csv(\"our.msmarco.gold.tsv\", sep=\"\\t\", usecols=[1, 3])\n",
    "\n",
    "# ── Optional: reduce dataset for faster iteration ──────────────────\n",
    "# Set FAST_MODE = True for quick runs (~100 train + 100 test queries, ~20k docs)\n",
    "# Set FAST_MODE = False for full dataset (500 train + 500 test, ~92k docs)\n",
    "FAST_MODE = True\n",
    "\n",
    "if FAST_MODE:\n",
    "    train_queries = queries.iloc[:100].copy()\n",
    "    test_queries = queries.iloc[500:600].copy()\n",
    "else:\n",
    "    train_queries = queries.iloc[:500].copy()\n",
    "    test_queries = queries.iloc[500:].copy()\n",
    "\n",
    "all_qids = set(train_queries[\"qid\"]) | set(test_queries[\"qid\"])\n",
    "gold = gold[gold[\"qid\"].isin(all_qids)]\n",
    "gold_dict = gold.groupby(\"qid\")[\"docid\"].apply(set).to_dict()\n",
    "\n",
    "# Keep gold docs + random distractors\n",
    "gold_docids = set(gold[\"docid\"])\n",
    "gold_mask = docs_full[\"docid\"].isin(gold_docids)\n",
    "gold_docs = docs_full[gold_mask]\n",
    "other_docs = docs_full[~gold_mask].sample(n=min(10000, (~gold_mask).sum()), random_state=42)\n",
    "docs = pd.concat([gold_docs, other_docs]).reset_index(drop=True)\n",
    "docs[\"text\"] = docs[\"title\"].fillna(\"\") + \" \" + docs[\"body\"].fillna(\"\")\n",
    "\n",
    "mode = \"FAST\" if FAST_MODE else \"FULL\"\n",
    "print(f\"[{mode}] Queries: {len(train_queries)}+{len(test_queries)}, Docs: {len(docs)} ({len(gold_docs)} gold + {len(other_docs)} distractor)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f38940a0",
   "metadata": {},
   "source": [
    "### Evaluation Functions from Class 1\n",
    "\n",
    "We paste the evaluation functions here so this notebook is self-contained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2db435",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Evaluation metrics from Class 1 ---\n",
    "\n",
    "def precision_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Precision at rank k.\"\"\"\n",
    "    retrieved_at_k = retrieved[:k]\n",
    "    n_relevant_at_k = len(set(retrieved_at_k) & relevant)\n",
    "    return n_relevant_at_k / k\n",
    "\n",
    "def recall_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Recall at rank k.\"\"\"\n",
    "    if len(relevant) == 0:\n",
    "        return 0.0\n",
    "    retrieved_at_k = retrieved[:k]\n",
    "    n_relevant_at_k = len(set(retrieved_at_k) & relevant)\n",
    "    return n_relevant_at_k / len(relevant)\n",
    "\n",
    "def reciprocal_rank(retrieved, relevant):\n",
    "    \"\"\"Reciprocal rank for a single query.\"\"\"\n",
    "    for rank, doc_id in enumerate(retrieved, start=1):\n",
    "        if doc_id in relevant:\n",
    "            return 1.0 / rank\n",
    "    return 0.0\n",
    "\n",
    "def average_precision(retrieved, relevant):\n",
    "    \"\"\"Average Precision for a single query.\"\"\"\n",
    "    if len(relevant) == 0:\n",
    "        return 0.0\n",
    "    ap_sum = 0.0\n",
    "    n_relevant_so_far = 0\n",
    "    for k, doc_id in enumerate(retrieved, start=1):\n",
    "        if doc_id in relevant:\n",
    "            n_relevant_so_far += 1\n",
    "            ap_sum += n_relevant_so_far / k\n",
    "    return ap_sum / len(relevant)\n",
    "\n",
    "def dcg_at_k(relevances, k):\n",
    "    \"\"\"Discounted Cumulative Gain at rank k.\"\"\"\n",
    "    relevances = np.array(relevances[:k], dtype=float)\n",
    "    positions = np.arange(1, len(relevances) + 1)\n",
    "    discounts = np.log2(positions + 1)\n",
    "    gains = (2 ** relevances - 1) / discounts\n",
    "    return float(np.sum(gains))\n",
    "\n",
    "def ndcg_at_k(retrieved, relevant, k):\n",
    "    \"\"\"Normalized DCG at rank k (binary relevance).\"\"\"\n",
    "    relevances = [1 if doc_id in relevant else 0 for doc_id in retrieved[:k]]\n",
    "    dcg = dcg_at_k(relevances, k)\n",
    "    ideal_relevances = sorted(relevances, reverse=True)\n",
    "    idcg = dcg_at_k(ideal_relevances, k)\n",
    "    if idcg == 0:\n",
    "        return 0.0\n",
    "    return dcg / idcg\n",
    "\n",
    "print(\"✓ Evaluation functions loaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c87cd475",
   "metadata": {},
   "source": [
    "### Generic Evaluation Helper\n",
    "\n",
    "This function evaluates any retrieval function that returns `(top_k_indices, scores)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4800cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_retriever(retrieve_fn, queries_df, gold_dict, docs_df, k=10):\n",
    "    \"\"\"Evaluate a retrieval function on a set of queries.\n",
    "\n",
    "    retrieve_fn: function(query_text) -> (top_k_doc_indices, scores)\n",
    "    \"\"\"\n",
    "    p_scores, mrr_scores, map_scores, ndcg_scores = [], [], [], []\n",
    "\n",
    "    for _, row in tqdm(queries_df.iterrows(), total=len(queries_df), desc=\"Evaluating\"):\n",
    "        qid, query_text = row[\"qid\"], row[\"query\"]\n",
    "        relevant = gold_dict.get(qid, set())\n",
    "        if len(relevant) == 0:\n",
    "            continue\n",
    "\n",
    "        top_k_idx, _ = retrieve_fn(query_text)\n",
    "        retrieved_ids = [docs_df.iloc[i][\"docid\"] for i in top_k_idx]\n",
    "\n",
    "        p_scores.append(precision_at_k(retrieved_ids, relevant, k))\n",
    "        mrr_scores.append(reciprocal_rank(retrieved_ids, relevant))\n",
    "        map_scores.append(average_precision(retrieved_ids, relevant))\n",
    "        ndcg_scores.append(ndcg_at_k(retrieved_ids, relevant, k))\n",
    "\n",
    "    return {\n",
    "        f\"P@{k}\": np.mean(p_scores),\n",
    "        \"MRR\": np.mean(mrr_scores),\n",
    "        \"MAP\": np.mean(map_scores),\n",
    "        f\"nDCG@{k}\": np.mean(ndcg_scores),\n",
    "    }\n",
    "\n",
    "\n",
    "def print_results(name, results):\n",
    "    \"\"\"Pretty-print evaluation results.\"\"\"\n",
    "    print(f\"\\n{'='*40}\")\n",
    "    print(f\"  {name}\")\n",
    "    print(f\"{'='*40}\")\n",
    "    for metric, value in results.items():\n",
    "        print(f\"  {metric:<12} {value:.4f}\")\n",
    "\n",
    "print(\"✓ Evaluation helper loaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9182e704",
   "metadata": {},
   "source": [
    "### Quick TF-IDF Baseline (from Class 1)\n",
    "\n",
    "Let's first reproduce our TF-IDF baseline so we have a reference point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "618ff24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build TF-IDF baseline\n",
    "tfidf_vectorizer = TfidfVectorizer(sublinear_tf=True, strip_accents=\"unicode\")\n",
    "tfidf_matrix = tfidf_vectorizer.fit_transform(docs[\"text\"])\n",
    "\n",
    "def tfidf_retrieve(query_text, k=10):\n",
    "    query_vec = tfidf_vectorizer.transform([query_text])\n",
    "    scores = (tfidf_matrix @ query_vec.T).toarray().flatten()\n",
    "    if k >= len(scores):\n",
    "        top_k = np.argsort(scores)[::-1][:k]\n",
    "    else:\n",
    "        top_k = np.argpartition(scores, -k)[-k:]\n",
    "        top_k = top_k[np.argsort(scores[top_k])[::-1]]\n",
    "    return top_k, scores[top_k]\n",
    "\n",
    "tfidf_results = evaluate_retriever(tfidf_retrieve, train_queries, gold_dict, docs)\n",
    "print_results(\"TF-IDF Baseline (Class 1)\", tfidf_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12b7b9e4",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 2: From TF-IDF to BM25 — Theory\n",
    "\n",
    "### 2.1 Why TF-IDF Falls Short for Retrieval\n",
    "\n",
    "TF-IDF was designed for document representation (classification, clustering), not retrieval. It has two critical problems when used as a ranking function:\n",
    "\n",
    "**Problem 1: Linear TF growth.**  \n",
    "If a document mentions \"python\" 100 times, its TF-IDF score for \"python\" is 100× higher than a doc mentioning it once (with raw tf) or $1 + \\log(100) = 3\\times$ higher (with sublinear tf). Is it really 100× or even 3× more relevant? Probably not.\n",
    "\n",
    "**Problem 2: No explicit document length normalization.**  \n",
    "`TfidfVectorizer` applies L2 normalization to the whole document vector, which implicitly normalizes by length. But this normalization isn't tunable — you can't control _how much_ length matters.\n",
    "\n",
    "### 2.2 The BM25 Formula\n",
    "\n",
    "BM25 fixes both problems with two parameters:\n",
    "\n",
    "$$\\text{BM25}(q, d) = \\sum_{t \\in q} \\text{IDF}(t) \\cdot \\frac{\\text{tf}(t, d) \\cdot (k_1 + 1)}{\\text{tf}(t, d) + k_1 \\cdot \\left(1 - b + b \\cdot \\frac{|d|}{\\text{avgdl}}\\right)}$$\n",
    "\n",
    "where:\n",
    "- $k_1$ controls **TF saturation** (typical: 1.2). Higher = more weight to raw frequency.\n",
    "- $b$ controls **length normalization** (typical: 0.75). Higher = more penalty for long docs.\n",
    "- $|d|$ = document length in tokens, $\\text{avgdl}$ = average document length in the corpus\n",
    "\n",
    "BM25's IDF is slightly different from TF-IDF's:\n",
    "\n",
    "$$\\text{IDF}(t) = \\log\\frac{N - \\text{df}(t) + 0.5}{\\text{df}(t) + 0.5}$$\n",
    "\n",
    "> 📖 **Read:** The [Wikipedia article on BM25](https://en.wikipedia.org/wiki/Okapi_BM25) has an excellent detailed breakdown. Also see [Robertson & Zaragoza (2009)](https://www.staff.city.ac.uk/~sbrp622/papers/foundations_bm25_review.pdf), Sections 1–3.\n",
    "\n",
    "### 🚧 TODO 2.1: Understanding the parameters\n",
    "\n",
    "Before implementing BM25, make sure you understand how the parameters work. Answer these questions:\n",
    "\n",
    "1. What happens to the BM25 score when $k_1 = 0$? What model does this reduce to?\n",
    "2. What happens when $k_1 \\to \\infty$?\n",
    "3. What happens when $b = 0$? When $b = 1$?\n",
    "4. For a document whose length equals the average ($|d| = \\text{avgdl}$), does $b$ matter? Why?\n",
    "5. A document is twice the average length. A query term appears 5 times in it. Compare the BM25 score contribution of this term for $b = 0$ vs $b = 1$ (use $k_1 = 1.2$). Show your calculation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4727e686",
   "metadata": {},
   "source": [
    "**Answers 2.1:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4934fcfd",
   "metadata": {},
   "source": [
    "### 2.3 Visualizing TF Saturation\n",
    "\n",
    "The following plots help you understand how BM25 differs from TF-IDF visually. Study the curves before implementing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4490ae99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize TF saturation for different k1 values\n",
    "\n",
    "tf_values = np.arange(0, 51)\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "# Left: BM25 TF component for different k1\n",
    "ax = axes[0]\n",
    "for k1 in [0.5, 1.2, 2.0, 5.0, 100.0]:\n",
    "    bm25_tf = tf_values * (k1 + 1) / (tf_values + k1)\n",
    "    ax.plot(tf_values, bm25_tf, label=f\"$k_1$={k1}\")\n",
    "ax.plot(tf_values, tf_values, \"--\", color=\"gray\", alpha=0.5, label=\"raw tf\")\n",
    "ax.set_xlabel(\"Term Frequency (tf)\")\n",
    "ax.set_ylabel(\"BM25 TF component\")\n",
    "ax.set_title(\"Effect of $k_1$ on TF Saturation\")\n",
    "ax.legend()\n",
    "ax.set_xlim(0, 50)\n",
    "ax.set_ylim(0, 12)\n",
    "\n",
    "# Right: Length normalization factor for different b\n",
    "ax = axes[1]\n",
    "doc_lengths = np.linspace(0.1, 4, 100)  # relative to avgdl\n",
    "for b in [0, 0.25, 0.5, 0.75, 1.0]:\n",
    "    length_factor = 1 - b + b * doc_lengths\n",
    "    ax.plot(doc_lengths, length_factor, label=f\"$b$={b}\")\n",
    "ax.axvline(1.0, color=\"gray\", linestyle=\":\", alpha=0.5, label=\"avgdl\")\n",
    "ax.set_xlabel(\"|d| / avgdl (relative document length)\")\n",
    "ax.set_ylabel(\"Length normalization factor\")\n",
    "ax.set_title(\"Effect of $b$ on Length Normalization\")\n",
    "ax.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3d48924",
   "metadata": {},
   "source": [
    "**🚧 TODO 2.2:** Look at the two plots above and answer:\n",
    "1. At what TF value does the BM25 curve for $k_1=1.2$ reach ~90% of its maximum?\n",
    "2. If a document is 3× the average length and $b=0.75$, what is the length normalization factor? How does this affect the effective \"TF cost\" of matching a term?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dacd76fb",
   "metadata": {},
   "source": [
    "**Answer 2.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "015a6de9",
   "metadata": {},
   "source": [
    "## Part 3: BM25 from Scratch (~30 min)\n",
    "\n",
    "Now implement BM25 yourself, using only NumPy and scikit-learn's `CountVectorizer` for tokenization.\n",
    "\n",
    "> 📖 **Reference:** [Manning et al., Chapter 11.4](https://nlp.stanford.edu/IR-book/html/htmledition/okapi-bm25-a-non-binary-model-1.html) — exact formulas with examples.\n",
    "\n",
    "### 🚧 TODO 3.1: Implement BM25 step by step\n",
    "\n",
    "The plan:\n",
    "1. Use `CountVectorizer` (not binary, not TF-IDF) to get raw term frequencies\n",
    "2. Compute document lengths and average document length\n",
    "3. Compute IDF for each term using the BM25 formula\n",
    "4. Write a `bm25_score(query_text, k1, b)` function that returns scores for all documents\n",
    "5. Write `bm25_retrieve(query_text, k=10)` that returns top-k doc indices + scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff36378c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.1: BM25 from scratch\n",
    "\n",
    "# Step 1: Raw term frequency matrix (NOT binary, NOT TF-IDF)\n",
    "count_vectorizer = CountVectorizer()\n",
    "tf_matrix = count_vectorizer.fit_transform(docs[\"text\"])  # sparse matrix, shape (n_docs, vocab_size)\n",
    "vocab = count_vectorizer.get_feature_names_out()\n",
    "print(f\"Vocabulary size: {len(vocab)}\")\n",
    "print(f\"TF matrix shape: {tf_matrix.shape}, nnz: {tf_matrix.nnz}\")\n",
    "\n",
    "# Step 2: Document lengths (number of terms per document)\n",
    "doc_lengths = tf_matrix.sum(axis=1).A1  # .A1 converts matrix to flat array\n",
    "avgdl = doc_lengths.mean()\n",
    "print(f\"Average document length: {avgdl:.1f} tokens\")\n",
    "\n",
    "# Step 3: Document frequencies and BM25 IDF\n",
    "N = tf_matrix.shape[0]\n",
    "df_array = (tf_matrix > 0).sum(axis=0).A1  # number of docs containing each term\n",
    "# BM25 IDF formula\n",
    "idf = np.log((N - df_array + 0.5) / (df_array + 0.5))\n",
    "# Clip negative IDFs (very common terms)\n",
    "idf = np.maximum(idf, 0)\n",
    "print(f\"IDF range: [{idf.min():.3f}, {idf.max():.3f}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0effb246",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 4 & 5: BM25 scoring and retrieval function\n",
    "\n",
    "def bm25_scores_all(query_text, k1=1.2, b=0.75):\n",
    "    \"\"\"Compute BM25 scores for all documents given a query.\n",
    "\n",
    "    Returns a 1D array of shape (n_docs,).\n",
    "    \"\"\"\n",
    "    # ... your code here ...\n",
    "\n",
    "\n",
    "def bm25_retrieve(query_text, k=10, k1=1.2, b=0.75):\n",
    "    \"\"\"Return top-k document indices and scores using our BM25.\"\"\"\n",
    "    # ... your code here ...\n",
    "\n",
    "\n",
    "# Quick test\n",
    "top_idx, top_scores = bm25_retrieve(\"what is a po box\", k=5)\n",
    "for rank, (idx, score) in enumerate(zip(top_idx, top_scores), 1):\n",
    "    print(f\"  #{rank} [score={score:.3f}] {docs.iloc[idx]['docid']}: {docs.iloc[idx]['text'][:120]}...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e23041c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate our BM25 from scratch\n",
    "bm25_scratch_results = evaluate_retriever(\n",
    "    lambda q: bm25_retrieve(q, k=10, k1=1.2, b=0.75),\n",
    "    train_queries, gold_dict, docs\n",
    ")\n",
    "print_results(\"BM25 from scratch (k1=1.2, b=0.75)\", bm25_scratch_results)\n",
    "print_results(\"TF-IDF baseline\", tfidf_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55b71697",
   "metadata": {},
   "source": [
    "**🚧 TODO 3.2:** Compare BM25 to TF-IDF on the training set. Answer:\n",
    "1. Which metrics improved the most?\n",
    "2. Why does BM25 handle long documents better than TF-IDF?\n",
    "3. Look at a specific query where BM25 succeeds but TF-IDF fails. Show the top-3 results from each system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462fab3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.2: Side-by-side comparison on specific queries\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eca04513",
   "metadata": {},
   "source": [
    "**Answer 3.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61356dc1",
   "metadata": {},
   "source": [
    "## Part 4: Using the `rank_bm25` Library\n",
    "\n",
    "The [`rank_bm25`](https://github.com/dorianbrown/rank_bm25) library provides a clean implementation of BM25 (and variants). Let's verify our implementation matches.\n",
    "\n",
    "> 📖 See the [rank_bm25 README](https://github.com/dorianbrown/rank_bm25) for usage examples.\n",
    "\n",
    "### 🚧 TODO 4.1: Compare with `rank_bm25`\n",
    "\n",
    "1. Tokenize all documents using the `CountVectorizer`'s built-in analyzer\n",
    "2. Build a `BM25Okapi` index\n",
    "3. Score the same queries and compare to our from-scratch implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "215c2058",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rank_bm25 import BM25Okapi\n",
    "\n",
    "# 🚧 TODO 4.1: Tokenize and build BM25Okapi index\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8798db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare scores on a few queries\n",
    "\n",
    "def rank_bm25_retrieve(query_text, k=10):\n",
    "    \"\"\"Retrieve using rank_bm25 library.\"\"\"\n",
    "    # ... your code here ...\n",
    "\n",
    "# Compare on a test query\n",
    "test_q = \"how long does it take to get to mars\"\n",
    "print(f\"Query: {test_q!r}\\n\")\n",
    "\n",
    "# ... your code here: compare the two implementations side by side ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb168f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Full evaluation of rank_bm25\n",
    "bm25_lib_results = evaluate_retriever(rank_bm25_retrieve, train_queries, gold_dict, docs)\n",
    "print_results(\"rank_bm25 (library)\", bm25_lib_results)\n",
    "print_results(\"BM25 from scratch\", bm25_scratch_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a211f22",
   "metadata": {},
   "source": [
    "**🚧 TODO 4.2:** Do the two implementations give exactly the same results? If there are small differences, why might that be?\n",
    "\n",
    "**Answer 4.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b897621",
   "metadata": {},
   "source": [
    "## Part 5: Preprocessing for Sparse Retrieval (~25 min)\n",
    "\n",
    "Stemming, stopword removal, and lowercasing can significantly affect BM25 performance. Let's experiment systematically.\n",
    "\n",
    "> 📖 **Resources:**\n",
    "> - [NLTK Stemming HOWTO](https://www.nltk.org/howto/stem.html) — Porter, Lancaster, Snowball stemmers\n",
    "> - [Martin Porter's page](https://tartarus.org/martin/PorterStemmer/) — original stemmer description\n",
    "> - [Wikipedia: Stemming](https://en.wikipedia.org/wiki/Stemming) — overview of algorithms\n",
    "> - [Blog: Lemmatization vs Stemming](https://nlp.stanford.edu/IR-book/html/htmledition/stemming-and-lemmatization-1.html) — Manning et al. textbook chapter\n",
    "\n",
    "### 🚧 TODO 5.1: Explore Different Stemmers\n",
    "\n",
    "NLTK provides three main stemmers. Explore them on a few examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "251ff628",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "nltk.download(\"punkt\", quiet=True)\n",
    "nltk.download(\"punkt_tab\", quiet=True)\n",
    "nltk.download(\"stopwords\", quiet=True)\n",
    "\n",
    "from nltk.stem import PorterStemmer, LancasterStemmer, SnowballStemmer\n",
    "from nltk.corpus import stopwords\n",
    "\n",
    "# 🚧 TODO 5.1: Compare stemmers on a set of test words\n",
    "\n",
    "porter = PorterStemmer()\n",
    "lancaster = LancasterStemmer()\n",
    "snowball = SnowballStemmer(\"english\")\n",
    "\n",
    "test_words = [\n",
    "    \"running\", \"organization\", \"university\", \"universe\", \"policy\", \"police\",\n",
    "    \"argument\", \"arguing\", \"retrieval\", \"retrieved\", \"information\", \"informational\",\n",
    "    \"better\", \"absorption\", \"absorb\", \"flies\", \"flying\", \"scientific\", \"science\"\n",
    "]\n",
    "\n",
    "# ... your code here: print a comparison table ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d347bff",
   "metadata": {},
   "source": [
    "**🚧 TODO 5.2:** Look at the stemmer outputs above and answer:\n",
    "1. Which stemmer is the most aggressive? Give an example where it over-stems.\n",
    "2. Find a pair of words that *should* have the same stem but don't (for Porter).\n",
    "3. Find a pair of words that have the same stem but *shouldn't* (for Porter).\n",
    "4. For the challenge (scientific articles), would you expect stemming to help or hurt? Why?\n",
    "\n",
    "**Answer 5.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### 🚧 TODO 5.3: Build BM25 with Preprocessing\n",
    "\n",
    "Now let's measure the actual impact. Build a pipeline that:\n",
    "1. Optional: removes stopwords using NLTK's stopword list\n",
    "2. Optional: applies a stemmer (Porter, Lancaster, or Snowball)\n",
    "3. Builds a new `CountVectorizer` with the preprocessed text\n",
    "4. Runs BM25 and evaluates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef5c74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 5.3: BM25 with preprocessing pipeline\n",
    "\n",
    "english_stopwords = set(stopwords.words(\"english\"))\n",
    "\n",
    "def make_preprocessor(use_stemmer=None, remove_stopwords=False):\n",
    "    \"\"\"Create a text preprocessing function.\n",
    "\n",
    "    use_stemmer: None, 'porter', 'lancaster', or 'snowball'\n",
    "    remove_stopwords: bool\n",
    "    Returns: function(text) -> list of tokens\n",
    "    \"\"\"\n",
    "    # ... your code here ...\n",
    "\n",
    "\n",
    "def build_bm25_with_preprocessing(use_stemmer=None, remove_stopwords=False, k1=1.2, b=0.75):\n",
    "    \"\"\"Build a BM25 system with specified preprocessing.\n",
    "\n",
    "    Returns: retrieve_fn(query_text, k=10) -> (top_k_indices, scores)\n",
    "    \"\"\"\n",
    "    # ... your code here ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d7f77c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test a few preprocessing combinations\n",
    "configs = [\n",
    "    {\"use_stemmer\": None,      \"remove_stopwords\": False, \"label\": \"No preprocessing\"},\n",
    "    {\"use_stemmer\": None,      \"remove_stopwords\": True,  \"label\": \"+ Stopwords removed\"},\n",
    "    {\"use_stemmer\": \"porter\",  \"remove_stopwords\": False, \"label\": \"+ Porter stemming\"},\n",
    "    {\"use_stemmer\": \"porter\",  \"remove_stopwords\": True,  \"label\": \"+ Porter + Stopwords\"},\n",
    "    {\"use_stemmer\": \"snowball\",\"remove_stopwords\": True,  \"label\": \"+ Snowball + Stopwords\"},\n",
    "    {\"use_stemmer\": \"lancaster\",\"remove_stopwords\": True, \"label\": \"+ Lancaster + Stopwords\"},\n",
    "]\n",
    "\n",
    "preprocessing_results = {}\n",
    "\n",
    "for cfg in configs:\n",
    "    label = cfg.pop(\"label\")\n",
    "    retrieve_fn = build_bm25_with_preprocessing(**cfg, k1=1.2, b=0.75)\n",
    "    results = evaluate_retriever(retrieve_fn, train_queries, gold_dict, docs)\n",
    "    preprocessing_results[label] = results\n",
    "    print_results(label, results)\n",
    "    cfg[\"label\"] = label  # restore for later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6a3103",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize preprocessing results\n",
    "prep_df = pd.DataFrame(preprocessing_results).T\n",
    "prep_df.plot(kind=\"bar\", figsize=(14, 5), rot=25)\n",
    "plt.title(\"BM25 with Different Preprocessing Strategies (k1=1.2, b=0.75)\")\n",
    "plt.ylabel(\"Score\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c28c5016",
   "metadata": {},
   "source": [
    "**🚧 TODO 5.4:** Based on the results, answer:\n",
    "1. Which preprocessing combination works best for MS MARCO?\n",
    "2. Does stopword removal help? By how much?\n",
    "3. Does stemming help? Which stemmer works best?\n",
    "4. Lancaster stemming is the most aggressive — did it perform best? If not, why?\n",
    "5. Pick the best preprocessing configuration and note it down. We'll use it for parameter tuning.\n",
    "\n",
    "**Answer 5.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f689addc",
   "metadata": {},
   "source": [
    "## Part 6: Parameter Tuning — Grid Search (~20 min)\n",
    "\n",
    "BM25 has two primary parameters: $k_1$ (TF saturation) and $b$ (length normalization). Finding optimal values is crucial and dataset-dependent.\n",
    "\n",
    "> 📖 **Video:** [BM25 tuning in practice (Pinecone)](https://www.youtube.com/watch?v=rOOJCih4Z-0) — 10 min overview including parameter effects\n",
    "\n",
    "### 🚧 TODO 6.1: Grid Search over $k_1$ and $b$\n",
    "\n",
    "Use the best preprocessing from Part 5. Evaluate BM25 on a grid of $k_1 \\in \\{0.5, 0.8, 1.0, 1.2, 1.5, 2.0, 3.0\\}$ and $b \\in \\{0.0, 0.25, 0.5, 0.65, 0.75, 0.85, 1.0\\}$.\n",
    "\n",
    "Visualize the results as a **heatmap**.\n",
    "\n",
    "> ⚠️ This takes a while because we're building a new BM25 index for each $(k_1, b)$ pair. Be patient — use only the training queries, not the full set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "065df182",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 6.1: Grid search\n",
    "\n",
    "# Use the best preprocessing (from Part 5)\n",
    "best_stemmer = \"porter\"  # adjust based on your Part 5 results\n",
    "best_stopwords = True\n",
    "preprocessor = make_preprocessor(use_stemmer=best_stemmer, remove_stopwords=best_stopwords)\n",
    "\n",
    "print(\"Tokenizing documents with best preprocessing...\")\n",
    "tokenized_docs_best = docs[\"text\"].progress_apply(preprocessor).tolist()\n",
    "\n",
    "# Parameter grid\n",
    "k1_values = [0.5, 0.8, 1.0, 1.2, 1.5, 2.0, 3.0]\n",
    "b_values = [0.0, 0.25, 0.5, 0.65, 0.75, 0.85, 1.0]\n",
    "\n",
    "# ... your code here: evaluate each (k1, b) pair and store nDCG@10 in a grid ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fcf523a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 6.1 (continued): Heatmap visualization\n",
    "\n",
    "# ... your code here: plot the heatmap ...\n",
    "\n",
    "# Print best configuration\n",
    "# best_idx = ...\n",
    "# print(f\"Best: k1=..., b=... → nDCG@10=...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1836a963",
   "metadata": {},
   "source": [
    "**🚧 TODO 6.2:** Analyze the heatmap and answer:\n",
    "1. Is the optimal $b$ close to the standard value of 0.75? What does this say about our MS MARCO documents?\n",
    "2. Is the heatmap \"smooth\" or are there sharp transitions? What does this imply about the robustness of BM25 to parameter choices?\n",
    "3. What happens at $b = 0$ (the leftmost column of the heatmap)? Why?\n",
    "4. Record your best $(k_1, b)$ here: ____ . You'll need it for the in-class session.\n",
    "\n",
    "**Answer 6.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b9c336b",
   "metadata": {},
   "source": [
    "### 🚧 TODO 6.3: Fine-grained search around the optimum\n",
    "\n",
    "Now that you know the rough region, do a finer grid search (e.g., $k_1 \\in [0.8, 1.6]$ with step 0.1 and $b \\in [0.5, 0.9]$ with step 0.05). Report your final best configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b24c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 6.3: Fine-grained grid search\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbc08fa0",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 7: Error Analysis & Discussion\n",
    "\n",
    "### 🚧 TODO 7.1: Where does BM25 still fail?\n",
    "\n",
    "Using your best configuration, find the 10 worst queries (lowest nDCG@10). For 3 of them:\n",
    "1. Print the query text\n",
    "2. Show the top-3 retrieved documents (first 200 chars)\n",
    "3. Show 2 gold documents (first 200 chars)\n",
    "4. Identify the failure mode: vocabulary mismatch? Ambiguity? Length bias? Other?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd6cccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 7.1: Error analysis with best BM25 config\n",
    "\n",
    "# ... your code here: build best BM25, evaluate per-query, find worst queries ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89516b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 7.1 (continued): Deep dive into 3 worst queries\n",
    "\n",
    "# ... your code here: for each query, print retrieved vs gold documents ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e888113",
   "metadata": {},
   "source": [
    "**🚧 TODO 7.2:** Based on your error analysis, categorize the failure modes you observed and answer:\n",
    "1. What is the single most common failure mode for BM25?\n",
    "2. Give a specific example of vocabulary mismatch from your results.\n",
    "3. Can parameter tuning or preprocessing fix these failures? Why or why not?\n",
    "4. What kind of approach might solve the vocabulary mismatch problem? (Hint: think about what we'll cover in Class 3.)\n",
    "\n",
    "**Answer 7.2:**\n",
    "\n",
    "... your answer here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "352a99bb",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### 🚧 TODO 7.3: Final Summary\n",
    "\n",
    "Record your best BM25 configuration and results here. You'll need this for the in-class session.\n",
    "\n",
    "| Setting | Value |\n",
    "|---------|-------|\n",
    "| Stemmer | ??? |\n",
    "| Stopwords | ??? |\n",
    "| $k_1$ | ??? |\n",
    "| $b$ | ??? |\n",
    "| nDCG@10 (train) | ??? |\n",
    "\n",
    "---\n",
    "\n",
    "## Further Reading & Things to Explore\n",
    "\n",
    "Here are optional resources if you want to go deeper:\n",
    "\n",
    "- 🎥 [Lecture by Stephen Robertson on BM25 (1h)](https://www.youtube.com/watch?v=a3sg6MH8m4k) — the inventor himself\n",
    "- 📖 [Elasticsearch BM25 similarity documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/index-modules-similarity.html) — how BM25 is configured in production search\n",
    "- 📖 [Lv & Zhai (2011) — BM25+ and BM25L](https://dl.acm.org/doi/10.1145/2063576.2063584) — the variants you'll implement in class\n",
    "- 📖 [Trotman et al. (2014) — Improvements to BM25](https://www.cs.otago.ac.nz/homepages/andrew/papers/2014-2.pdf) — comprehensive comparison of BM25 variants\n",
    "- 🔧 [Pyterrier](https://pyterrier.readthedocs.io/) — a full IR experimentation toolkit in Python (uses BM25 internally)\n",
    "- 📊 [MS MARCO Leaderboard](https://microsoft.github.io/msmarco/) — see how BM25 compares to neural methods\n",
    "\n",
    "*Next: In class on March 20 you'll implement BM25 variants (BM25+, BM25L, BM25F) and build the strongest sparse baseline. Bring your best configuration!*"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
