{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c5cb968",
   "metadata": {},
   "source": "# Information Retrieval — Class 5 Prep: Benchmarks (BEIR/MTEB), RAG & Challenge Wrap-up\n\n**Due:** April 10, 2026 (before class)\n\n**Goal:** Evaluate your retrieval pipeline across multiple domains using BEIR/MTEB benchmarks, understand the domain gap, build a Retrieval-Augmented Generation (RAG) pipeline, and explore the state of the art in retrieval (2025–2026).\n\n## Notebook Outline\n\n| Part | Topic | Description |\n|------|-------|-------------|\n| **1** | Data & Baselines | Load MS MARCO, build BM25 + dense baselines from Classes 2–4 |\n| **2** | IR Benchmarks: BEIR | Cross-domain evaluation on SciFact, NFCorpus, ArguAna |\n| **3** | IR Benchmarks: MTEB | MTEB leaderboard and local evaluation |\n| **4** | Retrieval-Augmented Generation (RAG) | Build a complete RAG pipeline with retriever + generator |\n| **5** | Chunking Strategies | Fixed-size, sentence-based, semantic chunking |\n| **6** | The State of the Art (2025–2026) | MTEB models, Matryoshka embeddings, LLM rerankers |\n\n**Key Concepts:**\n- BEIR: zero-shot cross-domain evaluation for retrieval\n- MTEB: massive text embedding benchmark and leaderboard\n- Domain gap: why MS MARCO performance ≠ scientific/financial/medical performance\n- RAG: retriever + generator architecture for QA\n- Chunking strategies for long documents\n- Matryoshka Representation Learning (MRL)\n\n**Required Libraries:** `sentence-transformers`, `faiss-cpu`, `rank_bm25`, `transformers`, `torch`, `beir`, `mteb`\n\n> 📌 **Installation:**\n> ```bash\n> pip install sentence-transformers faiss-cpu rank_bm25 beir mteb\n> ```\n\n**Essential Reading & Resources:**\n\n| Resource | Link |\n|----------|------|\n| Lewis et al. (2020), *\"Retrieval-Augmented Generation\"* | [Paper](https://arxiv.org/abs/2005.11401) |\n| Thakur et al. (2021), *\"BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of IR Models\"* | [Paper](https://arxiv.org/abs/2104.08663) |\n| Muennighoff et al. (2022), *\"MTEB: Massive Text Embedding Benchmark\"* | [Paper](https://arxiv.org/abs/2210.07316) |\n| MTEB Leaderboard (Hugging Face) | [Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) |\n| Kusupati et al. (2022), *\"Matryoshka Representation Learning\"* | [Paper](https://arxiv.org/abs/2205.13147) |\n| Barnett et al. (2024), *\"Seven Failure Points When Engineering a RAG System\"* | [Paper](https://arxiv.org/abs/2401.05856) |\n| BEIR GitHub & Documentation | [GitHub](https://github.com/beir-cellar/beir) |\n| Video: RAG Explained (IBM Technology) | [YouTube](https://www.youtube.com/watch?v=T-D1OfcDW1M) |\n| Video: BEIR Benchmark (Nils Reimers) | [YouTube](https://www.youtube.com/watch?v=gFPfSYJlk_Q) |\n| LangChain RAG Tutorial | [Docs](https://python.langchain.com/docs/tutorials/rag/) |\n| Manning, Raghavan & Schütze — *Introduction to Information Retrieval* | [Book (free)](https://nlp.stanford.edu/IR-book/) |"
  },
  {
   "cell_type": "markdown",
   "id": "93f2685d",
   "metadata": {},
   "source": [
    "<span style=\"color:magenta\">Student Name:</span>\n",
    "\n",
    "* Name\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "448bd95b",
   "metadata": {},
   "outputs": [],
   "source": "# !pip install sentence-transformers faiss-cpu rank_bm25 beir mteb\nimport os\nimport platform\nimport zipfile\nimport urllib.request\nimport time\nimport json\nfrom pathlib import Path\n\n# Force Hugging Face libraries to use cached models only (no network access).\n# This avoids SSL errors on networks with corporate proxies / self-signed certs.\nos.environ[\"HF_HUB_OFFLINE\"] = \"1\"\nos.environ[\"TRANSFORMERS_OFFLINE\"] = \"1\"\n\n# ── macOS memory fix ──────────────────────────────────────────────────────────\n# The nano malloc zone on Apple Silicon causes memory fragmentation that can\n# crash the kernel during SentenceTransformer encoding.  Setting this env var\n# BEFORE starting Jupyter prevents the issue.\nif platform.system() == \"Darwin\" and os.environ.get(\"MALLOC_NANO_ZONE\") != \"0\":\n    print(\"⚠️  macOS detected but MALLOC_NANO_ZONE is not set in your shell.\")\n    print(\"   To prevent kernel crashes, run this in your terminal, then restart Jupyter:\")\n    print(\"     echo 'export MALLOC_NANO_ZONE=0' >> ~/.zshrc && source ~/.zshrc\")\nos.environ[\"MALLOC_NANO_ZONE\"] = \"0\"  # helps child processes\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nfrom tqdm import tqdm\n\n# Limit PyTorch threads to prevent malloc arena fragmentation on macOS\nimport torch\ntorch.set_num_threads(1)\n\ntqdm.pandas()\nsns.set_style(\"darkgrid\")\nsns.set_context(\"notebook\")\npd.set_option(\"display.max_colwidth\", 120)"
  },
  {
   "cell_type": "markdown",
   "id": "2b46dee5",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Part 1: Data & Baselines (from Classes 2–4)\n",
    "\n",
    "### 1.1 Data Loading & Evaluation Toolkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "040502d9",
   "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": "8b77d3d0",
   "metadata": {},
   "outputs": [],
   "source": "queries = pd.read_csv(\"our.msmarco.queries.tsv\", sep=\"\\t\", usecols=[1, 2])\ngold = 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)\nFAST_MODE = False\n\nif FAST_MODE:\n    train_queries = queries.iloc[:100].copy()\n    test_queries  = queries.iloc[500:600].copy()\nelse:\n    train_queries = queries.iloc[:500].copy()\n    test_queries  = queries.iloc[500:].copy()\n\nall_qids = set(train_queries[\"qid\"]) | set(test_queries[\"qid\"])\ngold = gold[gold[\"qid\"].isin(all_qids)]\ngold_dict = gold.groupby(\"qid\")[\"docid\"].apply(set).to_dict()\ngold_docids = set(gold[\"docid\"])\n\n# ── Read docs in chunks to avoid 3 GB peak from loading all 92K rows ──\n# On macOS, freed memory is NOT returned to the OS, so peak RSS matters.\nimport gc\n\nMAX_DOC_WORDS = 200  # MiniLM truncates at 256 tokens (~200 words)\ngold_chunks = []\nother_chunks = []\nother_budget = 3000  # max distractor docs (keep low to limit memory on 24 GB Mac)\n\nfor chunk in pd.read_csv(\"our.msmarco.docs.tsv\", sep=\"\\t\", usecols=[1, 2, 3], chunksize=5000):\n    # Merge text and truncate INSIDE the loop to keep each chunk small\n    chunk[\"text\"] = (chunk[\"title\"].fillna(\"\") + \" \" + chunk[\"body\"].fillna(\"\")).apply(\n        lambda t: \" \".join(t.split()[:MAX_DOC_WORDS])\n    )\n    chunk.drop(columns=[\"title\", \"body\"], inplace=True)\n\n    mask = chunk[\"docid\"].isin(gold_docids)\n    gold_chunks.append(chunk[mask])\n    if other_budget > 0:\n        non_gold = chunk[~mask]\n        take = min(len(non_gold), other_budget)\n        other_chunks.append(non_gold.head(take))\n        other_budget -= take\n\ndocs = pd.concat(gold_chunks + other_chunks, ignore_index=True)\n\ndel gold_chunks, other_chunks\ngc.collect()\n\nmode = \"FAST\" if FAST_MODE else \"FULL\"\nprint(f\"[{mode}] Queries: {len(train_queries)}+{len(test_queries)}, Docs: {len(docs)}\")\nprint(f\"✓ Chunked loading + text truncated to {MAX_DOC_WORDS} words\")"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a83abadf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Evaluation metrics (from Classes 1-4) ---\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": "9ff6f3b2",
   "metadata": {},
   "source": [
    "### 1.2 Build BM25 & Dense Baselines (from Classes 2–4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bfb24e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this takes more than 17 minutes on the full dataset on my machine. You can set FAST_MODE = True to run on a smaller subset for quick iteration.\n",
    "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": "93e09c65",
   "metadata": {},
   "outputs": [],
   "source": "from sentence_transformers import SentenceTransformer, CrossEncoder\nimport faiss\nimport gc\n\n# Dense retrieval baseline\nmodel = SentenceTransformer(\"all-MiniLM-L6-v2\")\nd = model.get_sentence_embedding_dimension()\n\nprint(\"Encoding documents...\")\ndoc_texts = docs[\"text\"].tolist()\n# normalize_embeddings=True avoids allocating a second full-size array for normalization\n# batch_size=64 reduces peak PyTorch allocation (prevents macOS malloc fragmentation)\ndoc_embeddings_norm = model.encode(doc_texts, show_progress_bar=True, batch_size=64,\n                                   normalize_embeddings=True)\n\nprint(\"Encoding queries...\")\nquery_texts = queries[\"query\"].tolist()\nquery_embeddings_norm = model.encode(query_texts, show_progress_bar=True, batch_size=64,\n                                     normalize_embeddings=True)\n\n# Build FAISS index (used early in the notebook — safe at this point)\nindex_flat = faiss.IndexFlatIP(d)\nindex_flat.add(doc_embeddings_norm.astype(np.float32))\n\ndef dense_retrieve(query_text, k=100):\n    \"\"\"Dense retrieval using numpy dot product (avoids FAISS SIGSEGV on macOS after cumulative memory use).\"\"\"\n    qvec = model.encode([query_text], normalize_embeddings=True).astype(np.float32)\n    scores = (qvec @ doc_embeddings_norm.astype(np.float32).T)[0]\n    top_k = np.argsort(-scores)[:k]\n    return top_k, scores[top_k]\n\ndense_results = evaluate_retriever(lambda q, k=10: dense_retrieve(q, k), train_queries, gold_dict, docs)\nprint_results(\"Dense Baseline (MiniLM)\", dense_results)"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73a805ad",
   "metadata": {},
   "outputs": [],
   "source": "# Cross-encoder and full pipeline from Class 4\ncross_encoder = CrossEncoder(\"cross-encoder/ms-marco-MiniLM-L-6-v2\")\n\ndef rerank_with_cross_encoder(query_text, candidate_indices, docs_df, ce_model, k=10):\n    pairs = [[query_text, docs_df.iloc[idx][\"text\"][:512]] for idx in candidate_indices]\n    scores = ce_model.predict(pairs)\n    top_k = np.argsort(scores)[::-1][:k]\n    return candidate_indices[top_k], scores[top_k]\n\ndef reciprocal_rank_fusion(rankings, k=60):\n    rrf_scores = {}\n    for ranking in rankings:\n        for rank, item in enumerate(ranking, 1):\n            doc_id = item[0] if isinstance(item, (tuple, list)) else item\n            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + 1.0 / (k + rank)\n    return sorted(rrf_scores.items(), key=lambda x: -x[1])\n\nprint(\"✓ Cross-encoder and RRF loaded from Class 4\")\n\n# --- Memory cleanup before Part 2 ---\nimport gc\n\n# Delete large temporaries no longer needed\ndel tokenized_docs, doc_texts, query_texts, bm25\n\n# Ensure embeddings are float32 (encode may return float32 already, but be safe)\ndoc_embeddings_norm = doc_embeddings_norm.astype(np.float32)\nquery_embeddings_norm = query_embeddings_norm.astype(np.float32)\n\ngc.collect()\nprint(f\"✓ Memory cleanup done. docs columns: {list(docs.columns)}, embeddings dtype: {doc_embeddings_norm.dtype}\")\nprint(\"  Note: MS MARCO BM25 index will be rebuilt in Part 4 for RAG\")"
  },
  {
   "cell_type": "markdown",
   "id": "a192acc1",
   "metadata": {},
   "source": "---\n\n## Part 2: IR Benchmarks — BEIR (~30 min)\n\n### 2.1 Why Cross-Domain Evaluation?\n\nWe've spent 4 classes tuning our pipeline on **MS MARCO** — a web-search dataset. But does it generalize?\n\nThe **BEIR benchmark** (Thakur et al., 2021) answers this question by evaluating retrieval models on **18 diverse datasets** in a **zero-shot** setting (no training on the target domain).\n\n> 📖 The BEIR finding that shocked the IR community: **BM25 outperforms most dense models** in zero-shot cross-domain evaluation! Dense models trained on MS MARCO often overfit to web-search style queries.\n\n| Dataset | Domain | # Docs | # Queries | Why it matters |\n|---------|--------|--------|-----------|----------------|\n| MS MARCO | Web search | 8.8M | 6,980 | Our \"home turf\" — we know this one |\n| SciFact | Scientific claims | 5,183 | 300 | Can you verify claims from abstracts? |\n| NFCorpus | Nutrition/fitness | 3,633 | 323 | Medical domain, specialized vocabulary |\n| ArguAna | Argument retrieval | 8,674 | 1,406 | Find counter-arguments |\n| TREC-COVID | COVID-19 literature | 171K | 50 | Pandemic-era scientific search |\n\n> 📖 Read: [BEIR paper (Thakur et al., 2021)](https://arxiv.org/abs/2104.08663) — especially Sections 3 and 4 on the benchmark design and main results.\n>\n> 📖 [BEIR GitHub](https://github.com/beir-cellar/beir) — explore the codebase before the in-class session.\n\n### 🚧 TODO 2.1: Load and explore BEIR datasets\n\nUse the `beir` library to load SciFact, NFCorpus, and ArguAna. Explore the structure of each dataset."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18ce33f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 2.1: Load BEIR datasets\n",
    "from beir import util\n",
    "from beir.datasets.data_loader import GenericDataLoader\n",
    "\n",
    "# Download and load SciFact\n",
    "# Use beir's util.download_and_unzip() and GenericDataLoader().load(split=\"test\")\n",
    "# Print corpus size, number of queries, sample query and sample document\n",
    "\n",
    "# ... your code here ...\n",
    "# sf_corpus, sf_queries, sf_qrels = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13ca86a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load NFCorpus\n",
    "\n",
    "# ... your code here ...\n",
    "# nf_corpus, nf_queries, nf_qrels = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531a1b11",
   "metadata": {},
   "outputs": [],
   "source": "# Load ArguAna\n\n# ... your code here ...\n# ar_corpus, ar_queries, ar_qrels = ..."
  },
  {
   "cell_type": "markdown",
   "id": "684179e4",
   "metadata": {},
   "source": [
    "### 🚧 TODO 2.2: Evaluate BM25 on BEIR datasets\n",
    "\n",
    "Build a BM25 retriever for each BEIR dataset and measure nDCG@10. This shows how BM25 transfers across domains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "719e55ea",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 2.2: BM25 on BEIR datasets\n\nimport gc\n\nMAX_BEIR_WORDS = 200  # Truncate BEIR docs — matches MS MARCO truncation\n\ndef evaluate_bm25_on_beir(corpus, queries, qrels, k=10):\n    \"\"\"Evaluate BM25 on a BEIR dataset.\n    \n    Args:\n        corpus: dict {docid: {title, text}}\n        queries: dict {qid: query_text}\n        qrels: dict {qid: {docid: relevance}}\n    Returns:\n        dict with nDCG@10, MRR, MAP\n    \"\"\"\n    # ... your code here ...\n    # 1. Prepare doc_ids and doc_texts from corpus (truncate to MAX_BEIR_WORDS)\n    #      doc_texts = [\" \".join((corpus[did].get(\"title\",\"\") + \" \" + corpus[did].get(\"text\",\"\")).split()[:MAX_BEIR_WORDS]) ...]\n    # 2. Tokenize with preprocess() and build BM25Okapi index\n    # 3. Free tokenized list and doc_texts after building the index:\n    #      del tokenized, doc_texts; gc.collect()\n    # 4. For each query: get BM25 scores, find top-k, compute nDCG@10, MRR, MAP\n    # 5. Free bm25_index before returning:\n    #      del bm25_index; gc.collect()\n    # 6. Return mean scores\n    pass\n\n# Evaluate on all three datasets and print results\n# sf_bm25 = evaluate_bm25_on_beir(sf_corpus, sf_queries, sf_qrels)\n# nf_bm25 = evaluate_bm25_on_beir(nf_corpus, nf_queries, nf_qrels)\n# ar_bm25 = evaluate_bm25_on_beir(ar_corpus, ar_queries, ar_qrels)\n\n# ... your code here ..."
  },
  {
   "cell_type": "markdown",
   "id": "7897c821",
   "metadata": {},
   "source": [
    "### 🚧 TODO 2.3: Evaluate dense retrieval on BEIR datasets\n",
    "\n",
    "Now evaluate your MiniLM dense retriever on the same datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f415573",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 2.3: Dense retrieval on BEIR\n\ndef evaluate_dense_on_beir(corpus, queries, qrels, sent_model, k=10):\n    \"\"\"Evaluate dense retrieval (sentence-transformers) on a BEIR dataset.\n    \n    Args:\n        corpus: dict {docid: {title, text}}\n        queries: dict {qid: query_text}\n        qrels: dict {qid: {docid: relevance}}\n        sent_model: SentenceTransformer model\n    Returns:\n        dict with nDCG@10, MRR, MAP\n    \"\"\"\n    # ... your code here ...\n    # 1. Prepare doc_ids and doc_texts from corpus (truncate to MAX_BEIR_WORDS)\n    #      doc_texts = [\" \".join((corpus[did].get(\"title\",\"\") + \" \" + corpus[did].get(\"text\",\"\")).split()[:MAX_BEIR_WORDS]) ...]\n    # 2. Encode documents with normalize_embeddings=True, batch_size=64, .astype(np.float32)\n    # 3. Encode queries  with normalize_embeddings=True, batch_size=64, .astype(np.float32)\n    # 4. Cosine similarity via numpy dot product (both are L2-normalized):\n    #      scores = q_embs @ doc_embs.T\n    #      I = np.argsort(-scores, axis=1)[:, :k]\n    #    (We use numpy instead of FAISS here to avoid a macOS memory bug)\n    # 5. Compute nDCG@10, MRR, MAP per query\n    # 6. Return mean scores. Remember: del + gc.collect() after search\n    pass\n\n# Evaluate on all three datasets and print results\n# sf_dense = evaluate_dense_on_beir(sf_corpus, sf_queries, sf_qrels, model)\n# nf_dense = evaluate_dense_on_beir(nf_corpus, nf_queries, nf_qrels, model)\n# ar_dense = evaluate_dense_on_beir(ar_corpus, ar_queries, ar_qrels, model)\n\n# ... your code here ..."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "973c81e8",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 2.3b: Compare BM25 vs Dense across all datasets\n# Create a comparison table and bar chart showing nDCG@10 for each dataset\n# Datasets: MS MARCO (train), SciFact, NFCorpus, ArguAna\n\n# ... your code here ...\n# Use matplotlib to create a grouped bar chart"
  },
  {
   "cell_type": "markdown",
   "id": "87ece70f",
   "metadata": {},
   "source": [
    "### 🚧 TODO 2.4: Analysis questions\n",
    "\n",
    "1. On which dataset(s) does BM25 outperform dense retrieval? Why?\n",
    "2. On which dataset(s) does dense retrieval win? Why?\n",
    "3. What is the biggest \"domain gap\" — i.e., the biggest drop from MS MARCO performance to another dataset?\n",
    "4. Based on these results, would you trust your MS MARCO nDCG@10 as a measure of retrieval quality for a **new** domain?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5418fbad",
   "metadata": {},
   "source": [
    "**Answer 2.4:**\n",
    "\n",
    "... your answer here ...\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ab43cb",
   "metadata": {},
   "source": [
    "## Part 3: IR Benchmarks — MTEB (~20 min)\n",
    "\n",
    "### 3.1 What is MTEB?\n",
    "\n",
    "While BEIR focuses on **retrieval** across domains, **MTEB** (Muennighoff et al., 2022) is a broader benchmark that evaluates embeddings across **8 task categories**:\n",
    "\n",
    "| Category | Example task | What it measures |\n",
    "|----------|-------------|-----------------|\n",
    "| **Retrieval** | MS MARCO, SciFact | Finding relevant documents |\n",
    "| **STS** | STS Benchmark | Semantic text similarity |\n",
    "| **Classification** | Amazon reviews | Embedding → classifier |\n",
    "| **Clustering** | Reddit posts | Grouping similar texts |\n",
    "| **Pair Classification** | Paraphrase detection | Same/different meaning |\n",
    "| **Reranking** | AskUbuntu | Reorder candidates |\n",
    "| **Summarization** | SummEval | Summary quality embeddings |\n",
    "| **BitextMining** | Tatoeba | Cross-language matching |\n",
    "\n",
    "The MTEB Leaderboard on Hugging Face ranks embedding models across all these tasks: https://huggingface.co/spaces/mteb/leaderboard\n",
    "\n",
    "> 📖 Read: [MTEB paper (Muennighoff et al., 2022)](https://arxiv.org/abs/2210.07316)\n",
    ">\n",
    "> 📖 Explore the [MTEB leaderboard](https://huggingface.co/spaces/mteb/leaderboard) — filter by \"Retrieval\" and note which models are on top.\n",
    "\n",
    "### Current top models (early 2026):\n",
    "\n",
    "| Rank | Model | Size | Avg Score | Retrieval Score |\n",
    "|------|-------|------|-----------|-----------------|\n",
    "| 1 | NV-Embed-v2 | 7.8B | 72.3 | 62.7 |\n",
    "| 2 | voyage-3 | ? | 71.8 | 61.3 |\n",
    "| 3 | Cohere embed-v3 | ? | 71.5 | 60.9 |\n",
    "| 4 | GTE-Qwen2-7B | 7B | 70.2 | 60.2 |\n",
    "| ... | all-MiniLM-L6-v2 | 22M | 56.3 | 41.9 |\n",
    "\n",
    "Our MiniLM model is **350× smaller** than the top models! Yet with BM25 + RRF + cross-encoder, it's competitive.\n",
    "\n",
    "### 🚧 TODO 3.1: Run a small MTEB evaluation locally\n",
    "\n",
    "Use the `mteb` library to evaluate MiniLM on one retrieval task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c427b0d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🚧 TODO 3.1: MTEB evaluation\n",
    "\n",
    "import mteb\n",
    "\n",
    "# ... your code here ...\n",
    "# 1. Load MiniLM with mteb.get_model(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "# 2. Get the SciFact retrieval task with mteb.get_tasks(tasks=[\"SciFact\"], languages=[\"eng\"])\n",
    "# 3. Create an MTEB evaluation and run it\n",
    "# 4. Load the JSON results and print the nDCG@10 score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37339dca",
   "metadata": {},
   "source": "### 3.2 Compare two embedding models on MTEB (optional)\n\nIf you have time, evaluate a second model — e.g., `BAAI/bge-small-en-v1.5` or `thenlper/gte-small` — and compare with MiniLM on SciFact."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0520c1",
   "metadata": {},
   "outputs": [],
   "source": "# Optional: Compare models on MTEB\n# Uncomment and fill in if you have time\n\n# eval_model_2 = mteb.get_model(\"BAAI/bge-small-en-v1.5\")\n# tasks = mteb.get_tasks(tasks=[\"SciFact\"], languages=[\"eng\"])\n# evaluation = mteb.MTEB(tasks=tasks)\n# results_2 = evaluation.run(eval_model_2, output_folder=\"mteb_results/bge_small\")"
  },
  {
   "cell_type": "markdown",
   "id": "df559339",
   "metadata": {},
   "source": [
    "### 🚧 TODO 3.3: MTEB leaderboard analysis\n",
    "\n",
    "Go to https://huggingface.co/spaces/mteb/leaderboard and answer:\n",
    "\n",
    "1. What is the top-ranked model for the \"Retrieval\" category? What is its size (number of parameters)?\n",
    "2. What is the smallest model in the top 20 for retrieval? How does its score compare to the largest?\n",
    "3. Is there a strong correlation between model size and retrieval performance?\n",
    "4. Find a model that is good at retrieval but bad at classification (or vice versa). What does this tell you about embedding generalization?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ecf916",
   "metadata": {},
   "source": [
    "**Answer 3.3:**\n",
    "\n",
    "... your answer here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2d639c",
   "metadata": {},
   "source": "## Part 4: Retrieval-Augmented Generation — RAG (~30 min)\n\n### 4.1 What is RAG?\n\nEverything we've built so far does **retrieval** — finding relevant documents. **RAG** goes one step further: it feeds the retrieved passages to a language model to **generate an answer**.\n\n```\nUser: \"What causes coral bleaching?\"\n    ↓\n[Retriever] → top-3 relevant passages from the corpus\n    ↓\n[Prompt Builder]:\n  \"Given the following passages:\n   1. Coral bleaching occurs when rising water temperatures cause...\n   2. Ocean acidification from CO2 absorption weakens coral...\n   3. Pollutants and agricultural runoff damage reef ecosystems...\n   \n   Answer the question: What causes coral bleaching?\"\n    ↓\n[LLM Generator] → \"The main causes of coral bleaching are elevated \n   water temperatures, ocean acidification from CO2, and pollution \n   from agricultural runoff.\"\n```\n\n**Why RAG matters:**\n- Reduces LLM **hallucination** by grounding answers in evidence\n- Every major AI assistant uses RAG: Bing Chat, Perplexity, Google AI Overviews\n- The retriever we built is the **critical component** — garbage in, garbage out!\n\n> 📖 Read: [Lewis et al. (2020)](https://arxiv.org/abs/2005.11401) — the original RAG paper from Facebook AI.\n>\n> 🎥 Watch: [RAG Explained (IBM Technology)](https://www.youtube.com/watch?v=T-D1OfcDW1M) — 10 min overview.\n>\n> 📖 [Barnett et al. (2024)](https://arxiv.org/abs/2401.05856) — Seven Failure Points When Engineering a RAG System.\n\n### 🚧 TODO 4.1: Build a simple RAG pipeline on MS MARCO\n\nUse your dense retriever + a small LM (flan-t5-small) to answer questions."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5454f4b",
   "metadata": {},
   "outputs": [],
   "source": "# Rebuild BM25 index for MS MARCO (was freed in Part 1 to save memory)\nprint(\"Rebuilding BM25 index for MS MARCO...\")\ntokenized_docs = docs[\"text\"].progress_apply(preprocess).tolist()\nbm25 = BM25Okapi(tokenized_docs, k1=1.2, b=0.75)\ndel tokenized_docs\ngc.collect()\n\ndef 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\nprint(\"✓ BM25 index rebuilt\")\n\n# 🚧 TODO 4.1: Build a RAG pipeline\n\nfrom transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM\n\n# Load a small generator model — flan-t5-small is great for QA\nprint(\"Loading generator model (flan-t5-small)...\")\ngenerator = pipeline(\"text2text-generation\", model=\"google/flan-t5-small\", max_new_tokens=128)\n\ndef rag_answer(question, retriever_fn, docs_df, k=3):\n    \"\"\"Simple RAG: retrieve top-k passages, build a prompt, generate an answer.\n    \n    Args:\n        question: the user's question\n        retriever_fn: function(query_text, k) → (indices, scores)\n        docs_df: DataFrame with document text\n        k: number of passages to retrieve\n    Returns:\n        answer, retrieved_passages\n    \"\"\"\n    # ... your code here ...\n    # Step 1: Retrieve top-k passages using retriever_fn\n    # Step 2: Build prompt with context passages\n    # Step 3: Generate answer using the generator pipeline\n    # Return: (generated_answer, list_of_passages)\n    pass\n\n# Test with some questions\ntest_questions = [\n    \"what causes rain\",\n    \"who invented the telephone\",\n    \"how does photosynthesis work\",\n    \"what is the capital of France\",\n]\n\n# ... your code here ...\n# Call rag_answer for each test question and print results"
  },
  {
   "cell_type": "markdown",
   "id": "4e9e6511",
   "metadata": {},
   "source": "### 🚧 TODO 4.2: RAG failure modes\n\nAnswer these questions:\n\n1. What happens when the retriever returns irrelevant passages? Does the LLM refuse to answer, or does it hallucinate?\n2. Try a question where NO document in MS MARCO is relevant (e.g., something very recent). What does the RAG system do?\n3. What if two retrieved passages **contradict** each other? How does the LLM handle it?\n4. Why is the retriever the \"critical component\" of RAG?"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a43df85",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 4.2: Explore RAG failure modes\n\n# ... your code here ...\n# Test 1: Question with no relevant documents (e.g., \"what is the GDP of Mars in 2025\")\n# Test 2: Ambiguous question with potentially contradictory passages (e.g., \"is coffee good for you\")\n# Print the question, generated answer, and retrieved passages for each test"
  },
  {
   "cell_type": "markdown",
   "id": "afdd607b",
   "metadata": {},
   "source": "**Answer 4.2:**\n\n... your answer here ...\n\n---"
  },
  {
   "cell_type": "markdown",
   "id": "18b82a68",
   "metadata": {},
   "source": "---\n\n## Part 5: Chunking Strategies (~15 min)\n\n### 🚧 TODO 5.1: Implement chunking strategies\n\nImplement three chunking strategies: fixed-size with overlap, sentence-based, and semantic chunking."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1a8b8db",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 5.1: Implement chunking strategies\n\nimport re\n\ndef chunk_fixed_size(text, chunk_size=200, overlap=50):\n    \"\"\"Split text into fixed-size word chunks with overlap.\n    \n    Args:\n        text: input text\n        chunk_size: number of words per chunk\n        overlap: number of overlapping words between chunks\n    Returns:\n        list of chunk strings\n    \"\"\"\n    # ... your code here ...\n    pass\n\ndef chunk_sentences(text, max_chunk_size=200):\n    \"\"\"Split text at sentence boundaries, grouping sentences into chunks.\n    \n    Args:\n        text: input text\n        max_chunk_size: maximum words per chunk\n    Returns:\n        list of chunk strings\n    \"\"\"\n    # ... your code here ...\n    # Hint: use re.split(r'(?<=[.!?])\\s+', text) to split sentences\n    pass\n\ndef chunk_semantic(text, sent_model, threshold=0.5, max_chunk_size=200):\n    \"\"\"Split text using embedding similarity between consecutive sentences.\n    \n    Args:\n        text: input text\n        sent_model: SentenceTransformer model for encoding\n        threshold: similarity threshold below which to split\n        max_chunk_size: maximum words per chunk\n    Returns:\n        list of chunk strings\n    \"\"\"\n    # ... your code here ...\n    # 1. Split text into sentences\n    # 2. Encode each sentence with sent_model\n    # 3. Compute cosine similarity between consecutive sentences\n    # 4. Split where similarity < threshold or chunk exceeds max_chunk_size\n    pass\n\n# Test on a longer document\nlong_doc = docs.iloc[0][\"text\"]\nprint(f\"Original document: {len(long_doc.split())} words\\n\")\n\n# ... your code here ...\n# Test all three chunking strategies and compare:\n# - Number of chunks\n# - Chunk sizes\n# - First 2 chunks from each strategy"
  },
  {
   "cell_type": "markdown",
   "id": "f5f2fc6a",
   "metadata": {},
   "source": "### 🚧 TODO 5.2: Chunking discussion\n\nAnswer:\n1. What chunk size works best for RAG with a small LM (flan-t5-small with 512 token context)?\n2. Why does overlapping chunking help? What information could be lost without overlap?\n3. When would **sentence-based** chunking be better than fixed-size? When would it be worse?\n4. For scientific articles that have sections (Abstract, Methods, Results, Discussion), what chunking strategy would you use?"
  },
  {
   "cell_type": "markdown",
   "id": "8c263a8d",
   "metadata": {},
   "source": "**Answer 5.2:**\n\n... your answer here ...\n\n---"
  },
  {
   "cell_type": "markdown",
   "id": "e7ca7b4d",
   "metadata": {},
   "source": "## Part 6: The State of the Art (2025–2026) (~15 min)\n\n### 6.1 Matryoshka Representation Learning\n\nOne of the most practical innovations in recent embedding research:\n\n**Matryoshka embeddings** (Kusupati et al., 2022) train models so that **prefixes** of the embedding vector are useful:\n\n$$\\vec{v} = [v_1, v_2, \\ldots, v_{768}]$$\n\n- Use $[v_1, \\ldots, v_{64}]$ for coarse search (fast, 12× less storage)\n- Use $[v_1, \\ldots, v_{256}]$ for medium search\n- Use $[v_1, \\ldots, v_{768}]$ for fine search\n\nLike Russian nesting dolls 🪆 — each smaller representation is nested inside the larger one.\n\n| Dimensions | Storage (1M docs) | Typical Recall@100 |\n|-----------|-------------------|------------|\n| 768 | 3.0 GB | 95.2% |\n| 256 | 1.0 GB | 94.1% |\n| 64 | 0.25 GB | 89.7% |\n\n> 📖 [Kusupati et al. (2022)](https://arxiv.org/abs/2205.13147) — Matryoshka Representation Learning\n\n### 🚧 TODO 6.1: Test Matryoshka embeddings\n\nUse a Matryoshka-trained model and compare retrieval quality at different dimensions."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33566ce",
   "metadata": {},
   "outputs": [],
   "source": "# 🚧 TODO 6.1: Matryoshka embeddings experiment\n\n# ... your code here ...\n# Demonstrate Matryoshka-style dimension reduction:\n# 1. Use doc_embeddings_norm (384 dimensions for MiniLM)\n# 2. Test dimensions: [32, 64, 128, 256, 384]\n# 3. For each dimension: truncate, re-normalize\n# 4. Use numpy dot product for search (avoids FAISS crash on macOS):\n#      scores = query_trunc @ doc_trunc.T\n#      I = np.argsort(-scores, axis=1)[:, :10]\n# 5. Compute Recall@10 vs ground truth (full-dimension results)\n# 6. Print: dimension, Recall@10, storage in MB\n# Note: MiniLM is NOT trained with MRL; a true Matryoshka model would perform better"
  },
  {
   "cell_type": "markdown",
   "id": "8e51fe08",
   "metadata": {},
   "source": "### 6.2 Instruction-Tuned Embeddings\n\nA major trend in 2024–2026: **tell the model what type of search you want**. Models like E5-Mistral and GTE-Qwen accept task-specific instructions:\n\n```python\n# Without instruction\nmodel.encode(\"programmed cell death\")  # generic embedding\n\n# With instruction\nmodel.encode(\"Retrieve scientific papers about: programmed cell death\")  # retrieval-specific embedding\nmodel.encode(\"Classify this text: programmed cell death\")  # classification-specific embedding\n```\n\n> 📖 [E5-Mistral-7B-Instruct paper](https://arxiv.org/abs/2401.00368) — instruction-tuned embeddings for retrieval\n\n### 6.3 The Landscape: Commercial vs Open-Source\n\n| Category | Commercial leaders | Open-source alternative |\n|----------|-------------------|----------------------|\n| Embedding API | OpenAI, Cohere, Voyage | BGE, GTE, E5 (free) |\n| Vector DB | Pinecone (managed) | FAISS, Qdrant, Weaviate |\n| LLM for RAG | GPT-4, Claude | Llama 3, Mistral, Qwen |\n| End-to-end RAG | Perplexity.ai | LangChain, LlamaIndex |\n| Re-ranker API | Cohere Rerank | Cross-encoder (free) |\n\nFor research and this class: open-source is enough and often better!\n\n### 🚧 TODO 6.2: State-of-the-art quiz\n\nAnswer these questions (use the MTEB leaderboard and the links above):\n\n1. What is the current #1 model on the MTEB Retrieval leaderboard? How many parameters does it have?\n2. What is the best model with fewer than 500M parameters?\n3. Name one advantage of commercial embedding APIs (like Cohere embed-v3) over open-source models.\n4. Name one advantage of open-source models over commercial APIs.\n5. If you were building a RAG system for a startup with limited budget, what stack would you choose?"
  },
  {
   "cell_type": "markdown",
   "id": "52ee0a88",
   "metadata": {},
   "source": "**Answer 6.2:**\n\n... your answer here ...\n\n---"
  },
  {
   "cell_type": "markdown",
   "id": "d15f4f54",
   "metadata": {},
   "source": "## Summary\n\n### Key Takeaways\n\n1. **MS MARCO performance ≠ real-world performance.** Cross-domain evaluation (BEIR/MTEB) is essential.\n2. **BM25 is surprisingly hard to beat** in zero-shot cross-domain settings.\n3. **RAG = Retriever + Generator.** The retriever is the critical component — garbage in, garbage out.\n4. **Chunking matters.** How you split documents dramatically affects RAG quality.\n5. **State of the art is moving fast.** Matryoshka embeddings, instruction-tuned models, LLM re-rankers.\n\n### For the In-Class Session:\n- You'll discuss BEIR/MTEB results as a class\n- Build a mini-RAG system for MS MARCO\n- Make your final challenge submission with all insights applied\n\n### Further Reading\n\n- 📖 [Lewis et al. (2020)](https://arxiv.org/abs/2005.11401) — RAG (**required reading**)\n- 📖 [Thakur et al. (2021)](https://arxiv.org/abs/2104.08663) — BEIR benchmark\n- 📖 [Muennighoff et al. (2022)](https://arxiv.org/abs/2210.07316) — MTEB benchmark\n- 📖 [Kusupati et al. (2022)](https://arxiv.org/abs/2205.13147) — Matryoshka embeddings\n- 📖 [Barnett et al. (2024)](https://arxiv.org/abs/2401.05856) — Seven RAG failure points\n- 🎥 [RAG Explained (IBM Technology)](https://www.youtube.com/watch?v=T-D1OfcDW1M) — 10 min\n- 🔧 [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) — live rankings\n- 🔧 [BEIR GitHub](https://github.com/beir-cellar/beir) — code & data\n\n*Next: In class on April 10 you'll present your benchmark results, build a RAG demo, and make your final challenge submission!*"
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}