{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ab7ccd8",
   "metadata": {},
   "source": [
    "# Information Retrieval — Class 3 Prep: Dense Retrieval & Bi-Encoders\n",
    "\n",
    "**Due:** March 27, 2026 (before class)\n",
    "\n",
    "**Goal:** Move from lexical to semantic retrieval — encode queries and documents as dense vectors with sentence-transformers, build FAISS indexes for fast search, compare dense vs sparse retrieval on MS MARCO, and get your first look at the scientific article challenge dataset.\n",
    "\n",
    "## Notebook Outline\n",
    "\n",
    "| Part | Topic | Description |\n",
    "|------|-------|-------------|\n",
    "| **1** | Setup & Data | Load MS MARCO, evaluation toolkit, install new libraries |\n",
    "| **2** | From Word Embeddings to Sentence Embeddings | Recap Word2Vec, mean pooling, why raw BERT fails, sentence-transformers |\n",
    "| **3** | Dense Retrieval on MS MARCO | Encode corpus, brute-force cosine search, compare to BM25 |\n",
    "| **4** | FAISS for Fast Retrieval | Approximate nearest neighbors: IndexFlatIP, IVF, HNSW |\n",
    "| **5** | When Does Dense Beat Sparse? | Semantic matching analysis, failure modes |\n",
    "| **6** | 🔬 The Scientific Article Domain | Load challenge data, explore differences from web text |\n",
    "| **7** | Domain-Specific Models | SPECTER, SciBERT — encoding scientific text |\n",
    "| **8** | Reading & Resources | Karpukhin et al. (2020) |\n",
    "\n",
    "**Key Concepts:**\n",
    "- Bi-encoder architecture for retrieval\n",
    "- Sentence-transformers library and pre-trained models\n",
    "- Cosine similarity search at scale\n",
    "- FAISS indexing (exact and approximate nearest neighbors)\n",
    "- Dense vs sparse trade-offs\n",
    "- Domain transfer: web text → scientific articles\n",
    "- Domain-specific embedding models (SPECTER, SciBERT)\n",
    "\n",
    "**Required Libraries:** `sentence-transformers`, `faiss-cpu`, `torch`, `numpy`, `pandas`, `matplotlib`, `rank_bm25`\n",
    "\n",
    "> 📌 **Installation:**\n",
    "> ```bash\n",
    "> pip install sentence-transformers faiss-cpu\n",
    "> ```\n",
    "\n",
    "**Essential Reading & Resources:**\n",
    "\n",
    "| Resource | Link |\n",
    "|----------|------|\n",
    "| Karpukhin et al. (2020), *\"Dense Passage Retrieval for Open-Domain QA\"* | [Paper](https://arxiv.org/abs/2004.04906) |\n",
    "| Reimers & Gurevych (2019), *\"Sentence-BERT\"* | [Paper](https://arxiv.org/abs/1908.10084) |\n",
    "| Wikipedia: Bi-encoder | [Link](https://en.wikipedia.org/wiki/Sentence_embedding) |\n",
    "| Sentence-Transformers documentation | [Docs](https://www.sbert.net/) |\n",
    "| FAISS Wiki (Facebook Research) | [GitHub Wiki](https://github.com/facebookresearch/faiss/wiki) |\n",
    "| FAISS tutorial notebooks | [GitHub](https://github.com/facebookresearch/faiss/tree/main/tutorial/python) |\n",
    "| Hugging Face Model Hub: Sentence-Transformers | [Hub](https://huggingface.co/models?library=sentence-transformers) |\n",
    "| Video: Bi-Encoders vs Cross-Encoders (James Briggs) | [YouTube](https://www.youtube.com/watch?v=lGPRByYrmgk) |\n",
    "| SPECTER paper (Cohan et al., 2020) | [Paper](https://arxiv.org/abs/2004.07180) |\n",
    "| Blog: Getting Started with FAISS | [Pinecone](https://www.pinecone.io/learn/series/faiss/) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cdfd7b0",
   "metadata": {},
   "source": [
    "<span style=\"color:magenta\">Student Name:</span>\n",
    "\n",
    "* Name\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64f53389",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"] = \"TRUE\"  # prevent OpenMP conflict between PyTorch and FAISS on macOS\n",
    "\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 sklearn.feature_extraction.text import CountVectorizer\n",
    "from tqdm import tqdm\n",
    "\n",
    "import faiss  # import early, before sentence-transformers/PyTorch\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": "df11bb80",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1: Data Loading & Evaluation Toolkit\n",
    "\n",
    "Same as Class 2 — load MS MARCO and our evaluation functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f1439d",
   "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": "7df1778d",
   "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",
    "# ⚠️ Dense encoding of 92k docs takes ~5–10 min on CPU. Use FAST_MODE on weak machines!\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": "code",
   "execution_count": null,
   "id": "136293d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Evaluation metrics (from Class 1 & 2) ---\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 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": "17552fd3",
   "metadata": {},
   "source": [
    "### BM25 Baseline (from Class 2)\n",
    "\n",
    "We need a BM25 baseline to compare dense retrieval against. Use your best configuration from the Class 2 prep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ba3014",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FastBM25:\n",
    "    def __init__(self, corpus, k1=1.2, b=0.75):\n",
    "        self.k1 = k1\n",
    "        self.b = b\n",
    "        self.vectorizer = CountVectorizer(stop_words='english')\n",
    "        print(\"Building TF matrix...\")\n",
    "        self.tf_matrix = self.vectorizer.fit_transform(corpus)\n",
    "        self.doc_lens = self.tf_matrix.sum(axis=1).A1\n",
    "        self.avgdl = self.doc_lens.mean()\n",
    "        self.N = self.tf_matrix.shape[0]\n",
    "        df_arr = (self.tf_matrix > 0).sum(axis=0).A1\n",
    "        self.idf = np.maximum(np.log((self.N - df_arr + 0.5) / (df_arr + 0.5)), 0)\n",
    "        print(f\"✓ FastBM25 indexed {self.N} documents.\")\n",
    "\n",
    "    def get_scores(self, query_text):\n",
    "        query_vec = self.vectorizer.transform([query_text])\n",
    "        query_indices = query_vec.nonzero()[1]\n",
    "        query_freqs = query_vec.toarray().flatten()\n",
    "        \n",
    "        scores = np.zeros(self.N)\n",
    "        for idx in query_indices:\n",
    "            q_tf = query_freqs[idx]\n",
    "            tf_col = self.tf_matrix[:, idx].toarray().flatten()\n",
    "            numerator = tf_col * (self.k1 + 1)\n",
    "            denominator = tf_col + self.k1 * (1 - self.b + self.b * self.doc_lens / self.avgdl)\n",
    "            scores += q_tf * self.idf[idx] * (numerator / denominator)\n",
    "        return scores\n",
    "\n",
    "bm25 = FastBM25(docs[\"text\"])\n",
    "\n",
    "def bm25_retrieve(query_text, k=10):\n",
    "    scores = bm25.get_scores(query_text)\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(bm25_retrieve, train_queries, gold_dict, docs)\n",
    "print_results(\"BM25 Baseline\", bm25_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83d146b5",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 2: From Word Embeddings to Sentence Embeddings\n",
    "\n",
    "### 2.1 The Problem\n",
    "\n",
    "In HoNLP, you learned about Word2Vec: each **word** gets a vector, and similar words have similar vectors. For retrieval, we need a single vector for an entire **document** (or passage).\n",
    "\n",
    "> 📖 **Recap:** [Word2Vec — Mikolov et al. (2013)](https://arxiv.org/abs/1301.3781) — you implemented Skip-Gram in HoNLP!\n",
    "\n",
    "| Method | Idea | Quality for IR |\n",
    "|--------|------|----------------|\n",
    "| Average word vectors | $\\vec{d} = \\frac{1}{n}\\sum_i \\vec{w}_i$ | Mediocre — loses word order |\n",
    "| TF-IDF weighted average | $\\vec{d} = \\sum_i \\text{tfidf}_i \\cdot \\vec{w}_i$ | Slightly better |\n",
    "| BERT [CLS] token | Use the special [CLS] output from BERT | **Bad!** Not trained for similarity |\n",
    "| Sentence-BERT | Fine-tune BERT *specifically* for semantic similarity | ✅ The right approach |\n",
    "\n",
    "> ⚠️ **Surprise:** Raw BERT [CLS] embeddings are **terrible** for similarity search — even worse than GloVe averages! (Reimers & Gurevych, 2019). You need a model specifically fine-tuned for this task.\n",
    "\n",
    "### 🚧 TODO 2.1: Why does mean-pooling of Word2Vec fail for retrieval?\n",
    "\n",
    "Consider these two texts:\n",
    "- *\"The bank of the river was covered in flowers\"*\n",
    "- *\"She deposited money at the bank\"*\n",
    "\n",
    "Both contain \"bank\" — Word2Vec gives the same vector for both. Answer:\n",
    "1. What is this problem called?\n",
    "2. How does BERT solve it?\n",
    "3. Why isn't the BERT solution enough for retrieval (without fine-tuning)?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "529735e3",
   "metadata": {},
   "source": [
    "**Answer 2.1:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35e31005",
   "metadata": {},
   "source": [
    "### 2.2 The Bi-Encoder Architecture\n",
    "\n",
    "The key insight of Sentence-BERT: encode queries and documents **independently** through the same transformer, then compare with cosine similarity.\n",
    "\n",
    "```\n",
    "Query: \"what causes rain\"          Document: \"precipitation occurs when...\"\n",
    "    ↓                                  ↓\n",
    "  [BERT encoder]                    [BERT encoder]\n",
    "    ↓                                  ↓\n",
    "  query vector (384-d)              doc vector (384-d)\n",
    "    ↓                                  ↓\n",
    "              cosine similarity\n",
    "                   ↓\n",
    "              relevance score\n",
    "```\n",
    "\n",
    "**Why this is brilliant for retrieval:**\n",
    "- Encode all documents **once** at index time → store the vectors\n",
    "- At query time: encode only the query → compare to all stored vectors\n",
    "- $O(1)$ document encoding per query (already done!)\n",
    "\n",
    "The training uses **contrastive loss** (InfoNCE):\n",
    "\n",
    "$$\\mathcal{L} = -\\log \\frac{e^{\\text{sim}(q, d^+) / \\tau}}{e^{\\text{sim}(q, d^+) / \\tau} + \\sum_j e^{\\text{sim}(q, d_j^-) / \\tau}}$$\n",
    "\n",
    "> 📖 Read: [Sentence-BERT paper (Reimers & Gurevych, 2019)](https://arxiv.org/abs/1908.10084) — the foundation for everything we do today.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69b5abef",
   "metadata": {},
   "source": [
    "## Part 3: Dense Retrieval on MS MARCO (~30 min)\n",
    "\n",
    "Now let's actually encode our documents and do semantic search!\n",
    "\n",
    "### 🚧 TODO 3.1: Load a sentence-transformer model and encode documents\n",
    "\n",
    "> 📖 **Docs:** [sentence-transformers — Getting Started](https://www.sbert.net/docs/quickstart.html)\n",
    "\n",
    "We'll use `all-MiniLM-L6-v2` — a fast, lightweight model (384 dimensions, ~80MB).\n",
    "Pre-trained on 1 billion sentence pairs.\n",
    "\n",
    "> ⚠️ Encoding ~10k documents (FAST_MODE) takes ~1 minute on CPU; the full ~92k takes ~5 minutes. Be patient, or use GPU if available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5e3ed64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# Load model\n",
    "model_name = \"all-MiniLM-L6-v2\"\n",
    "model = SentenceTransformer(model_name)\n",
    "print(f\"Model: {model_name}\")\n",
    "print(f\"Embedding dimension: {model.get_sentence_embedding_dimension()}\")\n",
    "print(f\"Max sequence length: {model.max_seq_length}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba3bb1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.1: Encode all documents and queries\n",
    "\n",
    "# ... your code here ...\n",
    "# Encode documents (use model.encode with show_progress_bar=True, batch_size=256)\n",
    "# Encode queries\n",
    "# Normalize both for cosine similarity\n",
    "# ⚠️ Important: after normalizing, cast to float32 with .astype(np.float32)\n",
    "#    (numpy division produces float64, but FAISS requires float32)\n",
    "\n",
    "print(f\"Document embeddings shape: ???\")\n",
    "print(f\"Query embeddings shape: ???\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf0a451f",
   "metadata": {},
   "source": [
    "### 🚧 TODO 3.2: Brute-force dense retrieval\n",
    "\n",
    "Implement retrieval by computing the cosine similarity between a query embedding and all document embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0245fa6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.2: Brute-force dense retrieval\n",
    "\n",
    "def dense_retrieve(query_text, k=10):\n",
    "    \"\"\"Retrieve top-k documents using cosine similarity with pre-computed embeddings.\"\"\"\n",
    "    # ... your code here ...\n",
    "\n",
    "\n",
    "# Test on a query\n",
    "print(f\"Query: {queries.iloc[0]['query']!r}\\n\")\n",
    "top_k, scores = dense_retrieve(queries.iloc[0][\"query\"], k=5)\n",
    "for rank, (idx, score) in enumerate(zip(top_k, scores), 1):\n",
    "    print(f\"  #{rank} [cos={score:.3f}] {docs.iloc[idx]['docid']}: {docs.iloc[idx]['text'][:120]}...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b1d9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate dense retrieval\n",
    "dense_results = evaluate_retriever(dense_retrieve, train_queries, gold_dict, docs)\n",
    "print_results(\"Dense Retrieval (all-MiniLM-L6-v2)\", dense_results)\n",
    "print_results(\"BM25 Baseline\", bm25_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "811d5f76",
   "metadata": {},
   "source": [
    "### 🚧 TODO 3.3: Timing comparison\n",
    "\n",
    "Measure the wall-clock time for 100 queries with each method: BM25 vs brute-force dense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36645cb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.3: Timing comparison\n",
    "\n",
    "# ... your code here: time 100 queries with BM25 and with dense retrieval ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cf6b0f8",
   "metadata": {},
   "source": [
    "**🚧 TODO 3.4:** Answer these questions:\n",
    "1. Which system has higher nDCG@10 on the training queries? By how much?\n",
    "2. Which is faster at query time? Why?\n",
    "3. What is the \"hidden cost\" of dense retrieval that the timing above doesn't capture?\n",
    "\n",
    "**Answer 3.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44477eb4",
   "metadata": {},
   "source": [
    "## Part 4: FAISS for Fast Retrieval (~25 min)\n",
    "\n",
    "Brute-force search is $O(n \\cdot d)$ per query. For large corpora, we need **Approximate Nearest Neighbor (ANN)** search.\n",
    "\n",
    "> 📖 **Read:** [FAISS Getting Started guide](https://github.com/facebookresearch/faiss/wiki/Getting-started) — essential reference.\n",
    ">\n",
    "> 📖 **Tutorial:** [Pinecone's FAISS tutorial series](https://www.pinecone.io/learn/series/faiss/) — excellent visual explanations.\n",
    "\n",
    "### 🚧 TODO 4.1: Build and use a FAISS flat index\n",
    "\n",
    "The simplest FAISS index is `IndexFlatIP` (exact search using inner product)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c21f15ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.1: Build a FAISS flat index and search with it\n",
    "\n",
    "# ... your code here ...\n",
    "# Build IndexFlatIP with dimension d\n",
    "# Add doc_embeddings_norm (as float32)\n",
    "# Search for top-10 nearest neighbors of the first query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74056894",
   "metadata": {},
   "source": [
    "### 🚧 TODO 4.2: Compare FAISS index types\n",
    "\n",
    "Build three index types and compare speed vs recall:\n",
    "1. **Flat** (exact) — `IndexFlatIP`\n",
    "2. **IVF** (inverted file, approximate) — `IndexIVFFlat`\n",
    "3. **HNSW** (graph-based, approximate) — `IndexHNSWFlat`\n",
    "\n",
    "> 📖 [FAISS index types guide](https://github.com/facebookresearch/faiss/wiki/Faiss-indexes) — all index types explained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58da2a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.2: Build IVF and HNSW indexes, benchmark them\n",
    "\n",
    "# --- IVF index ---\n",
    "# nlist = 50 clusters, nprobe = 10\n",
    "# ... your code here ...\n",
    "\n",
    "# --- HNSW index ---\n",
    "# M = 16 connections, efSearch = 64\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e7f78ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.2 (continued): Benchmark speed and recall\n",
    "\n",
    "# Write a function to benchmark each index: measure query time and recall@10\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64673569",
   "metadata": {},
   "source": [
    "### 🚧 TODO 4.3: Explore the nprobe parameter for IVF\n",
    "\n",
    "The `nprobe` parameter controls how many IVF clusters to search. Higher = more accurate but slower. Test `nprobe` ∈ {1, 5, 10, 20, 50, 100}."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f14518",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 4.3: nprobe sweep and plot\n",
    "\n",
    "# ... your code here: test different nprobe values, plot speed vs recall ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a30187e3",
   "metadata": {},
   "source": [
    "**🚧 TODO 4.4:** Answer these questions:\n",
    "1. At `nprobe=10`, what is the recall@10? Is the ~5% loss acceptable?\n",
    "2. What would you choose for a production system with 10 million documents?\n",
    "3. Why is HNSW faster than IVF at the same recall level (typically)?\n",
    "4. What is the memory overhead of HNSW compared to flat index?\n",
    "\n",
    "**Answer 4.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2855ab23",
   "metadata": {},
   "source": [
    "## Part 5: When Does Dense Beat Sparse? (~15 min)\n",
    "\n",
    "### 🚧 TODO 5.1: Side-by-side comparison\n",
    "\n",
    "For 500 training queries, compute both BM25 and dense nDCG@10. Find:\n",
    "1. Queries where dense wins by a lot (nDCG gap > 0.3)\n",
    "2. Queries where BM25 wins by a lot\n",
    "3. The overall correlation between the two systems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3968ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 5.1: Dense vs Sparse per-query comparison\n",
    "\n",
    "# ... your code here: compute nDCG@10 for both systems on each query ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ebae76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 5.1 (continued): Scatter plot + show queries where each wins big\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d03da35d",
   "metadata": {},
   "source": [
    "**🚧 TODO 5.2:** Analyze the results and answer:\n",
    "1. Look at 2 queries where dense wins big. What do they have in common?\n",
    "2. Look at 2 queries where BM25 wins big. What pattern do you see?\n",
    "3. Based on this analysis, when would you recommend dense retrieval vs BM25 for a new application?\n",
    "\n",
    "**Answer 5.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea142983",
   "metadata": {},
   "source": [
    "## Part 6: 🔬 The Scientific Article Domain (~20 min)\n",
    "\n",
    "Now let's look at a different domain: **scientific articles**. This prepares you for the Codabench challenge.\n",
    "\n",
    "> ⚠️ **Challenge dataset:** Your instructor will provide the challenge data files. Place them in the same directory as this notebook.\n",
    "\n",
    "For this prep, we'll simulate the challenge format using a sample. The actual challenge data will be available in class.\n",
    "\n",
    "### 🚧 TODO 6.1: Compare web text vs scientific text\n",
    "\n",
    "Let's understand the key differences between MS MARCO and scientific articles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6700dc0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 6.1: Explore MS MARCO characteristics\n",
    "\n",
    "# ... your code here: compute and plot document length and query length distributions ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c730e39",
   "metadata": {},
   "source": [
    "### Scientific articles vs web passages\n",
    "\n",
    "| Aspect | MS MARCO (web) | Scientific articles |\n",
    "|--------|----------------|---------------------|\n",
    "| Language | Casual, varied | Formal, technical |\n",
    "| Query length | ~6 words | Longer, more specific |\n",
    "| Document length | ~60 words (passage) | 200–500 words (abstract) |\n",
    "| Vocabulary | General English | Domain-specific (\"apoptosis\", \"quasar\") |\n",
    "| Structure | Free text | Title + Abstract + Sections |\n",
    "| Ambiguity | High (\"python\", \"jaguar\") | Low (terms are precise) |\n",
    "\n",
    "### 🚧 TODO 6.2: Implications for retrieval\n",
    "\n",
    "Answer these questions about scientific article retrieval:\n",
    "1. Should you stem scientific terms like \"apoptosis\", \"mitochondria\", \"algorithm\"? Why or why not?\n",
    "2. Are scientific queries more like keyword queries or natural language queries?\n",
    "3. Would you expect BM25 or dense retrieval to be stronger on scientific text? Why?\n",
    "\n",
    "**Answer 6.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eeb9bb6",
   "metadata": {},
   "source": [
    "## Part 7: Domain-Specific Embedding Models (~15 min)\n",
    "\n",
    "General-purpose models like `all-MiniLM-L6-v2` are trained on web text. For scientific articles, domain-specific models can perform better.\n",
    "\n",
    "> 📖 **SPECTER** (Cohan et al., 2020): trained on citation graphs — if paper A cites paper B, their embeddings should be similar. [Paper](https://arxiv.org/abs/2004.07180)\n",
    ">\n",
    "> 📖 **SciBERT** (Beltagy et al., 2019): BERT pre-trained on 1.14M scientific papers. [Paper](https://arxiv.org/abs/1903.10676)\n",
    ">\n",
    "> 📖 **SPECTER2** (Singh et al., 2022): improved SPECTER with adapters for different tasks.\n",
    "\n",
    "### 🚧 TODO 7.1: Compare general vs domain-specific embeddings\n",
    "\n",
    "Test how well different models capture scientific semantics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d0765c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 7.1: Compare embedding models on scientific text pairs\n",
    "\n",
    "# Scientific term pairs that should be similar\n",
    "scientific_pairs = [\n",
    "    (\"programmed cell death\", \"apoptosis\"),\n",
    "    (\"climate change effects on coral\", \"global warming impact on reef ecosystems\"),\n",
    "    (\"machine learning for drug discovery\", \"AI-driven pharmaceutical research\"),\n",
    "    (\"CRISPR gene editing\", \"genome modification using cas9\"),\n",
    "    (\"reinforcement learning\", \"reward-based sequential decision making\"),\n",
    "    (\"black holes\", \"gravitational singularities\"),\n",
    "]\n",
    "\n",
    "# Load SPECTER: model_specter = SentenceTransformer(\"allenai/specter2_base\")\n",
    "# Compare cosine similarity for each pair with both MiniLM and SPECTER\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4615baef",
   "metadata": {},
   "source": [
    "### 🚧 TODO 7.2: Evaluate SPECTER on MS MARCO\n",
    "\n",
    "Even though SPECTER is designed for scientific text, let's see how it does on our MS MARCO queries. This will quantify the \"domain mismatch\" effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c518ea6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 7.2: Dense retrieval with SPECTER on MS MARCO\n",
    "# ⚠️ SPECTER is ~4× slower than MiniLM. On the full 92k docs it would take 30+ minutes!\n",
    "# In FAST_MODE you already have ~10k docs, but SPECTER is still slow.\n",
    "# Use a SUBSET of ~5000 docs instead:\n",
    "#   - Include 1 gold (relevant) doc per train query (so evaluation works)\n",
    "#   - Fill the rest with random docs to reach ~5000 total\n",
    "# Then encode only this subset + all queries with SPECTER.\n",
    "# Remember to map subset indices back to original doc indices when evaluating.\n",
    "\n",
    "# ... your code here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bab7322",
   "metadata": {},
   "source": [
    "**🚧 TODO 7.3:** Answer:\n",
    "1. Does SPECTER perform better or worse than MiniLM on MS MARCO? Why?\n",
    "2. On which domain would you expect SPECTER to outperform MiniLM?\n",
    "3. What is the trade-off of using a domain-specific model?\n",
    "\n",
    "**Answer 7.3:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81161f9",
   "metadata": {},
   "source": [
    "## Part 8: Visualization & Wrap-up\n",
    "\n",
    "### 🚧 TODO 8.1: Visualize the embedding space\n",
    "\n",
    "Use t-SNE to project document embeddings to 2D. Color by a simple topic classification (based on keywords in the document text).\n",
    "\n",
    "> 📖 [t-SNE — Van der Maaten & Hinton (2008)](https://lvdmaaten.github.io/tsne/) — the original paper. Warning: t-SNE preserves local structure but distances between clusters are meaningless."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcd40926",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 8.1: t-SNE visualization\n",
    "\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "# ... your code here:\n",
    "# 1. Sample 2000 documents (don't run t-SNE on all docs!)\n",
    "# 2. Assign a simple topic label based on keywords\n",
    "# 3. Run t-SNE on the sampled embeddings\n",
    "#    Note: use max_iter= (not n_iter= which was renamed in newer scikit-learn)\n",
    "# 4. Plot, colored by topic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cdaf529",
   "metadata": {},
   "source": [
    "**🚧 TODO 8.2:** Looking at the t-SNE plot:\n",
    "1. Do you see topic clusters? Are they well-separated?\n",
    "2. What does it mean if two clusters overlap?\n",
    "3. Why might the \"Other\" category be scattered everywhere?\n",
    "\n",
    "**Answer 8.2:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---\n",
    "\n",
    "### Final Summary\n",
    "\n",
    "Record what you've learned for the in-class session:\n",
    "\n",
    "| System | nDCG@10 (train) | Notes |\n",
    "|--------|:---:|-------|\n",
    "| BM25 (from Class 2) | | Your best k₁, b |\n",
    "| Dense (MiniLM) | | Brute-force |\n",
    "| Dense (SPECTER) | | Domain model |\n",
    "| Dense + FAISS-IVF | | Approximate |\n",
    "\n",
    "You'll need: the dense retrieval code and FAISS index building patterns for the in-class notebook.\n",
    "\n",
    "## Further Reading\n",
    "\n",
    "- 📖 [DPR paper (Karpukhin et al., 2020)](https://arxiv.org/abs/2004.04906) — the first dense retriever to beat BM25 at scale\n",
    "- 📖 [FAISS paper (Johnson, Douze & Jégou, 2019)](https://arxiv.org/abs/1702.08734) — billion-scale similarity search\n",
    "- 📖 [Sentence-BERT (Reimers & Gurevych, 2019)](https://arxiv.org/abs/1908.10084) — how to make BERT useful for similarity\n",
    "- 🎥 [FAISS tutorial (James Briggs)](https://www.youtube.com/watch?v=sKyvsdEv6rk) — 20 min hands-on\n",
    "- 🔧 [FAISS index selection guide](https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index)\n",
    "- 📊 [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) — see how different models rank\n",
    "\n",
    "*Next: In class on March 27 you'll build a complete dense retrieval pipeline, compare it to BM25 head-to-head, and start the scientific article challenge with your first Codabench submission!*"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
