{ "cells": [ { "cell_type": "markdown", "id": "dfe380cf", "metadata": { "id": "dfe380cf" }, "source": [ "# CS145 NanoMem\n", "\n", "**Starter release: October 16, 2026.**\n", "\n", "**Where to run:**\n", "\n", "1. **Colab (recommended):** Enable a GPU if available: **Runtime → Change runtime type → GPU**, then set `EMBED_DEVICE = \"cuda\"` in Configuration. CPU also works; leave `EMBED_DEVICE = \"cpu\"`.\n", "2. **Local Jupyter:** If you have an M4-or-newer Mac and are comfortable with Jupyter, you can run locally. Other capable machines work too.\n", "\n", "The optional bonus benchmark must run in the specified Colab runtime.\n", "\n", "[Project instructions and grading](https://cs145.web.app/projects/project2-memory.html)\n", "\n", "**Provided:** dataset loaders, embedding and storage helpers, SQLite baselines, evaluation helpers, and tests.\n", "\n", "**You implement:** BM25, LSH, fusion, routing, and benchmark drivers. For Tiers 93–100, also build and query within the memory budget. Incremental loading and disk partitioning are suggested approaches, not a required streaming framework.\n", "\n", "**Using the project page and notebook:** We’ve aimed to keep the project page and notebook consistent. The project page gives the full requirements and grading; the notebook provides compact guidance, function contracts, and starter code.\n", "\n", "The supplied helpers support the standard datasets through T87; you still implement the assigned components. For T93–T100, expect to adapt or replace the loading and indexing helpers to meet the larger-data and memory requirements. Test changes on a small subset before scaling up.\n", "\n", "You’re welcome to adapt the starter code within the project requirements. If something seems inconsistent or buggy, check the project page and ask privately on Ed; include the error and a small reproducible example for bugs. Note important assumptions in your notebook so we can take them into account when reviewing your work.\n", "\n", "## How to use this notebook\n", "\n", "Run Setup once per runtime; expand supplied code only when needed. Set your dataset and device in Configuration, then complete the phases below. Avoid **Run all** until your implementations are complete.\n", "\n", "- [Setup: run the supplied utilities](#scrollTo=477aa235)\n", "- [Phase 0: run the baseline](#scrollTo=d209b3a8) → [fill in the proposal](#scrollTo=proposal-summary)\n", "- Build your engine: [Phase 1: BM25](#scrollTo=20009eee) · [Phase 2: LSH](#scrollTo=425f0f26) · [Phase 3: fusion and planner](#scrollTo=24c77a3d)\n", "- Measure and explain: [Phase 4: benchmarks](#scrollTo=758e69cd) · [Phase 5: SQLite comparison](#scrollTo=a3086d4f) · [Phase 6: writeup](#scrollTo=e177ea65)\n", "- Higher tiers only: [Tier 93: fixed data and scaling](#scrollTo=4566c012) · [Tier 100: full corpus](#scrollTo=c4c85686)\n", "- [Optional class bonus competition](#scrollTo=bonus-v1-0)\n", "\n", "**Important: Optional checkpoints (Phase 1 onward).** Uncomment and run the [Drive checkpoint cell](#scrollTo=344b3faa) after Setup to save files written to `CACHE_DIR` and `INDEX_DIR` in your Google Drive. Recommended for longer runs; not required for Phase 0. This does not automatically save objects in memory.\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "477aa235", "metadata": { "id": "477aa235", "heading_collapsed": true }, "source": [ "## Setup\n", "\n", "**Supplied utilities — run these cells, but you do not need to read or rewrite them.** They define loaders, embeddings, SQLite references, and evaluation helpers. The large-data loader does not download anything until you call it in Tier 93.\n", "\n", "After running all cells in this Setup section, [jump to Phase 0](#scrollTo=d209b3a8). Run its baseline cells and fill in the proposal summary. Your retrieval implementation begins in Phase 1; look for `YOU IMPLEMENT` and `NotImplementedError`.\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "30186762", "metadata": { "id": "30186762", "cellView": "form" }, "outputs": [], "source": [ "# @title Install dependencies\n", "!pip install -q pyarrow numpy sentence-transformers datasets sqlite-vec tqdm psutil" ] }, { "cell_type": "code", "execution_count": null, "id": "8be3f5c1", "metadata": { "id": "8be3f5c1", "cellView": "form" }, "outputs": [], "source": [ "# @title Configuration and core utilities\n", "import gc\n", "import hashlib\n", "import json\n", "import os\n", "import pickle\n", "import random\n", "import re\n", "import sqlite3\n", "import sys\n", "import time\n", "from collections import Counter, defaultdict\n", "from pathlib import Path\n", "from typing import Optional\n", "\n", "import numpy as np\n", "import pyarrow as pa\n", "import pyarrow.parquet as pq\n", "import psutil\n", "import sqlite_vec\n", "from sentence_transformers import SentenceTransformer\n", "from tqdm import tqdm\n", "\n", "random.seed(42)\n", "np.random.seed(42)\n", "\n", "EMBEDDING_MODEL = \"all-MiniLM-L6-v2\" # 384-dim, fast, good enough\n", "EMBEDDING_DIM = 384\n", "INDEX_DIR = Path(\"./nanomem_index\")\n", "INDEX_DIR.mkdir(exist_ok=True)\n", "CACHE_DIR = Path(\"./bench_cache\")\n", "CACHE_DIR.mkdir(exist_ok=True)\n", "\n", "PROC = psutil.Process(os.getpid())\n", "\n", "\n", "def mem_mb() -> float:\n", " \"\"\"Current RSS in MB. Useful for the writeup section.\"\"\"\n", " return PROC.memory_info().rss / 1_000_000\n", "\n", "\n", "class MemoryBudget:\n", " \"\"\"Track peak memory (RSS) across a block of code plus across the whole run.\n", "\n", " Use as a context manager around any code where memory matters.\n", " Required at Tier 93 (target 2000) and Tier 100 (target 1000).\n", "\n", " with MemoryBudget(\"Phase 1 build\", limit_mb=2000) as mb:\n", " inv_idx.build(docs)\n", " mb.update() # sample again after allocation-heavy steps\n", " # prints: [mem] Phase 1 build: peak 1340 MB / budget 2000 MB\n", "\n", " At the end of the run:\n", " print(f\"Run peak: {MemoryBudget.run_peak():.0f} MB\")\n", "\n", " The eval harness uses this automatically and reports peak_mb in\n", " every result dict, so your scoreboard table can quote it.\n", "\n", " Note: this reports `psutil` peak RSS, not the conceptual working set.\n", " At Tier 93/100, peak RSS will often read higher than 2 GB even when\n", " you only have one shard live at a time, because Python doesn't return\n", " freed pages to the OS. We assess live memory use by reading\n", " your notebook; this number is a soft signal. **Keep the `[mem]`\n", " prints visible in the submitted notebook.**\n", " \"\"\"\n", "\n", " _run_peak: float = 0.0\n", "\n", " def __init__(self, label: str, limit_mb: float | None = None):\n", " self.label = label\n", " self.limit_mb = limit_mb\n", " self.start: float = 0.0\n", " self.peak: float = 0.0\n", "\n", " def __enter__(self):\n", " self.start = mem_mb()\n", " self.peak = self.start\n", " return self\n", "\n", " def update(self) -> float:\n", " cur = mem_mb()\n", " self.peak = max(self.peak, cur)\n", " MemoryBudget._run_peak = max(MemoryBudget._run_peak, self.peak)\n", " return cur\n", "\n", " def __exit__(self, *args):\n", " self.update()\n", " msg = f\"[mem] {self.label}: peak {self.peak:.0f} MB\"\n", " if self.limit_mb is not None:\n", " msg += f\" / budget {self.limit_mb:.0f} MB\"\n", " if self.peak > self.limit_mb:\n", " msg += \" (RSS above target; inspect live working set)\"\n", " print(msg)\n", " return False # do not suppress exceptions\n", "\n", " @classmethod\n", " def run_peak(cls) -> float:\n", " \"\"\"Highest RSS seen inside any MemoryBudget block so far.\"\"\"\n", " return cls._run_peak\n", "\n", " @classmethod\n", " def reset(cls):\n", " cls._run_peak = 0.0\n", "\n", "\n", "def memory_breakdown(components: dict, label: str = \"memory profile\") -> dict:\n", " \"\"\"Print a per-component memory breakdown table.\n", "\n", " Use this to figure out *where* your memory is going before you decide\n", " what to optimize at Tier 93 or Tier 100. Pass a dict of\n", " {label: object}; values can be:\n", "\n", " - numpy.ndarray: uses .nbytes (exact)\n", " - InvertedIndex: walks postings dict + doc_lengths\n", " - LSHIndex: walks hyperplane arrays + bucket dicts\n", " - int / float: treated as MB directly (use for measurements\n", " you took elsewhere, e.g., measured native allocations)\n", " - anything else: sys.getsizeof(obj) / 1_000_000 (approximate)\n", "\n", " Returns the {label: mb} dict so you can also persist it.\n", "\n", " **Keep the printed table visible in your notebook for Tier 93/100.**\n", " These are approximate component sizes, not proof of a peak-memory bound.\n", " Include embeddings and document text separately from their indexes.\n", " Account for shared state, native allocations, and temporary buffers.\n", "\n", " Measure components on your own data. Choose partition sizes from those\n", " measurements, leaving room for shared state and temporary buffers.\n", " \"\"\"\n", " import sys\n", " sizes = {}\n", " for k, v in components.items():\n", " if isinstance(v, np.ndarray):\n", " sizes[k] = v.nbytes / 1_000_000\n", " elif isinstance(v, (int, float)):\n", " sizes[k] = float(v)\n", " elif hasattr(v, 'postings') and hasattr(v, 'doc_lengths'):\n", " n = sys.getsizeof(v.postings)\n", " for term, plist in v.postings.items():\n", " n += sys.getsizeof(term) + sys.getsizeof(plist) + len(plist) * 56\n", " n += sys.getsizeof(v.doc_lengths)\n", " sizes[k] = n / 1_000_000\n", " elif hasattr(v, 'bands'):\n", " n = 0\n", " for hp, buckets in v.bands:\n", " n += hp.nbytes if isinstance(hp, np.ndarray) else 0\n", " n += sys.getsizeof(buckets)\n", " for _bid, doc_ids in buckets.items():\n", " n += sys.getsizeof(doc_ids) + len(doc_ids) * 28\n", " sizes[k] = n / 1_000_000\n", " else:\n", " sizes[k] = sys.getsizeof(v) / 1_000_000\n", "\n", " total = sum(sizes.values())\n", " print(f\"\\n=== {label} ===\")\n", " print(f\"{'COMPONENT':<28} {'MB':>10} {'%':>6}\")\n", " print(\"-\" * 46)\n", " for k in sorted(sizes, key=lambda x: -sizes[x]):\n", " mb = sizes[k]\n", " pct = (100 * mb / total) if total else 0\n", " print(f\"{k:<28} {mb:>10.1f} {pct:>5.1f}%\")\n", " print(\"-\" * 46)\n", " print(f\"{'TOTAL':<28} {total:>10.1f} 100.0%\")\n", " print(f\"process peak so far: {MemoryBudget.run_peak():>7.1f} MB\")\n", " return sizes" ] }, { "cell_type": "markdown", "id": "7babd163", "metadata": { "id": "7babd163" }, "source": [ "### Embedding model (supplied helpers, frozen)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1f02bd8b", "metadata": { "id": "1f02bd8b", "cellView": "form" }, "outputs": [], "source": [ "# @title Embedding helpers\n", "# CPU works locally and in free Colab; change to \"cuda\" if your Colab has a GPU.\n", "EMBED_DEVICE = \"cpu\"\n", "import torch\n", "torch.set_num_threads(min(4, os.cpu_count() or 1))\n", "_embed_model = None\n", "\n", "\n", "def get_embedder() -> SentenceTransformer:\n", " global _embed_model\n", " if _embed_model is None:\n", " print(\"Loading embedding model (one-time)...\")\n", " _embed_model = SentenceTransformer(EMBEDDING_MODEL, device=EMBED_DEVICE)\n", " return _embed_model\n", "\n", "\n", "def embed_text(text: str) -> np.ndarray:\n", " \"\"\"Return a (EMBEDDING_DIM,) float32 unit vector.\"\"\"\n", " v = get_embedder().encode([text], convert_to_numpy=True)[0].astype(np.float32)\n", " return v / (np.linalg.norm(v) + 1e-12)\n", "\n", "\n", "def embed_batch(texts: list[str], batch_size: int = 32) -> np.ndarray:\n", " \"\"\"Return an (N, EMBEDDING_DIM) float32 matrix of unit vectors.\"\"\"\n", " M = get_embedder().encode(\n", " texts, batch_size=batch_size, show_progress_bar=True, convert_to_numpy=True\n", " ).astype(np.float32)\n", " norms = np.linalg.norm(M, axis=1, keepdims=True) + 1e-12\n", " return M / norms\n", "\n", "\n", "def cached_embed_batch(doc_ids: list[str], texts: list[str], cache_path: Path) -> np.ndarray:\n", " \"\"\"Compute embeddings once, save to Parquet, reload on subsequent calls.\n", "\n", " Use this to reuse embeddings across runs. Time a small batch on your\n", " machine before estimating the full run; document length affects runtime.\n", " The cache key is `cache_path`; pick a name that includes the dataset and\n", " tier (e.g., `bench_cache/codesearchnet_250k_embs.parquet`).\n", "\n", " On Colab specifically: Colab's local disk does NOT persist across session\n", " restarts. Mount Google Drive and point `cache_path` there:\n", "\n", " from google.colab import drive\n", " drive.mount('/content/drive')\n", " CACHE_DIR = Path(\"/content/drive/MyDrive/nanomem_cache\")\n", " CACHE_DIR.mkdir(exist_ok=True)\n", " embs = cached_embed_batch(doc_ids, texts, CACHE_DIR / \"codesearchnet_250k.parquet\")\n", "\n", " For Tier 100, cache bounded batches separately so you can resume without\n", " holding the full embedding matrix in memory.\n", "\n", " The sidecar fingerprint includes document IDs, text, model configuration,\n", " and normalization. Changing any of these invalidates the cached embeddings.\n", " \"\"\"\n", " if len(doc_ids) != len(texts):\n", " raise ValueError(\"Provide one text per document ID.\")\n", " model = get_embedder()\n", " settings = {\"model\": EMBEDDING_MODEL, \"dimension\": EMBEDDING_DIM,\n", " \"max_seq_length\": model.max_seq_length,\n", " \"prompts\": getattr(model, \"prompts\", {}),\n", " \"default_prompt\": getattr(model, \"default_prompt_name\", None),\n", " \"normalization\": \"l2-float32-v1\"}\n", " key = hashlib.sha256()\n", " key.update(json.dumps(settings, sort_keys=True, default=str).encode())\n", " for doc_id, text in zip(doc_ids, texts):\n", " key.update(json.dumps([doc_id, text], ensure_ascii=False).encode(\"utf-8\"))\n", " key.update(b\"\\n\")\n", " key = key.hexdigest()\n", " cache_path.parent.mkdir(parents=True, exist_ok=True)\n", " key_path = cache_path.with_suffix(cache_path.suffix + \".key\")\n", " if cache_path.exists() and key_path.exists() and key_path.read_text().strip() == key:\n", " _, cached_embs = load_embeddings_parquet(cache_path)\n", " print(f\" [embed cache HIT] {cache_path.name} ({cached_embs.shape[0]} docs)\")\n", " return cached_embs\n", " reason = \"no cache\" if not cache_path.exists() else \"input or model changed\"\n", " print(f\" [embed cache MISS: {reason}] computing embeddings for {len(texts)} docs...\")\n", " embs = embed_batch(texts)\n", " save_embeddings_parquet(doc_ids, embs, cache_path)\n", " key_path.write_text(key)\n", " print(f\" [embed cache WROTE] {cache_path.name}\")\n", " return embs" ] }, { "cell_type": "markdown", "id": "552ff218", "metadata": { "id": "552ff218" }, "source": [ "### Tokenizer (supplied helpers, modifiable)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d6b07524", "metadata": { "id": "d6b07524", "cellView": "form" }, "outputs": [], "source": [ "# @title Tokenizer\n", "_TOKEN_RE = re.compile(r\"[A-Za-z_][A-Za-z0-9_]*\")\n", "\n", "\n", "def tokenize(text: str) -> list[str]:\n", " \"\"\"Lowercase identifier-style tokenizer.\n", "\n", " Splits on non-identifier characters; keeps `verify_jwt` as one token.\n", " Modify if you want, but you do not have to.\n", " \"\"\"\n", " return [t.lower() for t in _TOKEN_RE.findall(text)]" ] }, { "cell_type": "markdown", "id": "6538f4ea", "metadata": { "id": "6538f4ea" }, "source": [ "### Storage helpers (supplied helpers)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "830ef15e", "metadata": { "id": "830ef15e", "cellView": "form" }, "outputs": [], "source": [ "# @title Storage helpers\n", "def save_documents_json(docs: list[dict], path: Path) -> None:\n", " with open(path, \"w\") as f:\n", " for d in docs:\n", " f.write(json.dumps(d) + \"\\n\")\n", "\n", "\n", "def load_documents_json(path: Path) -> list[dict]:\n", " docs = []\n", " with open(path) as f:\n", " for line in f:\n", " docs.append(json.loads(line))\n", " return docs\n", "\n", "\n", "def save_embeddings_parquet(doc_ids: list[str], embeddings: np.ndarray, path: Path) -> None:\n", " \"\"\"One Parquet file. doc_id (string) + embedding (fixed-size list of float32).\"\"\"\n", " embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)\n", " if embeddings.shape != (len(doc_ids), EMBEDDING_DIM):\n", " raise ValueError(\"Provide one embedding of the expected dimension per document.\")\n", " table = pa.table({\n", " \"doc_id\": pa.array(doc_ids, type=pa.string()),\n", " \"embedding\": pa.FixedSizeListArray.from_arrays(\n", " pa.array(embeddings.reshape(-1), type=pa.float32()),\n", " EMBEDDING_DIM,\n", " ),\n", " })\n", " pq.write_table(table, path, compression=\"snappy\")\n", "\n", "\n", "def load_embeddings_parquet(path: Path) -> tuple[list[str], np.ndarray]:\n", " table = pq.read_table(path)\n", " doc_ids = table.column(\"doc_id\").to_pylist()\n", " column = table.column(\"embedding\").combine_chunks()\n", " embs = column.flatten().to_numpy(zero_copy_only=False).reshape(-1, EMBEDDING_DIM)\n", " return doc_ids, embs\n", "\n", "\n", "def brute_force_cosine(query_vec: np.ndarray, embeddings: np.ndarray, top_k: int = 10) -> list[int]:\n", " \"\"\"Return the top_k indices into `embeddings` by cosine similarity.\"\"\"\n", " scores = embeddings @ query_vec\n", " return list(np.argsort(-scores)[:top_k])" ] }, { "cell_type": "markdown", "id": "9e4c3a86", "metadata": { "id": "9e4c3a86" }, "source": [ "### LSH math helpers (supplied helpers)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "64257992", "metadata": { "id": "64257992", "cellView": "form" }, "outputs": [], "source": [ "# @title LSH math helpers\n", "def random_projection_matrix(n_hyperplanes: int, dim: int = EMBEDDING_DIM) -> np.ndarray:\n", " \"\"\"Returns (n_hyperplanes, dim) standard-normal hyperplanes.\"\"\"\n", " return np.random.randn(n_hyperplanes, dim).astype(np.float32)\n", "\n", "\n", "def hash_vector(vec: np.ndarray, hyperplanes: np.ndarray) -> int:\n", " \"\"\"Project vec onto each hyperplane; return bit signature as int.\"\"\"\n", " bits = (hyperplanes @ vec > 0).astype(np.uint8)\n", " sig = 0\n", " for b in bits:\n", " sig = (sig << 1) | int(b)\n", " return sig\n", "\n", "\n", "def hash_batch(M: np.ndarray, hyperplanes: np.ndarray) -> np.ndarray:\n", " \"\"\"Vectorized hash of an (N, dim) matrix. Returns (N,) int64 sigs.\"\"\"\n", " bits = (M @ hyperplanes.T > 0).astype(np.uint8)\n", " n_hp = bits.shape[1]\n", " weights = (1 << np.arange(n_hp - 1, -1, -1)).astype(np.int64)\n", " return (bits.astype(np.int64) * weights).sum(axis=1)" ] }, { "cell_type": "markdown", "id": "c8b42c00", "metadata": { "id": "c8b42c00" }, "source": [ "### SQLite reference helpers (used in Phase 0 and Phase 5)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6c979a7f", "metadata": { "id": "6c979a7f", "cellView": "form" }, "outputs": [], "source": [ "# @title SQLite reference helpers\n", "def open_sqlite_db(path: Path) -> sqlite3.Connection:\n", " \"\"\"Open a SQLite connection with the sqlite-vec extension loaded.\"\"\"\n", " path.parent.mkdir(parents=True, exist_ok=True)\n", " conn = sqlite3.connect(str(path))\n", " try:\n", " conn.execute(\"CREATE VIRTUAL TABLE temp.fts_check USING fts5(text)\")\n", " conn.execute(\"DROP TABLE temp.fts_check\")\n", " conn.enable_load_extension(True)\n", " sqlite_vec.load(conn)\n", " conn.enable_load_extension(False)\n", " except (AttributeError, sqlite3.Error) as exc:\n", " conn.close()\n", " raise RuntimeError(\"SQLite needs FTS5 and loadable extensions. Use the supplied Colab setup or a Python build with both enabled.\") from exc\n", " return conn\n", "\n", "\n", "def initialize_sqlite_corpus(conn: sqlite3.Connection) -> None:\n", " \"\"\"Create empty reference tables. Call once before inserting batches.\n", "\n", " This resets the reference corpus; do not call between batches.\n", " \"\"\"\n", " with conn:\n", " conn.execute(\"DROP TABLE IF EXISTS docs_fts\")\n", " conn.execute(\"DROP TABLE IF EXISTS docs_vec\")\n", " conn.execute(\"DROP TABLE IF EXISTS doc_id_map\")\n", " conn.execute(\"CREATE VIRTUAL TABLE docs_fts USING fts5(doc_id UNINDEXED, content)\")\n", " conn.execute(f\"CREATE VIRTUAL TABLE docs_vec USING vec0(doc_idx INTEGER PRIMARY KEY, embedding float[{EMBEDDING_DIM}])\")\n", " conn.execute(\"CREATE TABLE doc_id_map(doc_idx INTEGER PRIMARY KEY, doc_id TEXT UNIQUE)\")\n", "\n", "\n", "def append_sqlite_batch(conn: sqlite3.Connection, docs: list[dict], embeddings: np.ndarray) -> None:\n", " \"\"\"Append one bounded batch; retain earlier batches. IDs must be unique.\n", "\n", " Initialize once, then embed and append each batch. A failed batch rolls back.\n", " \"\"\"\n", " embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)\n", " if embeddings.shape != (len(docs), EMBEDDING_DIM) or not np.isfinite(embeddings).all():\n", " raise ValueError(\"Provide one finite embedding of the expected dimension per document.\")\n", " if len({d[\"doc_id\"] for d in docs}) != len(docs):\n", " raise ValueError(\"Document IDs must be unique.\")\n", " start = conn.execute(\"SELECT COALESCE(MAX(doc_idx), -1) + 1 FROM doc_id_map\").fetchone()[0]\n", " with conn:\n", " conn.executemany(\"INSERT INTO doc_id_map(doc_idx, doc_id) VALUES (?, ?)\",\n", " ((start+i, d[\"doc_id\"]) for i,d in enumerate(docs)))\n", " conn.executemany(\"INSERT INTO docs_fts(doc_id, content) VALUES (?, ?)\",\n", " ((d[\"doc_id\"], d[\"text\"]) for d in docs))\n", " conn.executemany(\"INSERT INTO docs_vec(doc_idx, embedding) VALUES (?, ?)\",\n", " ((start+i, embeddings[i].tobytes()) for i in range(len(docs))))\n", "\n", "\n", "def build_sqlite_corpus(conn: sqlite3.Connection, docs: list[dict], embeddings: np.ndarray) -> None:\n", " \"\"\"Rebuild a small reference corpus. Larger runs use initialize + append.\"\"\"\n", " initialize_sqlite_corpus(conn)\n", " append_sqlite_batch(conn, docs, embeddings)\n", "\n", "\n", "def fts5_query(conn: sqlite3.Connection, query: str, top_k: int = 10) -> list[tuple[str, float]]:\n", " \"\"\"BM25 query via FTS5. Returns [(doc_id, score), ...] sorted by relevance.\"\"\"\n", " tokens = re.findall(r\"[A-Za-z0-9_]+\", query)\n", " if not tokens:\n", " return []\n", " fts_query = \" OR \".join(f'\"{t}\"' for t in tokens)\n", " try:\n", " rows = conn.execute(\n", " \"SELECT doc_id, -bm25(docs_fts) FROM docs_fts WHERE docs_fts MATCH ? ORDER BY rank LIMIT ?\",\n", " (fts_query, top_k)).fetchall()\n", " return [(r[0], r[1]) for r in rows]\n", " except sqlite3.OperationalError as exc:\n", " raise RuntimeError(\"FTS5 query failed; inspect the index and query instead of treating this as zero matches.\") from exc\n", "\n", "\n", "def vec_query(conn: sqlite3.Connection, query_vec: np.ndarray, top_k: int = 10) -> list[tuple[str, float]]:\n", " \"\"\"Exact vector KNN via sqlite-vec. Returns [(doc_id, score), ...].\n", "\n", " Score is negative L2 distance: higher is better, but it is not cosine similarity.\n", " \"\"\"\n", " rows = conn.execute(\n", " \"\"\"SELECT m.doc_id, -v.distance FROM docs_vec v\n", " JOIN doc_id_map m ON v.doc_idx = m.doc_idx\n", " WHERE v.embedding MATCH ? AND k = ? ORDER BY distance LIMIT ?\"\"\",\n", " (query_vec.tobytes(), top_k, top_k)).fetchall()\n", " return [(r[0], r[1]) for r in rows]" ] }, { "cell_type": "markdown", "id": "e995d02f", "metadata": { "id": "e995d02f" }, "source": [ "### Corpus loaders\n", "\n", "**Benchmark note:** CodeSearchNet queries come from docstrings included in the indexed functions, so keyword matching has an advantage. Compare your implementation against its corresponding SQLite baseline; vector search need not outperform keyword search.\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c0c0243b", "metadata": { "id": "c0c0243b", "cellView": "form" }, "outputs": [], "source": [ "# @title Corpus loaders\n", "# These helpers load small corpora into memory. For Tiers 93/100, adapt loading to stay within the memory budget; fixed counts do not establish byte size.\n", "TIER_SIZES = {\"tiny\": 1_000, \"25k\": 25_000, \"250k\": 250_000}\n", "\n", "\n", "def load_codesearchnet(tier: str = \"tiny\") -> tuple[list[dict], list[dict]]:\n", " \"\"\"Load CodeSearchNet at the chosen tier. Returns (docs, eval_set).\n", "\n", " eval_set is built from docstrings: query is the first nonempty line of the\n", " docstring; relevant_doc_ids is [the function whose docstring it was].\n", " \"\"\"\n", " from datasets import load_dataset\n", "\n", " if tier not in TIER_SIZES:\n", " raise ValueError(f\"unknown tier {tier!r}; choose from {list(TIER_SIZES)}\")\n", " n = TIER_SIZES[tier]\n", " use_streaming = True\n", "\n", " ds = load_dataset(\n", " \"code-search-net/code_search_net\", \"python\", split=\"train\",\n", " streaming=use_streaming,\n", " )\n", " docs, eval_pool = [], []\n", " for row in ds:\n", " if len(docs) >= n:\n", " break\n", " code = row[\"func_code_string\"]\n", " if not code:\n", " continue\n", " doc_id = f\"python/{row['repository_name']}/{row['func_name']}::{len(docs)}\"\n", " docs.append({\"doc_id\": doc_id, \"text\": code})\n", " docstring = (row.get(\"func_documentation_string\") or \"\").strip()\n", " if docstring:\n", " first = next((line.strip() for line in docstring.split(\"\\n\") if line.strip()), \"\")\n", " if first and 20 <= len(first) <= 200:\n", " eval_pool.append({\"query\": first, \"relevant_doc_ids\": [doc_id]})\n", " if len(docs) < n:\n", " raise ValueError(f\"Requested {n:,} documents but source supplied {len(docs):,}. Use a smaller tier or extend the loader with documented sources.\")\n", " rng = random.Random(42)\n", " eval_n = min(500, len(eval_pool))\n", " eval_set = rng.sample(eval_pool, eval_n)\n", " return docs, eval_set\n", "\n", "\n", "def load_msmarco(tier: str = \"tiny\") -> tuple[list[dict], list[dict]]:\n", " \"\"\"Load MS MARCO at the chosen tier. Returns (docs, eval_set).\n", "\n", " eval_set is real Bing/SQ queries from the BeIR/msmarco-qrels training split,\n", " filtered to those whose relevant docs are in the loaded subset.\n", " \"\"\"\n", " from datasets import load_dataset\n", "\n", " if tier not in TIER_SIZES:\n", " raise ValueError(f\"unknown tier {tier!r}; choose from {list(TIER_SIZES)}\")\n", " n = TIER_SIZES[tier]\n", "\n", " corpus_ds = load_dataset(\"BeIR/msmarco\", \"corpus\", split=\"corpus\", streaming=True)\n", " queries_ds = load_dataset(\"BeIR/msmarco\", \"queries\", split=\"queries\", streaming=True)\n", " qrels_ds = load_dataset(\"BeIR/msmarco-qrels\", split=\"train\", streaming=True)\n", "\n", " docs = []\n", " docs_in_subset = set()\n", " for i, row in enumerate(corpus_ds):\n", " if i >= n:\n", " break\n", " doc_id = str(row[\"_id\"])\n", " text = ((row.get(\"title\") or \"\") + \" \" + (row.get(\"text\") or \"\")).strip()\n", " docs.append({\"doc_id\": doc_id, \"text\": text})\n", " docs_in_subset.add(doc_id)\n", "\n", " qid_to_text = {str(row[\"_id\"]): row[\"text\"] for row in queries_ds}\n", " qid_to_relevant = {}\n", " for row in qrels_ds:\n", " qid = str(row[\"query-id\"])\n", " did = str(row[\"corpus-id\"])\n", " if did in docs_in_subset and int(row[\"score\"]) > 0:\n", " qid_to_relevant.setdefault(qid, []).append(did)\n", " eligible = [q for q, r in qid_to_relevant.items() if r and q in qid_to_text]\n", " if len(docs) < n:\n", " raise ValueError(f\"Requested {n:,} documents but source supplied {len(docs):,}. Use a smaller tier or extend the loader with documented sources.\")\n", " rng = random.Random(42)\n", " eval_n = min(500, len(eligible))\n", " eval_set = [{\"query\": qid_to_text[q], \"relevant_doc_ids\": qid_to_relevant[q]}\n", " for q in rng.sample(eligible, eval_n)]\n", " return docs, eval_set" ] }, { "cell_type": "markdown", "id": "f805a071", "metadata": { "id": "f805a071" }, "source": [ "### Eval harness (supplied helpers)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5e294d79", "metadata": { "id": "5e294d79", "cellView": "form" }, "outputs": [], "source": [ "# @title Evaluation helpers\n", "def recall_at_k(retrieved: list[str], relevant: list[str], k: int = 5) -> float:\n", " \"\"\"Fraction of distinct relevant documents found in the first k results.\"\"\"\n", " relevant_ids = set(relevant)\n", " return len(set(retrieved[:k]) & relevant_ids) / len(relevant_ids) if relevant_ids else 0.0\n", "\n", "\n", "def mean_reciprocal_rank(retrieved: list[str], relevant: list[str]) -> float:\n", " rel_set = set(relevant)\n", " for i, doc in enumerate(retrieved, start=1):\n", " if doc in rel_set:\n", " return 1.0 / i\n", " return 0.0\n", "\n", "\n", "def run_eval(query_fn, eval_set: list[dict], top_k: int = 10,\n", " memory_limit_mb: float | None = None) -> dict:\n", " \"\"\"Run query_fn against an eval set. Reports R@1, R@5, R@10, MRR@10,\n", " p50/p95 latency, AND peak RSS during the run.\n", "\n", " Pass `memory_limit_mb=2000` for Tier 93 or `memory_limit_mb=1000`\n", " for Tier 100 to label the RSS report with the tier target. This\n", " does not enforce or measure the live working-set cap. Report live\n", " components separately with memory_breakdown and explain how your implementation stays within the memory budget.\n", " \"\"\"\n", " r1, r5, r10, mrrs, lat = [], [], [], [], []\n", " with MemoryBudget(\"eval run\", limit_mb=memory_limit_mb) as mb:\n", " for ex in eval_set:\n", " t0 = time.perf_counter()\n", " results = query_fn(ex[\"query\"])\n", " lat.append((time.perf_counter() - t0) * 1000)\n", " mb.update()\n", " retrieved = [d for d, _ in results] if (results and isinstance(results[0], tuple)) else (results or [])\n", " r1.append(recall_at_k(retrieved, ex[\"relevant_doc_ids\"], 1))\n", " r5.append(recall_at_k(retrieved, ex[\"relevant_doc_ids\"], 5))\n", " r10.append(recall_at_k(retrieved, ex[\"relevant_doc_ids\"], 10))\n", " mrrs.append(mean_reciprocal_rank(retrieved[:10], ex[\"relevant_doc_ids\"]))\n", " return {\n", " \"n\": len(eval_set),\n", " \"recall@1\": float(np.mean(r1)),\n", " \"recall@5\": float(np.mean(r5)),\n", " \"recall@10\": float(np.mean(r10)),\n", " \"mrr@10\": float(np.mean(mrrs)),\n", " \"p50_ms\": float(np.percentile(lat, 50)),\n", " \"p95_ms\": float(np.percentile(lat, 95)),\n", " \"peak_mb\": mb.peak,\n", " }" ] }, { "cell_type": "markdown", "id": "scale-utility-heading", "metadata": { "id": "scale-utility-heading" }, "source": [ "### Fixed Tier 93 data (supplied utility)\n", "\n", "Run this definition with Setup. Call it only when you reach Tier 93.\n", "\n", "\n" ] }, { "cell_type": "code", "id": "scale-utility", "metadata": { "id": "scale-utility", "cellView": "form" }, "source": [ "# @title Tier 93 data helper\n", "# Supplied utility: defines a loader; downloads nothing until you call it.\n", "def prepare_tier93_data(data_dir=\"nanomem_scale\"):\n", " import hashlib, importlib.util, urllib.request\n", " from pathlib import Path\n", " root = Path(data_dir)\n", " root.mkdir(parents=True, exist_ok=True)\n", " helper_path = root / \"nanomem_benchmark_v1.py\"\n", " payload = urllib.request.urlopen(\n", " \"https://cs145.web.app/projects/nanomem/bonus/benchmark_v1.py\"\n", " ).read()\n", " if hashlib.sha256(payload).hexdigest() != \"c3acbaf9bfb77dd4bb471df1b67b8a3d371df09b4e9afc9ba4af96c9baaed1c1\":\n", " raise ValueError(\"Helper version mismatch\")\n", " helper_path.write_bytes(payload)\n", " spec = importlib.util.spec_from_file_location(\"nanomem_benchmark\", helper_path)\n", " utility = importlib.util.module_from_spec(spec)\n", " spec.loader.exec_module(utility)\n", " data = utility.load_bonus_benchmark(root, config_sha256=\"75744dcb478a27f4ef75972361bcf015d626e88633cda85110da122d2ce2456f\")\n", " utility.prepare_documents(data)\n", " return utility, data\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "840817cd", "metadata": { "id": "840817cd" }, "source": [ "### Tiny dataset for unit tests\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d6ffe31a", "metadata": { "id": "d6ffe31a", "cellView": "form" }, "outputs": [], "source": [ "# @title Tiny test data\n", "TINY_CORPUS = [\n", " {\"doc_id\": \"auth/jwt.py:1\", \"text\": \"def verify_jwt(token, secret): return jwt.decode(token, secret, algorithms=['HS256'])\"},\n", " {\"doc_id\": \"auth/jwt.py:2\", \"text\": \"def issue_jwt(user_id, secret, ttl=3600): return jwt.encode({'sub': user_id, 'exp': time.time() + ttl}, secret)\"},\n", " {\"doc_id\": \"auth/scopes.py:1\", \"text\": \"def verify_jwt_with_scope(token, secret, required_scope): payload = verify_jwt(token, secret); assert required_scope in payload['scopes']\"},\n", " {\"doc_id\": \"config/parser.py:1\", \"text\": \"def parse_config(path): with open(path) as f: return json.load(f)\"},\n", " {\"doc_id\": \"config/parser.py:2\", \"text\": \"def merge_configs(base, override): merged = dict(base); merged.update(override); return merged\"},\n", " {\"doc_id\": \"utils/load_json.py:1\", \"text\": \"def load_json(path): with open(path, 'r') as fh: return json.load(fh)\"},\n", " {\"doc_id\": \"utils/load_json.py:2\", \"text\": \"def save_json(obj, path): with open(path, 'w') as fh: json.dump(obj, fh, indent=2)\"},\n", " {\"doc_id\": \"db/connection.py:1\", \"text\": \"def connect_postgres(dsn): conn = psycopg2.connect(dsn); conn.autocommit = False; return conn\"},\n", " {\"doc_id\": \"db/connection.py:2\", \"text\": \"def close_connection(conn): conn.commit(); conn.close()\"},\n", " {\"doc_id\": \"db/transactions.py:1\", \"text\": \"def with_transaction(conn, fn): try: result = fn(conn); conn.commit(); return result; except: conn.rollback(); raise\"},\n", " {\"doc_id\": \"search/tokenize.py:1\", \"text\": \"def tokenize(text): return [w.lower() for w in re.split(r'\\\\W+', text) if w]\"},\n", " {\"doc_id\": \"search/bm25.py:1\", \"text\": \"def bm25_score(tf, df, n_docs, doc_len, avg_len, k1=1.2, b=0.75): idf = math.log((n_docs - df + 0.5) / (df + 0.5))\"},\n", " {\"doc_id\": \"cache/redis_client.py:1\", \"text\": \"def cache_get(key): return redis.get(key)\"},\n", " {\"doc_id\": \"cache/redis_client.py:2\", \"text\": \"def cache_set(key, value, ttl=300): redis.setex(key, ttl, value)\"},\n", " {\"doc_id\": \"models/user.py:1\", \"text\": \"class User: def __init__(self, id, email): self.id = id; self.email = email\"},\n", " {\"doc_id\": \"models/user.py:2\", \"text\": \"def get_user_by_email(session, email): return session.query(User).filter_by(email=email).first()\"},\n", " {\"doc_id\": \"api/routes.py:1\", \"text\": \"@app.route('/login', methods=['POST']) def login(): user = authenticate(request.json); return issue_jwt(user.id, SECRET)\"},\n", " {\"doc_id\": \"api/routes.py:2\", \"text\": \"@app.route('/logout', methods=['POST']) def logout(): revoke_token(request.headers['Authorization'])\"},\n", " {\"doc_id\": \"tests/test_auth.py:1\", \"text\": \"def test_jwt_roundtrip(): token = issue_jwt('user123', 'secret'); assert verify_jwt(token, 'secret')['sub'] == 'user123'\"},\n", " {\"doc_id\": \"tests/test_config.py:1\", \"text\": \"def test_parse_config(): cfg = parse_config('test.json'); assert cfg['db'] == 'postgres'\"},\n", "]" ] }, { "cell_type": "markdown", "id": "visible-configuration", "metadata": { "id": "visible-configuration" }, "source": [ "## Configuration\n", "\n", "Choose the dataset and embedding device here after running Setup.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "459b1bc5", "metadata": { "id": "459b1bc5" }, "outputs": [], "source": [ "# Small-data setting: tiny = 1,000; 25k = 25,000; 250k = 250,000 documents per corpus.\n", "# Change this, then rerun Phase 0 and the phase execution cells you need.\n", "DATASET = \"tiny\"\n", "assert DATASET in {\"tiny\", \"25k\", \"250k\"}, \"Use the bounded-memory scaling sections for Tiers 93/100.\"\n", "\n", "# Match your runtime: \"cuda\" for a Colab GPU; otherwise \"cpu\".\n", "EMBED_DEVICE = \"cpu\"\n" ] }, { "cell_type": "markdown", "id": "d209b3a8", "metadata": { "id": "d209b3a8" }, "source": [ "## Phase 0: SQLite reference setup\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#proposal)\n", "\n", "Run the cells below with `DATASET = \"tiny\"` first (1,000 documents per corpus). Then change `DATASET` in Configuration to `\"25k\"`, rerun that setting cell, and rerun Phase 0 for the proposal.\n", "\n", "**No NanoMem implementation is needed here.** Inspect the results, [fill in your proposal summary](#scrollTo=proposal-summary), and stop here for the proposal.\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "606d9d6c", "metadata": { "id": "606d9d6c", "cellView": "form" }, "outputs": [], "source": [ "# @title Phase 0 baseline helpers\n", "# Supplied setup: later phases contain the search algorithms you implement.\n", "def setup_phase0(tier=\"tiny\", queries=None):\n", " state = {}\n", " for corpus, loader in [(\"codesearchnet\", load_codesearchnet), (\"msmarco\", load_msmarco)]:\n", " docs, eval_set = loader(tier)\n", " embs = cached_embed_batch([d[\"doc_id\"] for d in docs], [d[\"text\"] for d in docs],\n", " CACHE_DIR / f\"{corpus}_{tier}.parquet\")\n", " conn = open_sqlite_db(INDEX_DIR / f\"{corpus}_phase0.db\")\n", " build_sqlite_corpus(conn, docs, embs)\n", " query = (queries or {}).get(corpus) or (eval_set[0][\"query\"] if eval_set else None)\n", " if not query:\n", " raise ValueError(f\"Choose a query for {corpus}.\")\n", " results = {\"fts5\": fts5_query(conn, query, top_k=3),\n", " \"vector\": vec_query(conn, embed_text(query), top_k=3)}\n", " stats = {\"documents\": len(docs), \"text_bytes\": sum(len(d[\"text\"].encode(\"utf-8\")) for d in docs),\n", " \"embedding_bytes\": embs.nbytes, \"dimensions\": embs.shape[1]}\n", " print(corpus, stats, \"query:\", query)\n", " text_by_id = {d[\"doc_id\"]: d[\"text\"] for d in docs}\n", " for method, hits in results.items():\n", " print(method)\n", " for doc_id, score in hits:\n", " print(doc_id, \"score:\", round(score, 4), \"text:\", text_by_id[doc_id][:300])\n", " state[corpus] = dict(conn=conn, docs=docs, embs=embs, eval_set=eval_set, query=query, results=results, stats=stats)\n", " return state\n", "\n", "\n", "def run_phase_0_tests(state):\n", " assert set(state) == {\"codesearchnet\", \"msmarco\"}\n", " for corpus, s in state.items():\n", " assert len(s[\"docs\"]) > 0\n", " assert s[\"embs\"].shape == (len(s[\"docs\"]), EMBEDDING_DIM)\n", " assert s[\"results\"][\"vector\"], f\"{corpus}: vector search returned no results\"\n", " # A query can legitimately have no lexical hits. Test indexing separately\n", " # using a token drawn from this corpus rather than the word 'function'.\n", " token = next(t for d in s[\"docs\"] for t in tokenize(d[\"text\"]) if t)\n", " assert fts5_query(s[\"conn\"], token, top_k=1), f\"{corpus}: lexical index check failed\"\n", " print(corpus, \"baseline checks passed\")\n", "\n", "\n", "def search_demo(corpus, query, top_k=3, method=\"both\"):\n", " \"\"\"Display ranked results from the existing Phase 0 indexes; no fusion.\"\"\"\n", " from html import escape\n", " from IPython.display import HTML, display\n", " if method not in (\"keyword\", \"vector\", \"both\"):\n", " raise ValueError('Choose method=\"keyword\", \"vector\", or \"both\".')\n", " if \"phase0\" not in globals() or corpus not in phase0:\n", " raise ValueError(\"Run the Phase 0 baseline first; choose codesearchnet or msmarco.\")\n", " if not query.strip():\n", " raise ValueError(\"Enter a query.\")\n", " if not isinstance(top_k, int) or top_k < 1:\n", " raise ValueError(\"top_k must be a positive integer.\")\n", " s = phase0[corpus]\n", " print(f\"Searching {s['stats']['documents']:,} documents in {corpus}: {query}\")\n", " results = {}\n", " if method in (\"keyword\", \"both\"):\n", " results[\"keyword\"] = fts5_query(s[\"conn\"], query, top_k)\n", " if method in (\"vector\", \"both\"):\n", " results[\"vector\"] = vec_query(s[\"conn\"], embed_text(query), top_k)\n", " labels = {\"keyword\": \"Keyword · BM25 score\", \"vector\": \"Vector · negative L2 distance\"}\n", " headers = \"Rank\" + \"\".join(f\"{labels[m]}\" for m in results)\n", " rows = []\n", " for i in range(max(1, max(len(hits) for hits in results.values()))):\n", " columns = [f\"{i + 1}\"]\n", " for hits in results.values():\n", " if i >= len(hits):\n", " columns.append(\"No match\")\n", " continue\n", " doc_id, score = hits[i]\n", " text = s[\"conn\"].execute(\"SELECT content FROM docs_fts WHERE doc_id = ?\", (doc_id,)).fetchone()[0]\n", " columns.append(f'{escape(str(doc_id))}
Score: {score:.4f}
{escape(text[:350])}')\n", " rows.append(\"\" + \"\".join(columns) + \"\")\n", " display(HTML('' + headers + '' + \"\".join(rows) + '
'))\n", " return results\n" ] }, { "cell_type": "code", "execution_count": null, "id": "34fc3443", "metadata": { "id": "34fc3443" }, "outputs": [], "source": [ "# Run the supplied baseline. Rerunning closes the previous database connections first.\n", "if \"phase0\" in globals():\n", " for previous in phase0.values():\n", " previous[\"conn\"].close()\n", " del phase0\n", "phase0 = setup_phase0(tier=DATASET)\n", "run_phase_0_tests(phase0)\n" ] }, { "cell_type": "markdown", "metadata": { "id": "phase0-try-query" }, "id": "phase0-try-query", "source": [ "### Try query variants\n", "\n", "Keyword search matches words; vector search looks for similar meaning. Both retrieve documents, not a generated answer. Compare ranks and snippets; scores use different scales.\n", "\n", "For each dataset, run the supplied query with `method = \"both\"`, then write and run two variants. Keep all three ranking tables visible in your submitted notebook; no extra paragraph is needed. You can also try `\"keyword\"` or `\"vector\"` separately. These cells reuse your indexes. RRF comes later.\n" ] }, { "cell_type": "code", "metadata": { "id": "phase0-code-query" }, "id": "phase0-code-query", "source": [ "# CodeSearchNet: find a Python function\n", "method = \"both\" # \"keyword\", \"vector\", or \"both\"\n", "queries = [\n", " \"How do I check a password against its stored hash?\", # Query 1\n", " \"\", # Query 2: write your own variant.\n", " \"\", # Query 3: try another wording.\n", "]\n", "for number, query in enumerate(queries, 1):\n", " if query.strip():\n", " print(f\"\\nQuery {number}\")\n", " _ = search_demo(\"codesearchnet\", query, method=method)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": { "id": "phase0-passage-query" }, "id": "phase0-passage-query", "source": [ "# MS MARCO: find a passage\n", "method = \"both\" # \"keyword\", \"vector\", or \"both\"\n", "queries = [\n", " \"How can I roast coffee beans at home?\", # Query 1\n", " \"\", # Query 2: write your own variant.\n", " \"\", # Query 3: try another wording.\n", "]\n", "for number, query in enumerate(queries, 1):\n", " if query.strip():\n", " print(f\"\\nQuery {number}\")\n", " _ = search_demo(\"msmarco\", query, method=method)\n" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": { "id": "proposal-summary" }, "source": [ "### Proposal: answer these questions\n", "\n", "Fill in the answer column. Keep the query ranking tables and measured sizes above.\n", "\n", "| Question | Your answer |\n", "|---|---|\n", "| Who is on your team, and how will you divide the work? | [Answer] |\n", "| Which tier are you targeting? | [Answer] |\n", "| Will you use local Jupyter, Colab, or both? | [Answer] |\n", "| What is your initial custom-fusion idea? | [Answer] |\n", "| For Tier 93/100, what data and memory-management plan will you use? | [Answer, or not applicable] |\n", "\n", "**Stop here for the proposal.** Phase 1 begins your own implementation.\n", "\n", "\n" ], "id": "proposal-summary" }, { "cell_type": "markdown", "id": "344b3faa", "metadata": { "id": "344b3faa" }, "source": [ "### Save Colab checkpoints to Drive\n", "\n", "For Colab, checkpoints are optional and recommended from Phase 1 onward; they are not required for Phase 0. Local Jupyter users can skip this cell. Uncomment and run the cell below after Setup, then approve access to Google Drive.\n", "\n", "- **Embeddings:** the supplied cache calls write to `CACHE_DIR`.\n", "- **Indexes and results:** save your index files under `INDEX_DIR` and benchmark outputs under `CHECKPOINT_DIR` using your implementation.\n", "- **Next session:** run Setup, run the checkpoint cell again, and reload your saved files. Re-running Setup resets the paths to local storage.\n", "- **Long runs:** save numbered embedding chunks and record the last completed document so you can resume. Time a small batch before starting.\n", "\n", "Check that files appear in `MyDrive/nanomem_checkpoints` and that Drive has enough space. Enabling checkpoints does not move existing local files or save objects in memory.\n", "\n", "Without Drive checkpoints, embeddings go to `bench_cache/` and indexes to `nanomem_index/`. Downloads use the Hugging Face cache. Colab’s local files can disappear when the runtime resets; saving the notebook does not save them.\n", "\n", "\n" ] }, { "cell_type": "code", "id": "nanomem-drive-checkpoints", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# Optional: uncomment this cell to keep checkpoints in Google Drive.\n", "# Recommended from Phase 1 onward; not required for Phase 0.\n", "# Run AFTER Setup and BEFORE building indexes or computing embeddings.\n", "# from pathlib import Path\n", "# from google.colab import drive\n", "# drive.mount(\"/content/drive\")\n", "# CHECKPOINT_DIR = Path(\"/content/drive/MyDrive/nanomem_checkpoints\")\n", "# CACHE_DIR = CHECKPOINT_DIR / \"embeddings\"\n", "# INDEX_DIR = CHECKPOINT_DIR / \"indexes\"\n", "# for folder in (CACHE_DIR, INDEX_DIR):\n", "# folder.mkdir(parents=True, exist_ok=True)\n", "# print(\"Checkpoint folder:\", CHECKPOINT_DIR)\n" ] }, { "cell_type": "markdown", "id": "20009eee", "metadata": { "id": "20009eee" }, "source": [ "## Phase 1: Inverted index and BM25\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-1)\n", "\n", "\n", "\n", "Implement BM25, run the tests, and compare Recall@5 with FTS5; explain design choices and quality gaps.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1ae7a7d0", "metadata": { "id": "1ae7a7d0" }, "outputs": [], "source": [ "class InvertedIndex:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Required attributes after build:\n", " self.postings: dict[str, list[tuple[int, int]]] # term -> [(doc_idx, tf), ...]\n", " self.doc_lengths: list[int] # tokens per doc\n", " self.avg_doc_length: float\n", " self.n_docs: int\n", " self.doc_ids: list[str] # doc_idx -> doc_id\n", " \"\"\"\n", "\n", " def __init__(self) -> None:\n", " self.postings: dict[str, list[tuple[int, int]]] = {}\n", " self.doc_lengths: list[int] = []\n", " self.avg_doc_length: float = 0.0\n", " self.n_docs: int = 0\n", " self.doc_ids: list[str] = []\n", "\n", " def build(self, documents: list[dict]) -> None:\n", " \"\"\"Ingest a list of {'doc_id': str, 'text': str} into the index.\"\"\"\n", " # YOUR CODE HERE.\n", " # 1. Iterate documents, tokenize, count term frequencies.\n", " # 2. Append (doc_idx, tf) to self.postings[term] for each term.\n", " # 3. Track doc_lengths and doc_ids.\n", " # 4. Compute avg_doc_length.\n", " raise NotImplementedError(\"Phase 1: implement InvertedIndex.build\")\n", "\n", " def df(self, term: str) -> int:\n", " return len(self.postings.get(term, []))\n", "\n", " def save(self, path: Path) -> None:\n", " with open(path, \"wb\") as f:\n", " pickle.dump({\n", " \"postings\": self.postings, \"doc_lengths\": self.doc_lengths,\n", " \"avg_doc_length\": self.avg_doc_length, \"n_docs\": self.n_docs,\n", " \"doc_ids\": self.doc_ids,\n", " }, f)\n", "\n", " @classmethod\n", " def load(cls, path: Path) -> \"InvertedIndex\":\n", " idx = cls()\n", " with open(path, \"rb\") as f:\n", " d = pickle.load(f)\n", " idx.postings = d[\"postings\"]; idx.doc_lengths = d[\"doc_lengths\"]\n", " idx.avg_doc_length = d[\"avg_doc_length\"]; idx.n_docs = d[\"n_docs\"]\n", " idx.doc_ids = d[\"doc_ids\"]\n", " return idx\n", "\n", "\n", "def bm25_query(index: InvertedIndex, query: str, top_k: int = 10,\n", " k1: float = 1.2, b: float = 0.75) -> list[tuple[str, float]]:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Tokenize the query, look up each term's posting list, score each doc that\n", " contains at least one query term using BM25, return top_k as\n", " [(doc_id, score), ...] sorted by score descending.\n", " \"\"\"\n", " # YOUR CODE HERE.\n", " raise NotImplementedError(\"Phase 1: implement bm25_query\")" ] }, { "cell_type": "markdown", "id": "6e9bbf29", "metadata": { "id": "6e9bbf29" }, "source": [ "### Phase 1 tests\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "14a424db", "metadata": { "id": "14a424db", "cellView": "form" }, "outputs": [], "source": [ "# @title Phase 1 tests and comparison helpers\n", "def run_phase_1_tests() -> None:\n", " idx = InvertedIndex()\n", " idx.build(TINY_CORPUS)\n", " assert idx.n_docs == 20, f\"expected 20 docs, got {idx.n_docs}\"\n", " assert idx.df(\"verify_jwt\") >= 2, f\"expected df('verify_jwt')>=2, got {idx.df('verify_jwt')}\"\n", " assert idx.df(\"redis\") == 2, f\"expected df('redis')==2, got {idx.df('redis')}\"\n", " print(\" ✓ inverted_index_build\")\n", "\n", " save_path = INDEX_DIR / \"phase1_test.pkl\"\n", " idx.save(save_path)\n", " idx2 = InvertedIndex.load(save_path)\n", " assert idx2.n_docs == idx.n_docs\n", " assert idx2.df(\"verify_jwt\") == idx.df(\"verify_jwt\")\n", " print(\" ✓ inverted_index_save_load\")\n", "\n", " results = bm25_query(idx, \"verify jwt token\", top_k=5)\n", " assert len(results) > 0\n", " assert \"auth/jwt.py:1\" in [d for d, _ in results][:3], f\"top-3 missing JWT doc: {results}\"\n", " print(\" ✓ bm25_basic_recall\")\n", "\n", " results = bm25_query(idx, \"verify_jwt_with_scope\", top_k=5)\n", " assert results[0][0] == \"auth/scopes.py:1\", f\"expected scopes doc first, got {results}\"\n", " print(\" ✓ bm25_identifier_match\")\n", "\n", " results = bm25_query(idx, \"kubernetes deployment yaml\", top_k=5)\n", " assert results == [] or all(s == 0 for _, s in results)\n", " print(\" ✓ bm25_no_match\")\n", "\n", " print(\"\\nPhase 1: ALL PASS\")\n", "\n", "\n", "def phase_1_parity(corpus: str = \"codesearchnet\", tier: str = \"25k\") -> None:\n", " \"\"\"Compare your BM25 against FTS5 on the same eval set.\n", "\n", " Reports the gap for analysis; there is no fixed recall-gap cutoff.\n", " \"\"\"\n", " print(f\" loading {corpus} {tier}...\")\n", " if corpus == \"codesearchnet\":\n", " docs, eval_set = load_codesearchnet(tier)\n", " else:\n", " docs, eval_set = load_msmarco(tier)\n", " print(f\" {len(docs):,} docs, {len(eval_set):,} eval queries\")\n", "\n", " inv_idx = InvertedIndex()\n", " inv_idx.build(docs)\n", "\n", " # SQLite reference\n", " db_path = INDEX_DIR / f\"phase1_parity_{corpus}.db\"\n", " if db_path.exists():\n", " db_path.unlink()\n", " conn = open_sqlite_db(db_path)\n", " embs_dummy = np.zeros((len(docs), EMBEDDING_DIM), dtype=np.float32)\n", " build_sqlite_corpus(conn, docs, embs_dummy)\n", "\n", " yours = run_eval(lambda q: bm25_query(inv_idx, q, top_k=10), eval_set)\n", " fts5 = run_eval(lambda q: fts5_query(conn, q, 10), eval_set)\n", " print(f\"\\n Your BM25 R@5: {yours['recall@5']:.3f} (p50 {yours['p50_ms']:.1f} ms)\")\n", " print(f\" FTS5 BM25 R@5: {fts5['recall@5']:.3f} (p50 {fts5['p50_ms']:.1f} ms)\")\n", " gap = fts5['recall@5'] - yours['recall@5']\n", " print(f\" Gap: {gap*100:+.1f} pts\")\n", " conn.close()" ] }, { "cell_type": "code", "execution_count": null, "id": "c4c2382b", "metadata": { "id": "c4c2382b" }, "outputs": [], "source": [ "# Run after implementing Phase 1.\n", "run_phase_1_tests()\n", "for corpus in (\"codesearchnet\", \"msmarco\"):\n", " phase_1_parity(corpus=corpus, tier=DATASET)\n" ] }, { "cell_type": "markdown", "id": "phase1-notes", "metadata": { "id": "phase1-notes" }, "source": [ "### Your notes: FTS5 comparison and design choices\n", "\n", "[Your answer]\n" ] }, { "cell_type": "markdown", "id": "425f0f26", "metadata": { "id": "425f0f26" }, "source": [ "## Phase 2: LSH vector search\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-2)\n", "\n", "\n", "\n", "Implement LSH, run the tests, and compare Recall@5 with sqlite-vec; explain the speed/recall tradeoff.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "24c6990d", "metadata": { "id": "24c6990d" }, "outputs": [], "source": [ "class LSHIndex:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Suggested structure:\n", " self.bands: list[tuple[np.ndarray, dict[int, list[int]]]]\n", " # for each band: (hyperplanes (bits_per_band, dim), signature -> [doc_idx, ...])\n", " self.embeddings: np.ndarray # (N, dim) for exact re-rank\n", " self.doc_ids: list[str]\n", " \"\"\"\n", "\n", " def __init__(self, n_bands: int = 60, bits_per_band: int = 3) -> None:\n", " self.n_bands = n_bands\n", " self.bits_per_band = bits_per_band\n", " self.bands: list[tuple[np.ndarray, dict[int, list[int]]]] = []\n", " self.embeddings: np.ndarray = np.zeros((0, EMBEDDING_DIM), dtype=np.float32)\n", " self.doc_ids: list[str] = []\n", "\n", " def build(self, doc_ids: list[str], embeddings: np.ndarray) -> None:\n", " \"\"\"YOUR CODE.\n", "\n", " 1. Generate `n_bands` random hyperplane matrices (each (bits_per_band, dim)).\n", " 2. Hash all embeddings against each band; bucket doc_idx by signature.\n", " 3. Store self.embeddings, self.doc_ids, self.bands.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 2: implement LSHIndex.build\")\n", "\n", " def query(self, query_vec: np.ndarray, top_k: int = 10,\n", " candidate_budget: int = 1000) -> list[tuple[str, float]]:\n", " \"\"\"YOUR CODE.\n", "\n", " 1. Hash query_vec against each band; collect candidate doc_idx from any matching bucket.\n", " 2. RANK candidates by hit count (how many bands they appeared in).\n", " Naive set-union-then-truncate breaks at scale. See FAQ.\n", " 3. Take top `candidate_budget` candidates by hit count.\n", " 4. Exact-cosine re-rank candidates against self.embeddings.\n", " 5. Return top_k as [(doc_id, score), ...].\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 2: implement LSHIndex.query\")" ] }, { "cell_type": "markdown", "id": "f86f9904", "metadata": { "id": "f86f9904" }, "source": [ "### Phase 2 tests\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9ce974a6", "metadata": { "id": "9ce974a6", "cellView": "form" }, "outputs": [], "source": [ "# @title Phase 2 tests and comparison helpers\n", "def run_phase_2_tests() -> None:\n", " docs = TINY_CORPUS\n", " embs = embed_batch([d[\"text\"] for d in docs])\n", "\n", " idx = LSHIndex(n_bands=10, bits_per_band=4)\n", " idx.build([d[\"doc_id\"] for d in docs], embs)\n", " assert len(idx.bands) == 10\n", " assert idx.embeddings.shape == (20, EMBEDDING_DIM)\n", " print(\" ✓ lsh_build_no_crash\")\n", "\n", " idx2 = LSHIndex(n_bands=20, bits_per_band=4)\n", " idx2.build([d[\"doc_id\"] for d in docs], embs)\n", " target_idx = 0\n", " results = idx2.query(embs[target_idx], top_k=5)\n", " result_ids = [d for d, _ in results]\n", " assert docs[target_idx][\"doc_id\"] in result_ids, (\n", " f\"self-retrieval failed: {docs[target_idx]['doc_id']} not in {result_ids}\"\n", " )\n", " print(\" ✓ lsh_self_retrieval\")\n", "\n", " idx3 = LSHIndex(n_bands=30, bits_per_band=3)\n", " idx3.build([d[\"doc_id\"] for d in docs], embs)\n", " queries = [\"how do I parse a config file\", \"verify a JWT token\", \"redis cache\"]\n", " overlaps = []\n", " for q in queries:\n", " qv = embed_text(q)\n", " bf_top = set(docs[i][\"doc_id\"] for i in brute_force_cosine(qv, embs, top_k=5))\n", " lsh_top = set(d for d, _ in idx3.query(qv, top_k=5))\n", " overlaps.append(len(bf_top & lsh_top) / 5)\n", " avg = sum(overlaps) / len(overlaps)\n", " assert avg >= 0.4, f\"avg LSH/brute overlap {avg:.2f} too low\"\n", " print(f\" ✓ lsh_recall_vs_brute_force (avg overlap {avg:.2f})\")\n", "\n", " print(\"\\nPhase 2: ALL PASS\")\n", "\n", "\n", "def phase_2_parity(corpus: str = \"codesearchnet\", tier: str = \"25k\") -> None:\n", " \"\"\"Compare your LSH against sqlite-vec's exact KNN on the same eval set.\"\"\"\n", " print(f\" loading {corpus} {tier}...\")\n", " if corpus == \"codesearchnet\":\n", " docs, eval_set = load_codesearchnet(tier)\n", " else:\n", " docs, eval_set = load_msmarco(tier)\n", " print(f\" {len(docs):,} docs, {len(eval_set):,} queries\")\n", " embs = cached_embed_batch([d[\"doc_id\"] for d in docs], [d[\"text\"] for d in docs],\n", " CACHE_DIR / f\"{corpus}_{tier}.parquet\")\n", "\n", " lsh_idx = LSHIndex(n_bands=60, bits_per_band=3)\n", " lsh_idx.build([d[\"doc_id\"] for d in docs], embs)\n", "\n", " db_path = INDEX_DIR / f\"phase2_parity_{corpus}.db\"\n", " if db_path.exists():\n", " db_path.unlink()\n", " conn = open_sqlite_db(db_path)\n", " build_sqlite_corpus(conn, docs, embs)\n", "\n", " yours = run_eval(lambda q: lsh_idx.query(embed_text(q), top_k=10), eval_set)\n", " vec = run_eval(lambda q: vec_query(conn, embed_text(q), 10), eval_set)\n", " print(f\"\\n Your LSH R@5: {yours['recall@5']:.3f} (p50 {yours['p50_ms']:.1f} ms)\")\n", " print(f\" sqlite-vec KNN R@5: {vec['recall@5']:.3f} (p50 {vec['p50_ms']:.1f} ms)\")\n", " gap = vec['recall@5'] - yours['recall@5']\n", " print(f\" Gap: {gap*100:+.1f} pts\")\n", " conn.close()" ] }, { "cell_type": "code", "execution_count": null, "id": "97c38c2e", "metadata": { "id": "97c38c2e" }, "outputs": [], "source": [ "# Run after implementing Phase 2.\n", "run_phase_2_tests()\n", "for corpus in (\"codesearchnet\", \"msmarco\"):\n", " phase_2_parity(corpus=corpus, tier=DATASET)\n" ] }, { "cell_type": "markdown", "id": "phase2-notes", "metadata": { "id": "phase2-notes" }, "source": [ "### Your notes: sqlite-vec comparison and speed/recall tradeoff\n", "\n", "[Your answer]\n" ] }, { "cell_type": "markdown", "id": "24c77a3d", "metadata": { "id": "24c77a3d" }, "source": [ "## Phase 3: Hybrid retrieval and query planner\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-3)\n", "\n", "\n", "\n", "Implement RRF, routing, and custom fusion; run the tests and explain your custom method.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "98853915", "metadata": { "id": "98853915" }, "outputs": [], "source": [ "def reciprocal_rank_fusion(ranked_lists: list[list[tuple[str, float]]],\n", " k: int = 60, top_k: int = 10) -> list[tuple[str, float]]:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Standard RRF: for each (doc_id, score) appearing in any ranked_list at rank r,\n", " add 1.0 / (k + r) to its fused score. Sort, return top_k.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 3: implement reciprocal_rank_fusion\")\n", "\n", "\n", "def hybrid_query(inv_idx: InvertedIndex, lsh_idx: LSHIndex, query: str,\n", " top_k: int = 10) -> list[tuple[str, float]]:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Run BM25 + LSH, fuse with RRF, return top_k.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 3: implement hybrid_query\")\n", "\n", "\n", "def query_planner(query: str) -> str:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " Inspect the query text. Return one of: 'bm25', 'lsh', 'hybrid'.\n", "\n", " Hints:\n", " - Contains snake_case or camelCase identifiers? -> probably bm25\n", " - Looks like natural language (\"how do I\", \"what is\")? -> probably lsh\n", " - Mixed? -> hybrid\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 3: implement query_planner\")\n", "\n", "\n", "def custom_fusion(inv_idx: InvertedIndex, lsh_idx: LSHIndex, query: str,\n", " top_k: int = 10) -> list[tuple[str, float]]:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " A fusion variant of YOUR design that is meaningfully different from RRF.\n", " Examples: weighted score fusion (normalize then combine), confidence\n", " fall-through (use BM25 if it is confident, else fuse), agreement boost\n", " (RRF + extra weight when both signals agree on the top-1).\n", "\n", " Document what you did and why in a comment at the top of this function.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 3: implement custom_fusion\")" ] }, { "cell_type": "markdown", "id": "469f068c", "metadata": { "id": "469f068c" }, "source": [ "### Phase 3 tests\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e6ea80a4", "metadata": { "id": "e6ea80a4", "cellView": "form" }, "outputs": [], "source": [ "# @title Phase 3 tests\n", "def run_phase_3_tests() -> None:\n", " a = [(\"doc1\", 0.9), (\"doc2\", 0.5), (\"doc3\", 0.1)]\n", " b = [(\"doc2\", 0.8), (\"doc4\", 0.7), (\"doc1\", 0.3)]\n", " fused = reciprocal_rank_fusion([a, b], k=60, top_k=4)\n", " fused_ids = [d for d, _ in fused]\n", " assert set(fused_ids[:2]) == {\"doc1\", \"doc2\"}, f\"top-2 should be doc1,doc2: {fused_ids}\"\n", " print(\" ✓ rrf_basic\")\n", "\n", " a2 = [(\"doc1\", 0.9), (\"doc2\", 0.5)]\n", " fused2 = reciprocal_rank_fusion([a2], k=60, top_k=2)\n", " assert [d for d, _ in fused2] == [\"doc1\", \"doc2\"]\n", " print(\" ✓ rrf_single_list\")\n", "\n", " assert query_planner(\"verify_jwt_with_scope\") == \"bm25\"\n", " assert query_planner(\"how do I parse a config file\") == \"lsh\"\n", " print(\" ✓ planner_routes\")\n", "\n", " print(\"\\nPhase 3: ALL PASS\")" ] }, { "cell_type": "code", "execution_count": null, "id": "0daa3122", "metadata": { "id": "0daa3122" }, "outputs": [], "source": [ "# Run after implementing Phase 3.\n", "run_phase_3_tests()\n" ] }, { "cell_type": "markdown", "id": "phase3-notes", "metadata": { "id": "phase3-notes" }, "source": [ "### Your notes: How your custom fusion works\n", "\n", "[Your answer]\n" ] }, { "cell_type": "markdown", "id": "758e69cd", "metadata": { "id": "758e69cd" }, "source": [ "## Phase 4: Two-dataset benchmark\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-4)\n", "\n", "\n", "\n", "Benchmark all five methods on both datasets using identical queries; save recall, MRR@10, latency, and memory results.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cdbba183", "metadata": { "id": "cdbba183" }, "outputs": [], "source": [ "def run_two_corpus_benchmark(tier: str = \"25k\") -> dict:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " For each corpus in [codesearchnet, msmarco]:\n", " 1. Load with the supplied loader.\n", " 2. Compute embeddings.\n", " 3. Build both indexes (InvertedIndex, LSHIndex).\n", " 4. Run each retriever (bm25, lsh, hybrid, custom_fusion, planner-routed)\n", " through run_eval against the eval_set.\n", " 5. Print a side-by-side table.\n", "\n", " Return a nested dict: {corpus_name: {retriever_name: metrics}}.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 4: implement run_two_corpus_benchmark\")\n", "\n", "\n", "def run_phase_4_tests() -> None:\n", " \"\"\"Sanity check the driver runs end-to-end on tiny.\"\"\"\n", " out = run_two_corpus_benchmark(tier=\"tiny\")\n", " assert \"codesearchnet\" in out and \"msmarco\" in out\n", " for corpus, retrievers in out.items():\n", " assert \"bm25\" in retrievers\n", " assert \"lsh\" in retrievers\n", " for name, m in retrievers.items():\n", " assert \"recall@5\" in m and \"p50_ms\" in m\n", " print(\"\\nPhase 4: ALL PASS\")" ] }, { "cell_type": "code", "execution_count": null, "id": "57dcf4cd", "metadata": { "id": "57dcf4cd" }, "outputs": [], "source": [ "# Run after implementing the benchmark driver.\n", "run_phase_4_tests()\n", "benchmark_results = run_two_corpus_benchmark(tier=DATASET)\n" ] }, { "cell_type": "markdown", "id": "a3086d4f", "metadata": { "id": "a3086d4f" }, "source": [ "## Phase 5: SQLite comparison\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-5)\n", "\n", "\n", "\n", "Add FTS5 and sqlite-vec to the comparison using identical data and queries; reuse Phase 4 results.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "4eccf11d", "metadata": { "id": "4eccf11d" }, "outputs": [], "source": [ "def run_production_scoreboard(tier: str = \"25k\") -> dict:\n", " \"\"\"YOU IMPLEMENT.\n", "\n", " For each corpus, on top of your Phase 4 results, add:\n", " - FTS5 BM25 metrics (use the supplied fts5_query)\n", " - sqlite-vec exact KNN metrics (use the supplied vec_query)\n", "\n", " Print one row per retriever, with columns for recall, MRR, latency, and memory.\n", "\n", " Return a dict of the same shape as Phase 4, with the production rows added.\n", " \"\"\"\n", " raise NotImplementedError(\"Phase 5: implement run_production_scoreboard\")\n", "\n", "\n", "def run_phase_5_tests() -> None:\n", " out = run_production_scoreboard(tier=\"tiny\")\n", " assert out, \"Missing benchmark results\"\n", " for corpus, retrievers in out.items():\n", " assert \"FTS5 BM25\" in retrievers or \"fts5_bm25\" in retrievers, f\"{corpus}: missing FTS5 row\"\n", " assert \"sqlite-vec KNN\" in retrievers or \"sqlite_vec_knn\" in retrievers, f\"{corpus}: missing sqlite-vec row\"\n", " print(\"\\nPhase 5: ALL PASS\")" ] }, { "cell_type": "code", "execution_count": null, "id": "072b503e", "metadata": { "id": "072b503e" }, "outputs": [], "source": [ "# Run after implementing the SQLite comparison driver.\n", "run_phase_5_tests()\n", "comparison_results = run_production_scoreboard(tier=DATASET)\n" ] }, { "cell_type": "markdown", "id": "e177ea65", "metadata": { "id": "e177ea65" }, "source": [ "## Phase 6: Findings and AI use\n", "\n", "[Instructions](https://cs145.web.app/projects/project2-memory.html#phase-6)\n", "\n", "Answer these three questions in half a page to one page total. Reuse your notes from earlier phases. Figures, code, and benchmark tables are outside this limit.\n", "\n", "### 1. What did your results show?\n", "\n", "[Your findings, supported by the benchmark results.]\n", "\n", "### 2. How did you use and check AI assistance?\n", "\n", "[Give 2–3 examples. If you did not use AI, describe 2–3 debugging or design decisions and how you tested them.]\n", "\n", "### 3. How does one of your algorithms work?\n", "\n", "[Explain it in plain English.]\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "3295145f", "metadata": { "id": "3295145f" }, "source": [ "---\n", "## Final check (after Phases 0–5)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "40d2e964", "metadata": { "lines_to_next_cell": 2, "id": "40d2e964", "cellView": "form" }, "outputs": [], "source": [ "# @title Final test helper\n", "def run_all_tests(phase0_state) -> None:\n", " print(\"=\" * 60)\n", " print(\"RUNNING ALL PHASE TESTS\")\n", " print(\"=\" * 60)\n", " print(\"\\nPhase 0:\")\n", " run_phase_0_tests(phase0_state)\n", " print(\"\\nPhase 1:\")\n", " run_phase_1_tests()\n", " print(\"\\nPhase 2:\")\n", " run_phase_2_tests()\n", " print(\"\\nPhase 3:\")\n", " run_phase_3_tests()\n", " print(\"\\nPhase 4:\")\n", " run_phase_4_tests()\n", " print(\"\\nPhase 5:\")\n", " run_phase_5_tests()\n", " print(\"\\n\" + \"=\" * 60)\n", " print(\"ALL PHASES PASS\")\n", " print(\"=\" * 60)" ] }, { "cell_type": "code", "execution_count": null, "id": "b10c7ad1", "metadata": { "id": "b10c7ad1" }, "outputs": [], "source": [ "# Final smoke check after completing Phases 0–5. Reuses Phase 0 data.\n", "run_all_tests(phase0)\n" ] }, { "cell_type": "markdown", "id": "2e0826ed", "metadata": { "id": "2e0826ed" }, "source": [ "## Scaling (Tiers 93 and 100)\n", "\n", "[Scale, memory limits, and evaluation requirements](https://cs145.web.app/projects/project2-memory.html#tiers-at-a-glance)\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "4566c012", "metadata": { "id": "4566c012" }, "source": [ "### Tier 93: Build and query within the memory budget\n", "\n", "[Data source and evaluation](https://cs145.web.app/projects/project2-memory.html#tier-93-data) · [Partitioning](https://cs145.web.app/projects/project2-memory.html#partitioning-guide)\n", "\n", "Use the supplied loader below: **709,254 documents (5 GB of title/body text)**, **50 development queries**, and **500 final queries**. Everyone uses the same data and labels. No T4 or Colab runtime is required for this tier.\n", "\n", "Run the preparation cell below once, then add your bounded-memory build, seven-method comparison, and memory evidence in this section. Reuse these saved indexes for the optional bonus; no second build is needed if they cover this exact subset.\n", "\n", "### Fixed MS MARCO data and saved IDs\n", "\n", "### Bounded-memory build and retrieval\n", "\n", "### Development (50 queries) and final evaluation (500 separate queries)\n", "\n", "### Larger-corpus benchmark and SQLite comparison\n", "\n", "For SQLite, call `initialize_sqlite_corpus(conn)` once, then `append_sqlite_batch(conn, docs, embeddings)` for each batch. Both helpers are supplied.\n", "\n", "### Build/query memory and disk reads\n", "\n", "Keep the Tier 87 results in Phases 4–5. Add the larger run here.\n", "\n", "\n" ] }, { "cell_type": "code", "id": "tier93-prepare", "metadata": { "id": "tier93-prepare" }, "source": [ "# Higher tiers only: downloads/prepares about 5 GB of text on first use.\n", "# Leave False for Phase 0 and Tiers 80/87. Completed data is verified and reused.\n", "RUN_TIER93 = False\n", "if RUN_TIER93:\n", " scale_utils, tier93_data = prepare_tier93_data()\n", " # Iterate tier93_data.documents(); do not load it all with list().\n", " # Use tier93_data.development for tuning and tier93_data.evaluation for final results.\n", " print(\"Documents:\", tier93_data.config[\"document_count\"])\n", " print(\"Text bytes:\", tier93_data.config[\"text_bytes\"])\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "tier93-student-work" }, "source": [ "# YOUR WORK: Tier 93 build, evaluation, and memory measurements.\n", "# Add your implementation and keep the resulting outputs visible here.\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "c4c85686", "metadata": { "id": "c4c85686" }, "source": [ "### Tier 100: Full MS MARCO document corpus\n", "\n", "[Requirements](https://cs145.web.app/projects/project2-memory.html#tiers-at-a-glance) · [Memory reference](https://cs145.web.app/projects/project2-memory.html#bounded-memory-guide)\n", "\n", "### Full MS MARCO document build and checkpoints\n", "\n", "### Full-corpus benchmark and SQLite comparison\n", "\n", "### Build/query memory and disk reads\n", "\n", "Keep the Tier 93 results above. Add the full MS MARCO v1 document-corpus run here (3,213,835 documents). Reuse the Tier 93 development and final queries and their relevance labels; do not tune on final queries. Build and query with a working set below 1 GB.\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "9e98709e", "metadata": { "lines_to_next_cell": 2, "id": "9e98709e" }, "source": [ "### Full-corpus evaluation\n", "\n", "Reuse the saved Tier 93 final query set and document relevance labels with `run_eval()`. Evaluate all five methods and both SQLite references against the full MS MARCO document corpus.\n", "\n", "\n" ] }, { "cell_type": "code", "metadata": { "id": "tier100-student-work" }, "source": [ "# YOUR WORK: Tier 100 build, evaluation, and memory measurements.\n", "# Add your implementation and keep the resulting outputs visible here.\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "27b580fb", "metadata": { "id": "27b580fb" }, "source": [ "## Optional CLI scaffold\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9c5062a3", "metadata": { "id": "9c5062a3", "cellView": "form" }, "outputs": [], "source": [ "# @title Optional CLI template\n", "CLI_TEMPLATE = '''\n", "\"\"\"nanomem CLI. Usage:\n", " python nanomem.py index --corpus path/to/corpus/\n", " python nanomem.py query \"how do I parse a config file\"\n", "\"\"\"\n", "import argparse, json\n", "from pathlib import Path\n", "\n", "INDEX_DIR = Path(\"./nanomem_index\")\n", "\n", "def cmd_index(args):\n", " # Load corpus, build inverted index + LSH index, save to INDEX_DIR.\n", " raise NotImplementedError\n", "\n", "def cmd_query(args):\n", " # Load indexes, run hybrid_query, print top-K as JSON.\n", " raise NotImplementedError\n", "\n", "if __name__ == \"__main__\":\n", " p = argparse.ArgumentParser(prog=\"nanomem\")\n", " sub = p.add_subparsers(dest=\"cmd\", required=True)\n", " p_index = sub.add_parser(\"index\"); p_index.add_argument(\"--corpus\", required=True)\n", " p_query = sub.add_parser(\"query\"); p_query.add_argument(\"query_text\", type=str)\n", " args = p.parse_args()\n", " if args.cmd == \"index\": cmd_index(args)\n", " elif args.cmd == \"query\": cmd_query(args)\n", "'''\n" ] }, { "cell_type": "markdown", "id": "bonus-v1-0", "metadata": { "id": "bonus-v1-0" }, "source": [ "## Optional bonus benchmark\n", "\n", "**Skip this section unless entering the bonus competition.** It is disabled by default and adds no required work.\n", "\n", "- **Run:** Use the optional benchmark code below with your Tier 93 data and indexes.\n", "- **Save:** Keep the report and embedded benchmark record in your submitted notebook; no separate attachment is needed.\n", "- **Review:** Staff review the results and finalize awards for the top five valid entries (0–5 points each).\n", "\n", "### Benchmark setup and measurement rules\n", "\n", "**Reuse your Tier 93 data and saved indexes.** [Prepare the required dataset in Tier 93](#scrollTo=4566c012) before running this section. This competition times only your planner; the required Tier 93 evaluation compares all five methods and both SQLite references.\n", "\n", "1. **Runtime:** In Colab, choose **Runtime → Change runtime type → Python 3 → T4 GPU**, with **standard RAM**. Reconnect and rerun Setup. This helper checks the GPU/RAM and sets CPU math threads to 1; it records the actual CPU, GPU, RAM, and package versions. Local runtimes do not qualify.\n", "2. **Data and queries:** The supplied benchmark selects **709,254 MS MARCO v1 full documents**, containing **5,000,000,958 bytes of title/body text**. Everyone uses identical document IDs, 50 development queries, and 500 evaluation queries. Tune only on the development set. The first preparation downloads several GB and writes about 5 GB locally; interrupted downloads restart. Save completed data/indexes to Drive if needed. For timed runs, copy your index to `/content` so Drive latency is not part of the comparison.\n", "3. **Your engine:** Build or load your own index over this entire subset. Reuse your Tier 93 indexes if they cover this exact subset; otherwise build once using your Tier 93 code. Connect the adapter functions in the next cell. Save build-time memory evidence to a text file. Return original **document IDs**, merging chunk hits if you chunk documents. Count all live index/model/buffer components against the **2 GB** limit, including GPU allocations. Use the supplied MiniLM model; use `EMBED_DEVICE = \"cuda\"` before building/loading the embedding model if your planner uses vector retrieval.\n", "4. **Measure and submit:** Ten development queries warm up the engine, then all 500 evaluation queries run once in a fixed shuffled order. Timing includes query embedding, retrieval, disk reads, and GPU synchronization; index construction is excluded. Do not cache answers or query embeddings, tune on evaluation queries, or hard-code relevant IDs. Keep the displayed report and embedded benchmark record in this notebook; no separate attachment is needed. The score is **mean Recall@5 ÷ median query latency (ms)**. Staff review the saved results and finalize awards.\n", "\n", "Colab CPUs and host load vary even with a T4. The report records these differences for staff review; the runtime checks do not establish identical hardware or prove the live-memory bound. The top five valid entries receive 0–5 bonus points each.\n", "\n", "\n" ] }, { "cell_type": "code", "id": "bonus-v1-1", "metadata": { "id": "bonus-v1-1" }, "source": [ "# 1. Optional competition runtime. Tier 93 setup provides data and helpers.\n", "RUN_BONUS = False\n", "if RUN_BONUS:\n", " if \"tier93_data\" not in globals():\n", " raise RuntimeError(\"Run the Tier 93 data preparation cell first.\")\n", " bonus = scale_utils\n", " benchmark = tier93_data\n", " bonus.configure_bonus_runtime()\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "id": "bonus-v1-2", "metadata": { "id": "bonus-v1-2" }, "source": [ "# 2. Connect your Tier 93 implementation. These adapters are student work.\n", "# Build once; load saved indexes on later runs. Do not use the tiny-data engine.\n", "def load_or_build_bonus_engine(documents):\n", " # documents is an iterator of {\"doc_id\": ..., \"text\": ...}; do not list() it.\n", " # Build/load all documents, saving build-memory logs and source document IDs.\n", " raise NotImplementedError(\"Connect your bounded-memory index build/load code\")\n", "\n", "def bonus_search(query, k):\n", " # Route this raw-text query through your planner, then run your retrieval.\n", " # Return up to k original document IDs, or (document_id, score) pairs.\n", " raise NotImplementedError(\"Connect planner-routed retrieval using bonus_engine\")\n", "\n", "def bonus_indexed_document_ids():\n", " # Yield each original document ID actually indexed, once (not the input manifest).\n", " raise NotImplementedError(\"Read IDs from your saved index metadata\")\n", "\n", "def bonus_memory_components():\n", " # Return {\"component_name\": live_bytes}, including the model, index, and buffers.\n", " # Values are BYTES. Supplement these query-boundary snapshots with build logs.\n", " raise NotImplementedError(\"Connect your live-component memory accounting\")\n", "\n", "BUILD_MEMORY_EVIDENCE = \"/content/bonus_build_memory.txt\"\n", "if RUN_BONUS:\n", " bonus_engine = load_or_build_bonus_engine(benchmark.documents())\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "id": "bonus-v1-3", "metadata": { "id": "bonus-v1-3" }, "source": [ "# 3. Run and save results inside this notebook. Keep both outputs below.\n", "if RUN_BONUS:\n", " import base64, contextlib, io\n", " from IPython.display import display\n", " # Capture the helper's legacy ZIP instructions; submission is now notebook-only.\n", " with contextlib.redirect_stdout(io.StringIO()):\n", " bonus_report = bonus.run_bonus_benchmark(\n", " search=bonus_search,\n", " benchmark=benchmark,\n", " indexed_document_ids=bonus_indexed_document_ids(),\n", " memory_components=bonus_memory_components,\n", " build_memory_evidence=BUILD_MEMORY_EVIDENCE,\n", " output_dir=\"/content/bonus_results\",\n", " )\n", " bonus.display_bonus_report(bonus_report)\n", " archive = Path(\"/content/bonus_results.zip\").read_bytes()\n", " display({\n", " \"application/json\": {\"nanomem_bonus_archive\": base64.b64encode(archive).decode(\"ascii\")},\n", " \"text/plain\": \"Benchmark record saved in this output. Keep it in your submitted notebook; no separate upload is needed.\",\n", " }, raw=True)\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "id": "gradescope-final-submission", "metadata": {}, "source": [ "## Submitting to Gradescope\n", "\n", "- Submit your final notebook(s) as both `.ipynb` and PDF files, with saved outputs visible.\n", "- Tag the PDF pages to the associated Gradescope questions.\n", "- In Colab, choose **File → Print**, then **Save as PDF**.\n" ] } ], "metadata": { "jupytext": { "cell_metadata_filter": "-all", "main_language": "python", "notebook_metadata_filter": "-all" }, "colab": { "name": "cs145_nanomem_2026-10-12.ipynb" } }, "nbformat": 4, "nbformat_minor": 5 }