{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06d68a2d",
   "metadata": {},
   "source": [
    "# Information Retrieval — Class 4 Prep: Re-Ranking, Cross-Encoders & Hybrid Search\n",
    "\n",
    "**Due:** April 3, 2026 (before class)\n",
    "\n",
    "**Goal:** Build the modern retrieval pipeline — combine BM25 and dense retrieval with Reciprocal Rank Fusion, then re-rank with a cross-encoder. Explore ColBERT and neural query expansion.\n",
    "\n",
    "## Notebook Outline\n",
    "\n",
    "| Part | Topic | Description |\n",
    "|------|-------|-------------|\n",
    "| **1** | The Two-Stage Pipeline | Why retrieve-then-rerank, computational cost analysis |\n",
    "| **2** | Cross-Encoders | Architecture, training, `cross-encoder/ms-marco-MiniLM-L-6-v2` |\n",
    "| **3** | Implementing Re-Ranking | BM25 top-100 → cross-encoder → measure nDCG improvement |\n",
    "| **4** | Hybrid Retrieval: RRF | Reciprocal Rank Fusion: combine BM25 + dense without training |\n",
    "| **5** | Neural Query Expansion | doc2query-T5: generating synthetic queries for documents |\n",
    "| **6** | ColBERT & Late Interaction | Token-level matching with MaxSim |\n",
    "| **7** | Reading & Resources | Nogueira & Cho (2019) |\n",
    "\n",
    "**Key Concepts:**\n",
    "- Cross-encoder architecture and contrastive training\n",
    "- Two-stage retrieval pipeline (retrieve → re-rank)\n",
    "- Reciprocal Rank Fusion (RRF)\n",
    "- Score interpolation vs rank-based fusion\n",
    "- Neural document expansion (doc2query-T5)\n",
    "- Late interaction models (ColBERT, MaxSim)\n",
    "\n",
    "**Required Libraries:** `sentence-transformers`, `faiss-cpu`, `rank_bm25`, `transformers`, `torch`\n",
    "\n",
    "> 📌 **Installation:**\n",
    "> ```bash\n",
    "> pip install sentence-transformers faiss-cpu rank_bm25\n",
    "> ```\n",
    "\n",
    "**Essential Reading & Resources:**\n",
    "\n",
    "| Resource | Link |\n",
    "|----------|------|\n",
    "| Nogueira & Cho (2019), *\"Passage Re-ranking with BERT\"* | [Paper](https://arxiv.org/abs/1901.04085) |\n",
    "| Cormack, Clarke & Butt (2009), *\"Reciprocal Rank Fusion\"* | [Paper](https://dl.acm.org/doi/10.1145/1571941.1572114) |\n",
    "| Khattab & Zaharia (2020), *\"ColBERT\"* | [Paper](https://arxiv.org/abs/2004.12832) |\n",
    "| Nogueira & Lin (2019), *\"doc2query\"* | [Paper](https://arxiv.org/abs/1904.08375) |\n",
    "| Humeau et al. (2020), *\"Poly-encoders\"* | [Paper](https://arxiv.org/abs/1905.01969) |\n",
    "| sentence-transformers: Cross-Encoders | [Docs](https://www.sbert.net/docs/cross_encoder/usage/usage.html) |\n",
    "| Video: Bi-Encoders vs Cross-Encoders (James Briggs) | [YouTube](https://www.youtube.com/watch?v=lGPRByYrmgk) |\n",
    "| ColBERT v2 (Santhanam et al., 2022) | [Paper](https://arxiv.org/abs/2112.01488) |\n",
    "| RAGatouille library | [GitHub](https://github.com/bclavie/RAGatouille) |\n",
    "| Elasticsearch RRF docs | [Elastic](https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375dbb97",
   "metadata": {},
   "source": [
    "<span style=\"color:magenta\">Student Name:</span>\n",
    "\n",
    "* Name\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac76ec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "import urllib.request\n",
    "import time\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 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": "22fd1769",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1: Data, Baselines & the Two-Stage Pipeline\n",
    "\n",
    "### 1.1 Data Loading & Evaluation Toolkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6f16746",
   "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",
    "    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": "76a41bc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "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, ~10k docs)\n",
    "# Set FAST_MODE = False for full dataset (500 train + 500 test, ~92k docs)\n",
    "FAST_MODE = False\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": "code",
   "execution_count": null,
   "id": "1a7957eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Evaluation metrics (from Classes 1-3) ---\n",
    "\n",
    "def precision_at_k(retrieved, relevant, k):\n",
    "    return len(set(retrieved[:k]) & relevant) / k\n",
    "\n",
    "def recall_at_k(retrieved, relevant, k):\n",
    "    if len(relevant) == 0: return 0.0\n",
    "    return len(set(retrieved[:k]) & relevant) / len(relevant)\n",
    "\n",
    "def reciprocal_rank(retrieved, relevant):\n",
    "    for rank, doc_id in enumerate(retrieved, 1):\n",
    "        if doc_id in relevant: return 1.0 / rank\n",
    "    return 0.0\n",
    "\n",
    "def average_precision(retrieved, relevant):\n",
    "    if len(relevant) == 0: return 0.0\n",
    "    ap, rel = 0.0, 0\n",
    "    for k, doc_id in enumerate(retrieved, 1):\n",
    "        if doc_id in relevant:\n",
    "            rel += 1; ap += rel / k\n",
    "    return ap / len(relevant)\n",
    "\n",
    "def dcg_at_k(rels, k):\n",
    "    rels = np.array(rels[:k], dtype=float)\n",
    "    return float(np.sum((2**rels - 1) / np.log2(np.arange(1, len(rels)+1) + 1)))\n",
    "\n",
    "def ndcg_at_k(retrieved, relevant, k):\n",
    "    rels = [1 if d in relevant else 0 for d in retrieved[:k]]\n",
    "    dcg = dcg_at_k(rels, k)\n",
    "    idcg = dcg_at_k(sorted(rels, reverse=True), k)\n",
    "    return dcg / idcg if idcg > 0 else 0.0\n",
    "\n",
    "def evaluate_retriever(retrieve_fn, queries_df, gold_dict, docs_df, k=10):\n",
    "    p, mrr, ap, ndcg = [], [], [], []\n",
    "    for _, row in tqdm(queries_df.iterrows(), total=len(queries_df), desc=\"Eval\"):\n",
    "        rel = gold_dict.get(row[\"qid\"], set())\n",
    "        if not rel: continue\n",
    "        idx, _ = retrieve_fn(row[\"query\"])\n",
    "        ids = [docs_df.iloc[i][\"docid\"] for i in idx]\n",
    "        p.append(precision_at_k(ids, rel, k))\n",
    "        mrr.append(reciprocal_rank(ids, rel))\n",
    "        ap.append(average_precision(ids, rel))\n",
    "        ndcg.append(ndcg_at_k(ids, rel, k))\n",
    "    return {f\"P@{k}\": np.mean(p), \"MRR\": np.mean(mrr), \"MAP\": np.mean(ap), f\"nDCG@{k}\": np.mean(ndcg)}\n",
    "\n",
    "def evaluate_ranking(ranked_docids, queries_df, gold_dict, k=10):\n",
    "    \"\"\"Evaluate pre-computed rankings. ranked_docids: dict qid → [docid list].\"\"\"\n",
    "    p, mrr, ap, ndcg = [], [], [], []\n",
    "    for _, row in queries_df.iterrows():\n",
    "        rel = gold_dict.get(row[\"qid\"], set())\n",
    "        if not rel: continue\n",
    "        ids = ranked_docids.get(row[\"qid\"], [])\n",
    "        p.append(precision_at_k(ids, rel, k))\n",
    "        mrr.append(reciprocal_rank(ids, rel))\n",
    "        ap.append(average_precision(ids, rel))\n",
    "        ndcg.append(ndcg_at_k(ids, rel, k))\n",
    "    return {f\"P@{k}\": np.mean(p), \"MRR\": np.mean(mrr), \"MAP\": np.mean(ap), f\"nDCG@{k}\": np.mean(ndcg)}\n",
    "\n",
    "def print_results(name, r):\n",
    "    print(f\"\\n{'='*40}\\n  {name}\\n{'='*40}\")\n",
    "    for m, v in r.items(): print(f\"  {m:<12} {v:.4f}\")\n",
    "\n",
    "print(\"✓ Evaluation toolkit loaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e57e0ff",
   "metadata": {},
   "source": [
    "### 1.2 Build BM25 & Dense Baselines (from Classes 2–3)\n",
    "\n",
    "We need both systems for fusion experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de010752",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "from rank_bm25 import BM25Okapi\n",
    "import nltk\n",
    "from nltk.stem import PorterStemmer\n",
    "from nltk.corpus import stopwords\n",
    "\n",
    "nltk.download(\"stopwords\", quiet=True)\n",
    "english_stopwords = set(stopwords.words(\"english\"))\n",
    "porter = PorterStemmer()\n",
    "base_analyzer = CountVectorizer().build_analyzer()\n",
    "\n",
    "def preprocess(text):\n",
    "    tokens = base_analyzer(text)\n",
    "    tokens = [t for t in tokens if t not in english_stopwords]\n",
    "    tokens = [porter.stem(t) for t in tokens]\n",
    "    return tokens\n",
    "\n",
    "print(\"Tokenizing for BM25...\")\n",
    "tokenized_docs = docs[\"text\"].progress_apply(preprocess).tolist()\n",
    "bm25 = BM25Okapi(tokenized_docs, k1=1.2, b=0.75)\n",
    "\n",
    "def bm25_retrieve(query_text, k=100):\n",
    "    tokens = preprocess(query_text)\n",
    "    scores = bm25.get_scores(tokens)\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",
    "bm25_results = evaluate_retriever(lambda q, k=10: bm25_retrieve(q, k), train_queries, gold_dict, docs)\n",
    "print_results(\"BM25 Baseline\", bm25_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2454cfc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "import faiss\n",
    "\n",
    "# 🚧 TODO 1.1: Build the dense retrieval baseline\n",
    "\n",
    "# ... your code here ...\n",
    "# 1. Load \"all-MiniLM-L6-v2\" with SentenceTransformer()\n",
    "# 2. Encode all docs[\"text\"] with model.encode(..., show_progress_bar=True, batch_size=256)\n",
    "# 3. L2-normalize embeddings: emb / np.linalg.norm(emb, axis=1, keepdims=True)\n",
    "# 4. Build a FAISS inner-product index: faiss.IndexFlatIP(dimension)\n",
    "# 5. Add normalized doc embeddings to the index\n",
    "# 6. Define dense_retrieve(query_text, k=100) that encodes, normalizes, and searches\n",
    "# 7. Evaluate with evaluate_retriever() and print_results()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be80345e",
   "metadata": {},
   "source": [
    "### 1.3 The Two-Stage Pipeline: Why?\n",
    "\n",
    "Consider the cost of evaluating every query-document pair:\n",
    "\n",
    "| System | Cost per query (92k docs) | Why |\n",
    "|--------|:---:|-------|\n",
    "| BM25 | ~1 ms | Inverted index — only score docs with matching terms |\n",
    "| Dense (FAISS) | ~2 ms | ANN search |\n",
    "| Cross-Encoder | **~90 seconds** | Must pass each (query, doc) through BERT |\n",
    "\n",
    "The solution: **retrieve** a small candidate set (top-100) with a fast method, then **re-rank** with the accurate method.\n",
    "\n",
    "```\n",
    "92,000 documents  →  BM25 (1ms)  →  Top 100  →  Cross-Encoder (100ms)  →  Top 10\n",
    "```\n",
    "\n",
    "This gives us BM25 speed **and** neural accuracy.\n",
    "\n",
    "> 📖 This pipeline was formalized by Nogueira & Cho (2019). They showed BERT re-ranking could improve BM25 by **27%** on MS MARCO.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "868cd2d5",
   "metadata": {},
   "source": [
    "## Part 2: Cross-Encoders (~20 min)\n",
    "\n",
    "### 2.1 Bi-Encoder vs Cross-Encoder\n",
    "\n",
    "**Bi-Encoder** (Class 3):\n",
    "```\n",
    "\"what is rain\" → [BERT] → vec_q     \"precipitation...\" → [BERT] → vec_d\n",
    "                          cosine(vec_q, vec_d) = 0.72\n",
    "```\n",
    "- Independent encoding — fast (encode docs once), but query and document never \"see\" each other.\n",
    "\n",
    "**Cross-Encoder:**\n",
    "```\n",
    "[CLS] what is rain [SEP] precipitation occurs when water vapor... [SEP]\n",
    "                          → [BERT] → single score: 0.94\n",
    "```\n",
    "- Joint encoding — attention flows between query and document tokens.\n",
    "- Can match \"heart attack\" → \"myocardial infarction\", understand negation, handle complex reasoning.\n",
    "- But must run on **every** (query, document) pair — infeasible for the full corpus.\n",
    "\n",
    "> 📖 Humeau et al. (2020) showed cross-encoders outperform bi-encoders by **10–20%** on most benchmarks but are **1000× slower**.\n",
    "\n",
    "### 🚧 TODO 2.1: Why can a cross-encoder understand that \"heart attack\" and \"myocardial infarction\" are similar, while a bi-encoder might miss this relationship?\n",
    "\n",
    "Think about what happens inside the transformer's attention mechanism."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65791d13",
   "metadata": {},
   "source": [
    "**Answer 2.1:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1730e99",
   "metadata": {},
   "source": [
    "### 2.2 Cross-Encoder Training\n",
    "\n",
    "The training objective is simple — binary classification or regression:\n",
    "\n",
    "$$\\mathcal{L} = -\\sum_i \\left[ y_i \\log \\sigma(s_i) + (1 - y_i) \\log(1 - \\sigma(s_i)) \\right]$$\n",
    "\n",
    "where $s_i$ = model output score, $y_i$ = relevance label, $\\sigma$ = sigmoid.\n",
    "\n",
    "For graded relevance, MSE loss:\n",
    "\n",
    "$$\\mathcal{L} = \\sum_i (s_i - y_i)^2$$\n",
    "\n",
    "We'll use a pre-trained cross-encoder: `cross-encoder/ms-marco-MiniLM-L-6-v2` — trained on MS MARCO, same base model as our bi-encoder.\n",
    "\n",
    "> 📖 [sentence-transformers Cross-Encoder docs](https://www.sbert.net/docs/cross_encoder/usage/usage.html)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "469fd8a4",
   "metadata": {},
   "source": [
    "## Part 3: Implementing Re-Ranking (~25 min)\n",
    "\n",
    "### 🚧 TODO 3.1: Load the cross-encoder and score a few pairs\n",
    "\n",
    "Load `cross-encoder/ms-marco-MiniLM-L-6-v2` and score some (query, document) pairs to understand how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63618fc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import CrossEncoder\n",
    "\n",
    "cross_encoder = CrossEncoder(\"cross-encoder/ms-marco-MiniLM-L-6-v2\")\n",
    "\n",
    "# Test: score a few (query, document) pairs\n",
    "test_query = \"what causes rain\"\n",
    "test_pairs = [\n",
    "    (test_query, \"Precipitation occurs when water vapor in the atmosphere condenses into droplets.\"),\n",
    "    (test_query, \"The stock market experienced heavy rain in selling activity.\"),\n",
    "    (test_query, \"Rain forests are tropical biomes with heavy rainfall.\"),\n",
    "    (test_query, \"Water evaporates, rises, cools, condenses into clouds, and falls as precipitation.\"),\n",
    "]\n",
    "\n",
    "scores = cross_encoder.predict([list(pair) for pair in test_pairs])\n",
    "print(f\"Query: {test_query!r}\\n\")\n",
    "for (_, doc), score in sorted(zip(test_pairs, scores), key=lambda x: -x[1]):\n",
    "    print(f\"  [{score:+.4f}] {doc[:100]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac41d0a",
   "metadata": {},
   "source": [
    "### 🚧 TODO 3.2: Build the full re-ranking pipeline\n",
    "\n",
    "1. Use BM25 to retrieve top-100 candidates\n",
    "2. Score each candidate with the cross-encoder\n",
    "3. Re-rank by cross-encoder score\n",
    "4. Evaluate the re-ranked top-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4fd0be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.2: BM25 → Cross-Encoder re-ranking pipeline\n",
    "\n",
    "def rerank_with_cross_encoder(query_text, candidate_indices, docs_df, ce_model, k=10):\n",
    "    \"\"\"Re-rank candidate documents using a cross-encoder.\"\"\"\n",
    "    # ... your code here ...\n",
    "    # 1. Build pairs: [[query_text, docs_df.iloc[idx][\"text\"][:512]] for each candidate]\n",
    "    # 2. Score with ce_model.predict(pairs)\n",
    "    # 3. Sort by score descending: np.argsort(scores)[::-1][:k]\n",
    "    # 4. Return candidate_indices[top_k], scores[top_k]\n",
    "    pass\n",
    "\n",
    "# Test the pipeline on one query and display BM25 top-5 vs re-ranked top-5\n",
    "# Hint: use train_queries.iloc[0], call bm25_retrieve(k=100), then rerank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c3026f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.2 (continued): Evaluate the full pipeline on train queries\n",
    "\n",
    "# ... your code here ...\n",
    "# Loop over train_queries, for each query:\n",
    "#   1. candidate_idx, _ = bm25_retrieve(query, k=100)\n",
    "#   2. reranked_idx, _ = rerank_with_cross_encoder(query, candidate_idx, docs, cross_encoder)\n",
    "#   3. Store docids in a dict: results[qid] = [docs.iloc[i][\"docid\"] for i in reranked_idx]\n",
    "# Then evaluate with evaluate_ranking() and compare to bm25_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ee57e6",
   "metadata": {},
   "source": [
    "### 🚧 TODO 3.3: Timing analysis\n",
    "\n",
    "Measure the time per query for each pipeline component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d836ad98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.3: Timing analysis\n",
    "\n",
    "# ... your code here ...\n",
    "# Time ~20 queries for each stage using time.time():\n",
    "#   t0 = time.time(); [bm25_retrieve(q, k=100) for each query]; elapsed = time.time() - t0\n",
    "# Measure: BM25 alone, Dense alone, BM25 + Cross-Encoder (retrieve + rerank)\n",
    "# Print average ms per query for each"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7759ac8",
   "metadata": {},
   "source": [
    "**🚧 TODO 3.4:** Answer these questions:\n",
    "\n",
    "1. How much did cross-encoder re-ranking improve nDCG@10 over BM25 alone?\n",
    "2. What is the time cost? Is it acceptable?\n",
    "3. What happens if you increase the candidate pool from top-100 to top-200?\n",
    "\n",
    "**Answer 3.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f57f448b",
   "metadata": {},
   "source": [
    "## Part 4: Hybrid Retrieval with RRF (~25 min)\n",
    "\n",
    "### 4.1 Reciprocal Rank Fusion\n",
    "\n",
    "What if we combine BM25 and dense retrieval rankings **before** re-ranking?\n",
    "\n",
    "**RRF** (Cormack, Clarke & Butt, 2009):\n",
    "\n",
    "$$\\text{RRF}(d) = \\sum_{r \\in R} \\frac{1}{k + \\text{rank}_r(d)}$$\n",
    "\n",
    "where $R$ is the set of ranking systems, $\\text{rank}_r(d)$ is the rank of document $d$ in system $r$, and $k$ is a constant (typically 60).\n",
    "\n",
    "**Why RRF is great:**\n",
    "- No training needed — just ranks\n",
    "- No score normalization — ranks are already comparable\n",
    "- Works across heterogeneous systems (BM25, dense, keyword, anything)\n",
    "- Robust — even if one system is bad, it barely hurts\n",
    "\n",
    "> 📖 The original [RRF paper](https://dl.acm.org/doi/10.1145/1571941.1572114) is only **2 pages long**. The $k=60$ constant was found empirically.\n",
    "\n",
    "### 🚧 TODO 4.1: Implement RRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b742808",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.1: Implement Reciprocal Rank Fusion\n",
    "\n",
    "def reciprocal_rank_fusion(rankings, k=60):\n",
    "    \"\"\"Combine multiple rankings using RRF.\n",
    "\n",
    "    Args:\n",
    "        rankings: list of lists of docids\n",
    "        k: RRF constant (default: 60)\n",
    "    Returns:\n",
    "        sorted list of (docid, rrf_score) pairs\n",
    "    \"\"\"\n",
    "    # ... your code here ...\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bb3b946",
   "metadata": {},
   "source": [
    "### 🚧 TODO 4.2: Evaluate RRF on train queries\n",
    "\n",
    "Compare: BM25 alone, Dense alone, RRF(BM25 + Dense)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5858fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.2: Evaluate RRF\n",
    "\n",
    "# ... your code here ...\n",
    "# For each train query:\n",
    "#   1. bm25_idx, _ = bm25_retrieve(query, k=100)\n",
    "#   2. dense_idx, _ = dense_retrieve(query, k=100)\n",
    "#   3. Convert indices to docids: [docs.iloc[i][\"docid\"] for i in idx]\n",
    "#   4. fused = reciprocal_rank_fusion([bm25_docids, dense_docids], k=60)\n",
    "#   5. Store top-10: [docid for docid, _ in fused[:10]]\n",
    "# Evaluate with evaluate_ranking() and compare BM25, Dense, and RRF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d1d020",
   "metadata": {},
   "source": [
    "### 🚧 TODO 4.3: RRF vs Score Interpolation\n",
    "\n",
    "Implement score interpolation as an alternative to RRF:\n",
    "\n",
    "$$s(d) = \\alpha \\cdot \\hat{s}_{\\text{BM25}}(d) + (1-\\alpha) \\cdot \\hat{s}_{\\text{dense}}(d)$$\n",
    "\n",
    "where $\\hat{s}$ denotes min-max normalized scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda97d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.3: Score interpolation\n",
    "\n",
    "# ... your code here ...\n",
    "# 1. Define score_interpolation(bm25_idx, bm25_scores, dense_idx, dense_scores, alpha=0.5, k=10)\n",
    "#    - Min-max normalize each: (scores - min) / (max - min + 1e-9)\n",
    "#    - combined[idx] = alpha * bm25_norm + (1 - alpha) * dense_norm\n",
    "#    - Sort by combined score, return top-k indices and scores\n",
    "# 2. Sweep alpha from 0.0 to 1.0 in steps of 0.2, evaluate each\n",
    "# 3. Compare best alpha result to RRF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7ba895f",
   "metadata": {},
   "source": [
    "**🚧 TODO 4.4:** Compare RRF and score interpolation. Answer:\n",
    "1. Which method gives better nDCG@10?\n",
    "2. Which is more convenient to use in practice? Why?\n",
    "3. What is the best alpha value for score interpolation? What does it tell you about the relative strength of BM25 vs dense?\n",
    "\n",
    "**Answer 4.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9095805b",
   "metadata": {},
   "source": [
    "### 🚧 TODO 4.5: The Full Pipeline — RRF + Cross-Encoder\n",
    "\n",
    "Now combine everything: RRF fusion, then cross-encoder re-ranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd2fb87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.5: Full pipeline: BM25 + Dense → RRF → Cross-Encoder\n",
    "\n",
    "# ... your code here ...\n",
    "# For each train query:\n",
    "#   Stage 1: bm25_retrieve(k=100) + dense_retrieve(k=100) → convert to docid lists\n",
    "#   Stage 2: reciprocal_rank_fusion([bm25_ids, dense_ids], k=60) → take top-100 fused docids\n",
    "#   Stage 3: Convert fused docids back to indices, then rerank_with_cross_encoder(k=10)\n",
    "# Hint: you'll need a docid→index mapping: {docid: i for i, docid in enumerate(docs[\"docid\"])}\n",
    "# Evaluate and print comparison of ALL methods"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5633ec97",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 5: Neural Query Expansion — doc2query-T5 (~15 min)\n",
    "\n",
    "### 5.1 The Idea\n",
    "\n",
    "Instead of expanding the **query**, expand the **document** at indexing time by generating synthetic queries:\n",
    "\n",
    "1. Train a T5 model on (document → query) pairs from MS MARCO\n",
    "2. For each document, generate 5–20 potential queries\n",
    "3. Append these queries to the document text\n",
    "4. Index with BM25 as usual\n",
    "\n",
    "> 📖 Nogueira & Lin (2019): [doc2query paper](https://arxiv.org/abs/1904.08375). The T5 version (docTTTTTquery, 2020) showed **+10% MAP** on MS MARCO with BM25 — no neural retrieval needed!\n",
    "\n",
    "> ⚠️ Running doc2query-T5 on 92k documents is slow without a GPU. We'll demonstrate the concept on a small sample and discuss the impact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "911cc6ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate doc2query-T5 concept on a few documents\n",
    "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "\n",
    "# Load a small doc2query model\n",
    "print(\"Loading doc2query-T5 model (this may take a moment)...\")\n",
    "try:\n",
    "    d2q_tokenizer = T5Tokenizer.from_pretrained(\"doc2query/msmarco-t5-base-v1\", legacy=True)\n",
    "    d2q_model = T5ForConditionalGeneration.from_pretrained(\"doc2query/msmarco-t5-base-v1\")\n",
    "\n",
    "    # Generate queries for a sample document\n",
    "    sample_doc = docs.iloc[0][\"text\"][:512]\n",
    "    print(f\"\\nDocument: {sample_doc[:200]}...\\n\")\n",
    "\n",
    "    input_ids = d2q_tokenizer.encode(sample_doc, return_tensors=\"pt\", max_length=512, truncation=True)\n",
    "    outputs = d2q_model.generate(\n",
    "        input_ids=input_ids,\n",
    "        max_length=64,\n",
    "        do_sample=True,\n",
    "        top_k=10,\n",
    "        num_return_sequences=5,\n",
    "    )\n",
    "\n",
    "    print(\"Generated queries:\")\n",
    "    for i, output in enumerate(outputs, 1):\n",
    "        query = d2q_tokenizer.decode(output, skip_special_tokens=True)\n",
    "        print(f\"  {i}. {query}\")\n",
    "except Exception as e:\n",
    "    print(f\"Could not load doc2query model: {e}\")\n",
    "    print(\"This is expected if the model isn't cached. The concept is demonstrated in the text above.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b63d8d8",
   "metadata": {},
   "source": [
    "### 🚧 TODO 5.1: How does document expansion help BM25?\n",
    "\n",
    "Consider this document: *\"The Great Barrier Reef is the world's largest coral reef system, located off the coast of Queensland, Australia.\"*\n",
    "\n",
    "doc2query generates: *\"where is the great barrier reef\"*, *\"largest coral reef in the world\"*, *\"great barrier reef australia\"*\n",
    "\n",
    "Answer:\n",
    "1. Before expansion, which query would BM25 match well? Which would it miss?\n",
    "2. After expansion, which new query becomes matchable?\n",
    "3. What is the downside of adding too many synthetic queries?\n",
    "\n",
    "**Answer 5.1:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8caadc9f",
   "metadata": {},
   "source": [
    "## Part 6: ColBERT & Late Interaction (~15 min)\n",
    "\n",
    "### 6.1 The ColBERT Architecture\n",
    "\n",
    "ColBERT is a middle ground between bi-encoders and cross-encoders:\n",
    "\n",
    "- **Bi-encoder:** 1 vector per document → fast, coarse\n",
    "- **Cross-encoder:** no pre-computation → slow, precise\n",
    "- **ColBERT:** many vectors per document (1 per token) → medium speed, fine-grained\n",
    "\n",
    "```\n",
    "Query: \"what causes rain\"\n",
    "  → [vec_what, vec_causes, vec_rain]  (3 token vectors)\n",
    "\n",
    "Document: \"precipitation occurs when moisture...\"\n",
    "  → [vec_precipitation, vec_occurs, vec_when, vec_moisture, ...]  (N token vectors)\n",
    "```\n",
    "\n",
    "**MaxSim scoring:**\n",
    "\n",
    "$$\\text{score}(q, d) = \\sum_{i=1}^{|q|} \\max_{j=1}^{|d|} \\text{sim}(q_i, d_j)$$\n",
    "\n",
    "For each query token, find its best-matching document token, then sum.\n",
    "\n",
    "> 📖 [ColBERT paper (Khattab & Zaharia, 2020)](https://arxiv.org/abs/2004.12832) — Stanford.\n",
    ">\n",
    "> 📖 [RAGatouille library](https://github.com/bclavie/RAGatouille) by Ben Clavié — makes ColBERT easy to use.\n",
    "\n",
    "### 🚧 TODO 6.1: Implement MaxSim scoring\n",
    "\n",
    "Given pre-computed token embeddings, implement the MaxSim scoring function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "077844ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 6.1: MaxSim scoring\n",
    "\n",
    "def maxsim_score(query_token_embeddings, doc_token_embeddings):\n",
    "    \"\"\"Compute ColBERT-style MaxSim score.\"\"\"\n",
    "    # ... your code here ...\n",
    "    # 1. Compute cosine similarity matrix: query_embs @ doc_embs.T → shape (n_q, n_d)\n",
    "    # 2. For each query token, take max similarity across all doc tokens: .max(axis=1)\n",
    "    # 3. Sum the max similarities\n",
    "    pass\n",
    "\n",
    "# Test with random embeddings\n",
    "np.random.seed(42)\n",
    "q_embs = np.random.randn(3, 128).astype(np.float32)   # 3 query tokens\n",
    "d_embs = np.random.randn(20, 128).astype(np.float32)   # 20 doc tokens\n",
    "q_embs = q_embs / np.linalg.norm(q_embs, axis=1, keepdims=True)\n",
    "d_embs = d_embs / np.linalg.norm(d_embs, axis=1, keepdims=True)\n",
    "# score = maxsim_score(q_embs, d_embs)\n",
    "# print(f\"MaxSim score: {score:.4f}\")  # Should be a small positive number"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a208ba3",
   "metadata": {},
   "source": [
    "### 🚧 TODO 6.2: ColBERT comparison table\n",
    "\n",
    "Fill in this table comparing the three architectures:\n",
    "\n",
    "| | Bi-Encoder | ColBERT | Cross-Encoder |\n",
    "|--|:-:|:-:|:-:|\n",
    "| Vectors per document | 1 | N (# tokens) | 0 (on-the-fly) |\n",
    "| Index-time cost | Medium | High | None |\n",
    "| Query-time cost | Low | Medium | High |\n",
    "| Storage (1M docs) | ~1.5 GB | ~15–50 GB | 0 |\n",
    "| Can pre-compute docs? | Yes | Yes | No |\n",
    "| Accuracy (relative) | Good | Very good | Best |\n",
    "\n",
    "**🚧 TODO 6.2:** Why is ColBERT's storage so much higher than a bi-encoder? What is ColBERTv2's solution?\n",
    "\n",
    "**Answer 6.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc444a31",
   "metadata": {},
   "source": [
    "## Part 7: Summary & Further Reading\n",
    "\n",
    "### Pipeline Summary\n",
    "\n",
    "Record your results:\n",
    "\n",
    "| System | nDCG@10 | MRR | Time/query |\n",
    "|--------|:---:|:---:|:---:|\n",
    "| BM25 | | | |\n",
    "| Dense (MiniLM) | | | |\n",
    "| RRF (BM25 + Dense) | | | |\n",
    "| BM25 → Cross-Encoder | | | |\n",
    "| BM25 + Dense → RRF → CE | | | |\n",
    "\n",
    "### For the In-Class Session:\n",
    "- You'll implement the full pipeline and apply it to the challenge dataset\n",
    "- Key code to bring: RRF function, cross-encoder re-ranking, evaluation toolkit\n",
    "- The in-class challenge goal: combine BM25 + Dense + RRF + CE on scientific articles\n",
    "\n",
    "### Further Reading\n",
    "\n",
    "- 📖 [Nogueira & Cho (2019)](https://arxiv.org/abs/1901.04085) — first cross-encoder for IR (**required reading**)\n",
    "- 📖 [Cormack, Clarke & Butt (2009)](https://dl.acm.org/doi/10.1145/1571941.1572114) — the 2-page RRF paper\n",
    "- 📖 [Khattab & Zaharia (2020)](https://arxiv.org/abs/2004.12832) — ColBERT\n",
    "- 📖 [Nogueira & Lin (2019)](https://arxiv.org/abs/1904.08375) — doc2query\n",
    "- 🎥 [Two-stage retrieval explained (James Briggs)](https://www.youtube.com/watch?v=lGPRByYrmgk) — 15 min video\n",
    "- 🔧 [RAGatouille](https://github.com/bclavie/RAGatouille) — ColBERT made easy\n",
    "- 📊 [MS MARCO leaderboard](https://microsoft.github.io/msmarco/) — where all these models compete\n",
    "\n",
    "*Next: In class on April 3 you'll build the full pipeline for the challenge and do error analysis on scientific queries!*"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}