Modern LLMs17 / 20

Module 17. Embeddings, Retrieval, and RAG

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 2588.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Notebook 17. Retrieval and RAG: from documents to cited context

Main module notebook

Download .ipynb module_17_notebook_rag.ipynb · 64.6 KB

Show notebook text

Notebook 17. Retrieval and RAG: from documents to cited context

Module 17 — hands-on notebook

We will turn source documents into traceable chunks, build two search branches, and inspect how their ranks contribute to a fused list. A reranker revisits only the surviving candidates. After measuring what was retained, we assemble whole passages and citations within a budget measured in regex word units. Each stage exposes the decisions that the next stage inherits.

The dense branch uses TF-IDF followed by truncated SVD (LSA). It is genuine exact cosine search over a dense matrix, but not a neural embedder. The proxy keeps the mechanics inspectable and needs no model download. A production system may replace LSA with a bi-encoder while preserving the surrounding shape, ranking, metric, fusion, and context-assembly contracts.

Colab contract: Runtime → Run all completes the mandatory route on CPU in a clean session. The notebook reads no repository files and needs no upload, Google Drive, API key, or paid service. A frozen Module-17 numerical contract is embedded and checked by SHA-256.

Route:

safe chunking with an explicit empty-document policy;
BM25 with deterministic tie-breaking;
a dense LSA index and exact cosine retrieval;
RRF accumulation, contribution traces, and final sorting;
a transparent MaxSim-style reranker;
MRR, Recall@5, and nDCG@5 with qrel validation;
context packing with deduplication, source limits, and citations;
MiB/GiB, ANN arithmetic, and RAG-versus-long-context accounting;
negative tests for every public contract.

Python

import importlib.util
import subprocess
import sys


def ensure_package(import_name: str, pip_name: str | None = None) -> None:
    """Install one missing package without touching an already prepared runtime."""

    if importlib.util.find_spec(import_name) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pip_name or import_name])


ensure_package("numpy")

import hashlib
import json
import math
import re
from collections import Counter
from dataclasses import dataclass
from typing import Mapping, Sequence

import numpy as np

np.set_printoptions(precision=5, suppress=True)
print('Environment ready.')

Thirteen authored diagnostic documents and seven development queries are frozen independently of the lecture. They are not a held-out benchmark. Tuning on these queries cannot establish generalization. The language collections use different wording, so their scores do not measure a language advantage. Our tokenizer extracts words with a regular expression; context budgets below use these word units, not any LLM's subword tokens.

Data and scope

Python

import json, math, re, hashlib
from collections import Counter
from dataclasses import dataclass, asdict
from typing import Mapping, Sequence
import numpy as np
WORD_RE = re.compile(r"[^\W_]+", re.UNICODE)
REFERENCE_JSON = '{"ann": {"ann_flops": 4096000.0, "ann_milliseconds": 0.40959999999999996, "exact_flops": 20480000000.0, "exact_seconds": 2.048, "speedup": 5000.0}, "cascade": {"ann_ms": 0.40959999999999996, "context_tokens": 3712, "decode_ms": 1432.6647564469915, "decode_share": 0.8639290480052048, "embed_ms": 0.07757575757575758, "prefill_ms": 120.43377777777778, "rerank_ms": 104.72727272727272, "retrieval_before_prefill_share": 0.06344667718450364, "total_ms": 1658.3129827096177}, "chunk_example": [{"chunk_id": "demo:c000", "end_char": 53, "end_token": 8, "source_id": "demo", "start_char": 0, "start_token": 0, "text": "retrieval first splits a long source into overlapping", "token_count": 8}, {"chunk_id": "demo:c001", "end_char": 94, "end_token": 13, "source_id": "demo", "start_char": 30, "start_token": 5, "text": "source into overlapping windows while keeping source identifiers", "token_count": 8}, {"chunk_id": "demo:c002", "end_char": 125, "end_token": 18, "source_id": "demo", "start_char": 68, "start_token": 10, "text": "keeping source identifiers offsets and citations so later", "token_count": 8}, {"chunk_id": "demo:c003", "end_char": 159, "end_token": 22, "source_id": "demo", "start_char": 107, "start_token": 15, "text": "citations so later context packing remains traceable", "token_count": 7}], "corpus": {"binary_mib": 7.26611328125, "chunk_size": 400, "chunks": 59524, "encoder_seconds_floor": 57.72024242424242, "fp32_mib": 232.515625, "int8_mib": 58.12890625, "overlap": 64, "total_tokens": 20000000}, "datasets": {"en": {"documents": [{"document_id": "d00", "source_id": "ann", "text": "HNSW builds a multi-layer small-world graph for approximate nearest-neighbor search over vectors."}, {"document_id": "d01", "source_id": "ann", "text": "IVF PQ partitions vectors into clusters and quantizes residuals, reducing index memory for hundreds of millions of objects."}, {"document_id": "d02", "source_id": "hybrid", "text": "Reciprocal rank fusion RRF combines BM25 and dense-search rankings without calibrating their scores."}, {"document_id": "d03", "source_id": "sparse", "text": "BM25 uses term frequency, inverse document frequency, and document-length normalization."}, {"document_id": "d04", "source_id": "dense", "text": "A dense bi encoder encodes the query and document independently as vectors and compares them with cosine similarity."}, {"document_id": "d05", "source_id": "rerank", "text": "A cross encoder reads the query and document jointly and scores relevance more accurately, but needs a separate forward pass for every pair."}, {"document_id": "d06", "source_id": "chunking", "text": "Overlap between adjacent chunks preserves local context at a document boundary."}, {"document_id": "d07", "source_id": "chunking", "text": "Late chunking first encodes the long document as a whole, so pronouns and anaphora retain context after splitting."}, {"document_id": "d08", "source_id": "routing", "text": "Long context places the entire long document in the prompt and suits holistic analysis of one file."}, {"document_id": "d09", "source_id": "rag", "text": "RAG retrieves relevant corpus passages, assembles a cited context, and sends it to the generator."}, {"document_id": "d10", "source_id": "metrics", "text": "nDCG supports graded relevance and applies a logarithmic discount to lower ranks."}, {"document_id": "d11", "source_id": "metrics", "text": "MRR measures the reciprocal rank of the first relevant result and ignores additional relevant documents."}, {"document_id": "d12", "source_id": "packing", "text": "Context packing removes duplicates, respects a token budget, and preserves source citations."}], "queries": [{"query_id": "q0", "relevance": {"d02": 3, "d09": 1}, "text": "how to merge bm25 and dense results"}, {"query_id": "q1", "relevance": {"d00": 3, "d01": 1}, "text": "a graph for approximate nearest neighbor search"}, {"query_id": "q2", "relevance": {"d10": 1, "d11": 3}, "text": "position of the first useful result"}, {"query_id": "q3", "relevance": {"d06": 2, "d07": 3}, "text": "preserve anaphora after splitting a document"}, {"query_id": "q4", "relevance": {"d04": 1, "d05": 3}, "text": "why a cross encoder cannot score the whole corpus"}, {"query_id": "q5", "relevance": {"d08": 3, "d09": 2}, "text": "when to read the whole long file instead of retrieval"}, {"query_id": "q6", "relevance": {"d09": 2, "d12": 3}, "text": "pack retrieved passages with citations into a token limit"}]}}, "long_context_vs_rag": {"kv_bytes_per_token": 131072, "long_kv_gib": 24.4140625, "long_prefill_seconds": 6.488888888888889, "long_tokens": 200000, "rag_kv_gib": 0.453125, "rag_prefill_seconds": 0.12043377777777778, "rag_tokens": 3712, "ratio": 53.87931034482759}, "retrieval_results": {"en": {"context_pack_q6": {"citations": ["[packing:d12]", "[rag:d09]"], "max_tokens": 42, "rendered_text": "[packing:d12] Context packing removes duplicates, respects a token budget, and preserves source citations.\\n\\n[rag:d09] RAG retrieves relevant corpus passages, assembles a cited context, and sends it to the generator.", "selected_ids": ["d12", "d09"], "skipped": {"d01": "budget", "d03": "budget", "d04": "budget", "d05": "budget", "d06": "budget", "d10": "budget"}, "used_tokens": 31}, "metrics": {"bm25": {"mrr": 1.0, "ndcg_at_5": 0.9706318690479474, "recall_at_5": 0.9285714285714286}, "dense_lsa": {"mrr": 0.9047619047619048, "ndcg_at_5": 0.7301304500524601, "recall_at_5": 0.7142857142857143}, "reranked": {"mrr": 1.0, "ndcg_at_5": 0.969811453183001, "recall_at_5": 0.9285714285714286}, "rrf": {"mrr": 0.9285714285714286, "ndcg_at_5": 0.8546325517881126, "recall_at_5": 0.8571428571428571}}, "rankings": {"bm25": {"q0": ["d02", "d03", "d10", "d09", "d04", "d05", "d12", "d11", "d08", "d01", "d07", "d00", "d06"], "q1": ["d00", "d02", "d05", "d01", "d06", "d10", "d12", "d09", "d04", "d07", "d03", "d08", "d11"], "q2": ["d11", "d08", "d07", "d01", "d09", "d04", "d05", "d00", "d02", "d03", "d06", "d10", "d12"], "q3": ["d07", "d06", "d05", "d03", "d04", "d08", "d10", "d12", "d00", "d09", "d01", "d02", "d11"], "q4": ["d05", "d09", "d07", "d04", "d11", "d08", "d06", "d10", "d12", "d00", "d01", "d02", "d03"], "q5": ["d08", "d07", "d09", "d11", "d10", "d01", "d04", "d05", "d00", "d02", "d03", "d06", "d12"], "q6": ["d12", "d09", "d04", "d01", "d05", "d06", "d10", "d00", "d07", "d02", "d03", "d08", "d11"]}, "dense_lsa": {"q0": ["d02", "d09", "d03", "d11", "d04", "d00", "d12", "d10", "d06", "d07", "d05", "d08", "d01"], "q1": ["d00", "d02", "d11", "d10", "d04", "d09", "d03", "d05", "d01", "d12", "d06", "d07", "d08"], "q2": ["d11", "d08", "d09", "d07", "d12", "d03", "d06", "d04", "d01", "d02", "d10", "d00", "d05"], "q3": ["d06", "d07", "d12", "d03", "d08", "d04", "d09", "d11", "d10", "d05", "d02", "d01", "d00"], "q4": ["d04", "d03", "d10", "d12", "d06", "d05", "d07", "d09", "d08", "d11", "d02", "d00", "d01"], "q5": ["d08", "d07", "d12", "d06", "d09", "d11", "d03", "d04", "d01", "d10", "d05", "d02", "d00"], "q6": ["d04", "d03", "d12", "d10", "d06", "d11", "d09", "d08", "d07", "d01", "d05", "d00", "d02"]}, "reranked": {"q0": ["d02", "d03", "d04", "d09", "d10", "d11", "d12", "d00", "d05"], "q1": ["d00", "d05", "d10", "d09", "d01", "d04", "d02", "d11", "d06", "d03", "d12"], "q2": ["d11", "d08", "d07", "d01", "d04", "d09", "d03", "d12", "d05", "d06", "d00"], "q3": ["d07", "d06", "d05", "d04", "d12", "d08", "d03", "d09", "d10", "d11"], "q4": ["d05", "d07", "d09", "d04", "d10", "d06", "d12", "d03", "d11", "d08"], "q5": ["d08", "d07", "d09", "d11", "d10", "d04", "d12", "d06", "d01", "d03", "d05"], "q6": ["d12", "d09", "d04", "d05", "d10", "d01", "d06", "d03", "d11", "d00", "d08"]}, "rrf": {"q0": ["d02", "d03", "d09", "d04", "d10", "d11", "d12", "d00", "d05"], "q1": ["d00", "d02", "d10", "d05", "d09", "d11", "d01", "d04", "d06", "d03", "d12"], "q2": ["d11", "d08", "d07", "d09", "d04", "d01", "d12", "d03", "d05", "d06", "d00"], "q3": ["d06", "d07", "d03", "d12", "d04", "d08", "d05", "d09", "d10", "d11"], "q4": ["d04", "d05", "d09", "d07", "d10", "d06", "d03", "d12", "d11", "d08"], "q5": ["d08", "d07", "d09", "d11", "d04", "d12", "d06", "d10", "d01", "d03", "d05"], "q6": ["d04", "d12", "d09", "d10", "d06", "d03", "d01", "d05", "d11", "d00", "d08"]}}, "rrf_trace_q0": [{"contributions": {"bm25": 0.01639344262295082, "dense_lsa": 0.01639344262295082}, "document_id": "d02", "rank": 1, "score": 0.03278688524590164, "source_ranks": {"bm25": 1, "dense_lsa": 1}}, {"contributions": {"bm25": 0.016129032258064516, "dense_lsa": 0.015873015873015872}, "document_id": "d03", "rank": 2, "score": 0.03200204813108039, "source_ranks": {"bm25": 2, "dense_lsa": 3}}, {"contributions": {"bm25": 0.015625, "dense_lsa": 0.016129032258064516}, "document_id": "d09", "rank": 3, "score": 0.031754032258064516, "source_ranks": {"bm25": 4, "dense_lsa": 2}}, {"contributions": {"bm25": 0.015384615384615385, "dense_lsa": 0.015384615384615385}, "document_id": "d04", "rank": 4, "score": 0.03076923076923077, "source_ranks": {"bm25": 5, "dense_lsa": 5}}, {"contributions": {"bm25": 0.015873015873015872, "dense_lsa": 0.014705882352941176}, "document_id": "d10", "rank": 5, "score": 0.03057889822595705, "source_ranks": {"bm25": 3, "dense_lsa": 8}}]}}, "schema_version": 2, "storage": {"binary_gib": 1.1920928955078125, "dimension": 1024, "fp16_gib": 19.073486328125, "fp32_gib": 38.14697265625, "int8_gib": 9.5367431640625, "mrl256_int8_gib": 2.384185791015625, "num_vectors": 10000000}}'
assert hashlib.sha256(REFERENCE_JSON.encode()).hexdigest() == 'ee7e2475f2d3a9020be6c5abb112b503e4dfa553ba715b6fff87217fcb4b8109'
REFERENCE = json.loads(REFERENCE_JSON)
LANG = 'en'
DATASETS = REFERENCE["datasets"]
data = DATASETS[LANG]
ids = [row["document_id"] for row in data["documents"]]
texts = [row["text"] for row in data["documents"]]
documents = {row["document_id"]: row for row in data["documents"]}
queries = data["queries"]
stored = {}

1. Chunking: provenance, offsets, and complete coverage

chunk_document records source and chunk IDs, word offsets, and start_char / end_char offsets into the original Python string. These are character positions, not byte offsets. Quoted text is sliced from that string; search normalization is separate. Whitespace-only input has an explicit empty result. Invalid overlap is rejected before processing. Every non-empty document is covered without gaps or a duplicated final window.

Python

def _positive_int(name: str, value: int) -> int:
    """Validate and return one strictly positive integer."""

    if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)

def _positive_finite(name: str, value: float) -> float:
    """Validate and return one strictly positive finite float."""

    value = float(value)
    if not math.isfinite(value) or value <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return value

def tokenize(text: str) -> tuple[str, ...]:
    """Find word spans in the original text, then casefold each word for search."""

    if not isinstance(text, str):
        raise TypeError("text must be a string")
    return tuple(match.group().casefold() for match in WORD_RE.finditer(text))

@dataclass(frozen=True)
class Chunk:
    """One contiguous token window with source and offset metadata."""

    chunk_id: str
    source_id: str
    start_token: int
    end_token: int
    start_char: int
    end_char: int
    text: str
    token_count: int

def chunk_document(
    source_id: str,
    text: str,
    *,
    chunk_size: int,
    overlap: int,
) -> tuple[Chunk, ...]:
    """Split one document into deterministic overlapping token windows.

    Whitespace-only documents return an empty tuple explicitly.  Non-empty documents
    are covered without gaps; the final chunk is not duplicated when the last regular
    window already reaches the end.
    """

    if not isinstance(source_id, str) or not source_id.strip():
        raise ValueError("source_id must be a non-empty string")
    chunk_size = _positive_int("chunk_size", chunk_size)
    if not isinstance(overlap, (int, np.integer)) or isinstance(overlap, bool):
        raise TypeError("overlap must be an integer")
    overlap = int(overlap)
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap must satisfy 0 <= overlap < chunk_size")
    if not isinstance(text, str):
        raise TypeError("text must be a string")
    spans = tuple(WORD_RE.finditer(text))
    if not spans:
        return ()
    stride = chunk_size - overlap
    chunks: list[Chunk] = []
    start = 0
    index = 0
    while start < len(spans):
        end = min(start + chunk_size, len(spans))
        # Character offsets refer to the original Python string, not bytes or
        # casefolded text. Preserve case, punctuation, and spacing inside the span.
        char_start = spans[start].start()
        char_end = spans[end - 1].end()
        chunks.append(
            Chunk(
                chunk_id=f"{source_id}:c{index:03d}",
                source_id=source_id,
                start_token=start,
                end_token=end,
                start_char=char_start,
                end_char=char_end,
                text=text[char_start:char_end],
                token_count=end - start,
            )
        )
        if end == len(spans):
            break
        start += stride
        index += 1
    return tuple(chunks)

def analytical_chunk_count(total_tokens: int, chunk_size: int, overlap: int) -> int:
    """Return the exact sliding-window count for a tokenized document."""

    total_tokens = _positive_int("total_tokens", total_tokens)
    chunk_size = _positive_int("chunk_size", chunk_size)
    if not isinstance(overlap, (int, np.integer)) or isinstance(overlap, bool):
        raise TypeError("overlap must be an integer")
    overlap = int(overlap)
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap must satisfy 0 <= overlap < chunk_size")
    if total_tokens <= chunk_size:
        return 1
    stride = chunk_size - overlap
    return (total_tokens - chunk_size + stride - 1) // stride + 1

def _validate_ids(ids):
    """IDs remain nonempty strings; coercion must not merge 1 and '1'."""
    if any(not isinstance(i, str) or not i.strip() for i in ids):
        raise ValueError("document IDs must be nonempty strings")
    if len(ids) != len(set(ids)):
        raise ValueError("document IDs must be unique")

def _validate_judgments(ranking, relevance):
    """Reject duplicate ranks, fractional/negative grades and empty gold sets."""
    _validate_ids(ranking)
    _validate_ids(list(relevance))
    if any(isinstance(g, bool) or not isinstance(g, (int, np.integer))
           or not 0 <= g <= 30 for g in relevance.values()):
        raise ValueError("grades must be integers from 0 through 30")
    if not any(g > 0 for g in relevance.values()):
        raise ValueError("at least one positive judgment is required")

def stable_ranking(scores: Sequence[float], document_ids: Sequence[str]) -> tuple[str, ...]:
    """Sort scores descending with document ID as a deterministic tie-break."""

    score_array = np.asarray(scores, dtype=np.float64)
    if score_array.ndim != 1 or score_array.size != len(document_ids):
        raise ValueError("scores and document_ids must be aligned rank-1 sequences")
    if not np.all(np.isfinite(score_array)):
        raise ValueError("scores contain NaN or infinity")
    _validate_ids(document_ids)
    if len(set(document_ids)) != len(document_ids):
        raise ValueError("document_ids must be unique")
    return tuple(
        document_id
        for _, document_id in sorted(
            zip(score_array.tolist(), map(str, document_ids)),
            key=lambda pair: (-pair[0], pair[1]),
        )
    )

@dataclass(frozen=True)
class SearchResult:
    """One ranked document and its branch score."""

    document_id: str
    score: float
    rank: int
sample = "one two three four five six seven eight nine ten"
chunks = chunk_document("sample", sample, chunk_size=4, overlap=1)
assert len(chunks) == analytical_chunk_count(10, 4, 1) == 3
assert set(i for c in chunks for i in range(c.start_token,c.end_token)) == set(range(10))
print([(c.start_token,c.end_token,c.text) for c in chunks])
original = "TS-999:  UPDATE\nfailed — İstanbul; Straße."
quotes = chunk_document("case", original, chunk_size=4, overlap=1)
assert all(c.text == original[c.start_char:c.end_char] for c in quotes)
assert quotes[0].text == "TS-999:  UPDATE\nfailed"
assert tokenize(quotes[-1].text) == tokenize(original)[quotes[-1].start_token:quotes[-1].end_token]

2. Sparse branch: exact and inspectable BM25

Document frequencies are computed once at index construction. Empty documents are rejected, and a blank query does not become an arbitrary all-zero ranking. Equal scores are ordered by document_id, making the result reproducible.

Python

class BM25Index:
    """A small exact BM25 index with explicit empty-input policy."""

    def __init__(
        self,
        document_ids: Sequence[str],
        texts: Sequence[str],
        *,
        k1: float = 1.5,
        b: float = 0.75,
    ) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        self.tokens = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in self.tokens):
            raise ValueError("empty documents must be removed before indexing")
        self.k1 = _positive_finite("k1", k1)
        self.b = float(b)
        if not math.isfinite(self.b) or not 0.0 <= self.b <= 1.0:
            raise ValueError("b must lie in [0,1]")
        n_documents = len(self.tokens)
        self.average_length = sum(map(len, self.tokens)) / n_documents
        self.document_frequency = Counter(
            token for document_tokens in self.tokens for token in set(document_tokens)
        )
        self.idf = {
            token: math.log(
                (n_documents - frequency + 0.5) / (frequency + 0.5) + 1.0
            )
            for token, frequency in self.document_frequency.items()
        }

    def scores(self, query: str) -> np.ndarray:
        """Return one BM25 score per indexed document."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        if not any(token in self.idf for token in query_tokens):
            raise ValueError("query has no indexed lexical terms")
        scores = np.zeros(len(self.tokens), dtype=np.float64)
        for row, document_tokens in enumerate(self.tokens):
            frequencies = Counter(document_tokens)
            length_norm = 1.0 - self.b + self.b * len(document_tokens) / self.average_length
            for token in query_tokens:
                tf = frequencies.get(token, 0)
                if tf == 0 or token not in self.idf:
                    continue
                scores[row] += (
                    self.idf[token]
                    * tf
                    * (self.k1 + 1.0)
                    / (tf + self.k1 * length_norm)
                )
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank all documents or the requested top-k with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )
bm25 = BM25Index(ids, texts)
for q in queries:
    stored[q["query_id"]] = {"bm25": bm25.search(q["text"])}
assert [r.document_id for r in stored["q0"]["bm25"]] == REFERENCE["retrieval_results"][LANG]["rankings"]["bm25"]["q0"]
print(stored["q0"]["bm25"][:3])

3. Dense branch: TF-IDF → SVD → exact cosine search

The [N_docs, |V|] matrix is projected into [N_docs, D_lsa] and L2-normalized. A query follows the same path. Blank or fully out-of-vocabulary input raises a declared error rather than dividing by zero. This is an educational dense index, not a claim about neural-embedding quality.

Python

class LSAIndex:
    """Exact cosine search over dense TF-IDF + SVD document vectors."""

    def __init__(self, document_ids: Sequence[str], texts: Sequence[str], *, dimension: int = 4) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        tokenized = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in tokenized):
            raise ValueError("empty documents must be removed before indexing")
        vocabulary = sorted({token for tokens in tokenized for token in tokens})
        self.vocabulary = {token: index for index, token in enumerate(vocabulary)}
        counts = np.zeros((len(tokenized), len(vocabulary)), dtype=np.float64)
        document_frequency = np.zeros(len(vocabulary), dtype=np.float64)
        for row, tokens in enumerate(tokenized):
            frequencies = Counter(tokens)
            for token, count in frequencies.items():
                counts[row, self.vocabulary[token]] = count
                document_frequency[self.vocabulary[token]] += 1.0
        self.idf = np.log((1.0 + len(tokenized)) / (1.0 + document_frequency)) + 1.0
        tf = np.where(counts > 0.0, 1.0 + np.log(np.maximum(counts, 1.0)), 0.0)
        tfidf = tf * self.idf
        _, singular_values, vt = np.linalg.svd(tfidf, full_matrices=False)
        numerical_rank = int(np.sum(singular_values > 1e-10))
        dimension = min(_positive_int("dimension", dimension), numerical_rank)
        self.projection = vt[:dimension].T
        dense = tfidf @ self.projection
        norms = np.linalg.norm(dense, axis=1, keepdims=True)
        if np.any(norms <= 0.0):
            raise ValueError("at least one document has a zero dense representation")
        self.document_vectors = dense / norms

    def encode_query(self, query: str) -> np.ndarray:
        """Map one non-empty in-vocabulary query to a unit dense vector."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        counts = Counter(query_tokens)
        sparse = np.zeros(len(self.vocabulary), dtype=np.float64)
        for token, count in counts.items():
            position = self.vocabulary.get(token)
            if position is not None:
                sparse[position] = (1.0 + math.log(count)) * self.idf[position]
        dense = sparse @ self.projection
        norm = float(np.linalg.norm(dense))
        if norm <= 0.0:
            raise ValueError("query contains no terms represented by the dense index")
        return dense / norm

    def scores(self, query: str) -> np.ndarray:
        """Return exact cosine scores against every document vector."""

        scores = self.document_vectors @ self.encode_query(query)
        if not np.all(np.isfinite(scores)):
            raise FloatingPointError("dense search produced NaN or infinity")
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank exact dense cosine scores with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )
dense = LSAIndex(ids, texts, dimension=4)
assert dense.document_vectors.shape == (13,4)
assert np.allclose(np.linalg.norm(dense.document_vectors,axis=1),1)
for q in queries:
    stored[q["query_id"]]["dense_lsa"] = dense.search(q["text"])
print(dense.document_vectors.shape)

4. RRF: accumulating contributions and sorting candidates

The standard formula uses one-based ranks:

\operatorname{RRF}(d)=\sum_b \frac{w_b}{k+r_b(d)}.

The implementation scans every list once, rejects duplicate IDs inside a branch, and retains source ranks plus per-branch contributions. Accumulation is linear in total list length, followed by $O(U\log U)$ sorting of $U$ unique candidates. Hybrid retrieval is measured rather than assumed to dominate its components.

Python

@dataclass(frozen=True)
class FusedResult:
    """One RRF result with a per-branch contribution trace."""

    document_id: str
    score: float
    rank: int
    contributions: Mapping[str, float]
    source_ranks: Mapping[str, int]

def reciprocal_rank_fusion(
    rankings: Mapping[str, Sequence[str]],
    *,
    k: float = 60.0,
    weights: Mapping[str, float] | None = None,
    depth: int | None = None,
) -> tuple[FusedResult, ...]:
    """Accumulate ranks in linear time, then sort unique candidates.

    Ranks are one-based. Duplicate document IDs inside one branch are rejected because
    they would double-count evidence.  The implementation builds rank maps once and
    never calls ``list.index`` inside a loop.
    """

    k = _positive_finite("k", k)
    if not rankings:
        raise ValueError("at least one ranking is required")
    if depth is not None:
        depth = _positive_int("depth", depth)
    branch_weights = {name: 1.0 for name in rankings}
    if weights is not None:
        if set(weights) != set(rankings):
            raise ValueError("weights must contain exactly the ranking branch names")
        branch_weights = {
            name: _positive_finite(f"weight[{name}]", weights[name]) for name in rankings
        }
    scores: dict[str, float] = {}
    contributions: dict[str, dict[str, float]] = {}
    source_ranks: dict[str, dict[str, int]] = {}
    for branch_name, ranking in rankings.items():
        _validate_ids(ranking)
        sequence = tuple(ranking)
        if len(sequence) != len(set(sequence)):
            raise ValueError(f"ranking {branch_name!r} contains duplicate document IDs")
        if depth is not None:
            sequence = sequence[:depth]
        for rank, document_id in enumerate(sequence, start=1):
            contribution = branch_weights[branch_name] / (k + rank)
            scores[document_id] = scores.get(document_id, 0.0) + contribution
            contributions.setdefault(document_id, {})[branch_name] = contribution
            source_ranks.setdefault(document_id, {})[branch_name] = rank
    ordered = sorted(scores, key=lambda document_id: (-scores[document_id], document_id))
    return tuple(
        FusedResult(
            document_id=document_id,
            score=float(scores[document_id]),
            rank=rank,
            contributions=dict(sorted(contributions[document_id].items())),
            source_ranks=dict(sorted(source_ranks[document_id].items())),
        )
        for rank, document_id in enumerate(ordered, start=1)
    )
for q in queries:
    row = stored[q["query_id"]]
    row["rrf"] = reciprocal_rank_fusion(
        {name:[r.document_id for r in row[name]] for name in ("bm25","dense_lsa")},
        depth=8)
print([asdict(r) for r in stored["q0"]["rrf"][:3]])

5. Interaction reranking over a bounded candidate set

To keep the mechanism visible, the reranker uses a MaxSim-style teaching score: each query token takes its best character-trigram similarity to document tokens, and the values are averaged. This is a transparent late-interaction proxy, not a cross-encoder. The production invariant is the same: expensive pairwise scoring is applied only after candidate generation.

Token similarity is character-trigram Jaccard. This untrained lexical heuristic averages each query word's best match; it is neither ColBERT nor a cross-encoder. Truncating the candidate set to eight limits attainable recall.

Python

def _character_trigrams(token: str) -> frozenset[str]:
    """Return boundary-aware character trigrams for one token."""

    wrapped = f"^{token}$"
    return frozenset(wrapped[index : index + 3] for index in range(max(0, len(wrapped) - 2)))

def _token_similarity(left: str, right: str) -> float:
    """Return Jaccard similarity between two token trigram sets."""

    a = _character_trigrams(left)
    b = _character_trigrams(right)
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)

def token_interaction_score(query: str, document: str) -> float:
    """Score a query/document pair with a transparent MaxSim-style proxy.

    This is not a neural cross-encoder.  It deliberately exposes the per-query-token
    maximum interaction used by late-interaction retrieval families.
    """

    query_tokens = tokenize(query)
    document_tokens = tokenize(document)
    if not query_tokens:
        raise ValueError("query must contain at least one word token")
    if not document_tokens:
        raise ValueError("document must contain at least one word token")
    return float(
        sum(
            max(_token_similarity(query_token, document_token) for document_token in document_tokens)
            for query_token in query_tokens
        )
        / len(query_tokens)
    )

def rerank_candidates(
    query: str,
    candidate_ids: Sequence[str],
    documents: Mapping[str, str],
) -> tuple[SearchResult, ...]:
    """Rerank a candidate list with the token-interaction proxy."""

    _validate_ids(candidate_ids)
    candidate_ids = tuple(candidate_ids)
    if not candidate_ids:
        raise ValueError("candidate_ids must be non-empty")
    if len(candidate_ids) != len(set(candidate_ids)):
        raise ValueError("candidate_ids must be unique")
    missing = [document_id for document_id in candidate_ids if document_id not in documents]
    if missing:
        raise KeyError(f"missing documents: {missing}")
    scores = [token_interaction_score(query, documents[document_id]) for document_id in candidate_ids]
    ranking = stable_ranking(scores, candidate_ids)
    score_by_id = dict(zip(candidate_ids, scores))
    return tuple(
        SearchResult(document_id, float(score_by_id[document_id]), rank)
        for rank, document_id in enumerate(ranking, start=1)
    )
for q in queries:
    row = stored[q["query_id"]]
    row["reranked"] = rerank_candidates(q["text"], [r.document_id for r in row["rrf"][:8]],
                                       dict(zip(ids,texts)))
print(stored["q0"]["reranked"][:3])

6. Retrieval metrics answer different questions

MRR uses the first relevant hit;
Recall@5 measures how much relevant evidence survives the top five;
nDCG@5 incorporates graded relevance and position discount.

Qrels without a positive grade and rankings with duplicate IDs are errors. On the teaching corpus, RRF need not dominate every branch; reranking can improve only the candidates that reached it.

Python

def reciprocal_rank(ranking: Sequence[str], relevance: Mapping[str, int]) -> float:
    """Return reciprocal rank of the first positively relevant document."""
    _validate_judgments(ranking, relevance)

    if not any(int(grade) > 0 for grade in relevance.values()):
        raise ValueError("relevance must contain at least one positive grade")
    for rank, document_id in enumerate(ranking, start=1):
        if int(relevance.get(document_id, 0)) > 0:
            return 1.0 / rank
    return 0.0

def recall_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return set recall among all positively relevant documents."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    relevant = {document_id for document_id, grade in relevance.items() if int(grade) > 0}
    if not relevant:
        raise ValueError("relevance must contain at least one positive grade")
    return len(set(ranking[:k]) & relevant) / len(relevant)

def ndcg_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return nDCG@k with gains 2^grade-1 and log2 discount."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    grades = [int(grade) for grade in relevance.values()]
    if any(grade < 0 for grade in grades):
        raise ValueError("relevance grades must be non-negative")
    if not any(grade > 0 for grade in grades):
        raise ValueError("relevance must contain at least one positive grade")
    dcg = sum(
        (2 ** int(relevance.get(document_id, 0)) - 1.0) / math.log2(rank + 1.0)
        for rank, document_id in enumerate(ranking[:k], start=1)
    )
    ideal = sorted(grades, reverse=True)[:k]
    idcg = sum((2**grade - 1.0) / math.log2(rank + 1.0) for rank, grade in enumerate(ideal, start=1))
    return dcg / idcg

def evaluate_rankings(
    rankings: Mapping[str, Sequence[str]],
    query_records: Sequence[Mapping[str, object]],
    *,
    k: int = 5,
) -> dict[str, float]:
    """Aggregate MRR, Recall@k, and nDCG@k over one frozen query set."""

    k = _positive_int("k", k)
    if not query_records or len({r["query_id"] for r in query_records}) != len(query_records):
        raise ValueError("queries must be nonempty with unique IDs")
    if len(rankings) != len(query_records):
        raise ValueError("rankings must contain exactly one entry per query")
    reciprocal_ranks: list[float] = []
    recalls: list[float] = []
    ndcgs: list[float] = []
    for record in query_records:
        query_id = str(record["query_id"])
        if query_id not in rankings:
            raise KeyError(f"missing ranking for {query_id}")
        relevance = dict(record["relevance"])
        ranking = tuple(rankings[query_id])
        if len(ranking) != len(set(ranking)):
            raise ValueError(f"ranking for {query_id} contains duplicates")
        reciprocal_ranks.append(reciprocal_rank(ranking, relevance))
        recalls.append(recall_at_k(ranking, relevance, k))
        ndcgs.append(ndcg_at_k(ranking, relevance, k))
    return {
        "mrr": float(np.mean(reciprocal_ranks)),
        f"recall_at_{k}": float(np.mean(recalls)),
        f"ndcg_at_{k}": float(np.mean(ndcgs)),
    }
# No search or scoring calls inside evaluation: consume the stored records.
metrics = {}
for method in ("bm25","dense_lsa","rrf","reranked"):
    rankings = {q["query_id"]:[r.document_id for r in stored[q["query_id"]][method]]
                for q in queries}
    metrics[method] = evaluate_rankings(rankings,queries,k=5)
    assert all(np.isclose(value, REFERENCE["retrieval_results"][LANG]["metrics"][method][name])
               for name,value in metrics[method].items())
print(json.dumps(metrics,indent=2))
assert len(stored) == 7 and sum(len(row["reranked"]) for row in stored.values()) == 56

7. Context assembly: budget, deduplication, source limits, and citations

pack_context accepts ranked IDs and score maps, deduplicates by normalized search words, limits chunks per source, and admits only whole chunks. Every omission is recorded in skipped; there is no silent truncation. The rendered output contains stable citations of the form [source:document].

Each fragment is charged after its citation is added, and the final serialized string is counted again. For an actual model, use its tokenizer and reserve instructions, the query, chat-template tokens, and output capacity.

Python

@dataclass(frozen=True)
class ContextItem:
    """One selected chunk and its explicit budget contribution."""

    citation: str
    document_id: str
    source_id: str
    token_count: int
    score: float
    text: str

@dataclass(frozen=True)
class ContextPack:
    """A budgeted retrieval context with selected and skipped trace data."""

    items: tuple[ContextItem, ...]
    used_tokens: int
    max_tokens: int
    skipped: Mapping[str, str]
    rendered_text: str

def pack_context(
    ranked_ids: Sequence[str],
    scores: Mapping[str, float],
    documents: Mapping[str, Mapping[str, object]],
    *,
    max_tokens: int,
    max_per_source: int = 1,
    citation_overhead_tokens: int | None = None,
) -> ContextPack:
    """Deduplicate and pack whole chunks under an explicit token budget.

    Exact normalized-text duplicates and per-source overflow are skipped.  Chunks are
    never silently truncated; a chunk either fits in full or is traced as skipped.
    """

    max_tokens = _positive_int("max_tokens", max_tokens)
    max_per_source = _positive_int("max_per_source", max_per_source)
    if citation_overhead_tokens is not None:
        raise ValueError("count the serialized citation instead of a guessed overhead")
    _validate_ids(ranked_ids)
    ranked_ids = tuple(ranked_ids)
    if not ranked_ids:
        raise ValueError("ranked_ids must be nonempty")
    selected: list[ContextItem] = []
    skipped: dict[str, str] = {}
    seen_texts: set[str] = set()
    per_source: Counter[str] = Counter()
    used = 0
    for document_id in ranked_ids:
        if document_id not in documents or document_id not in scores:
            raise KeyError(f"missing document or score for {document_id}")
        record = documents[document_id]
        text = record["text"]
        source_id = record["source_id"]
        _validate_ids([source_id])
        if not math.isfinite(float(scores[document_id])):
            raise ValueError("context scores must be finite")
        tokens = tokenize(text)
        if not tokens:
            skipped[document_id] = "empty"
            continue
        normalized = " ".join(tokens)
        if normalized in seen_texts:
            skipped[document_id] = "duplicate_text"
            continue
        if per_source[source_id] >= max_per_source:
            skipped[document_id] = "source_limit"
            continue
        citation = f"[{source_id}:{document_id}]"
        cost = len(tokenize(citation + " " + text))
        if used + cost > max_tokens:
            skipped[document_id] = "budget"
            continue
        citation = f"[{source_id}:{document_id}]"
        selected.append(
            ContextItem(
                citation=citation,
                document_id=document_id,
                source_id=source_id,
                token_count=cost,
                score=float(scores[document_id]),
                text=text,
            )
        )
        used += cost
        seen_texts.add(normalized)
        per_source[source_id] += 1
    if not selected:
        raise ValueError("no chunk fits the context budget")
    rendered = "\n\n".join(f"{item.citation} {item.text}" for item in selected)
    assert len(tokenize(rendered)) == used
    return ContextPack(tuple(selected), used, max_tokens, dict(skipped), rendered)
head = stored["q6"]["reranked"]
pack = pack_context([r.document_id for r in head],{r.document_id:r.score for r in head},
                    documents,max_tokens=42)
assert len(tokenize(pack.rendered_text)) == pack.used_tokens <= 42
expected = REFERENCE["retrieval_results"][LANG]["context_pack_q6"]
assert [r.document_id for r in pack.items] == expected["selected_ids"]
print(pack.rendered_text)
print(pack.used_tokens, pack.skipped)

End-to-end request: source text to cited excerpt

This time the chunks themselves enter the index. The thirteen source documents become twelve-word windows with three-word overlap. The answer is an exact excerpt from a selected window: this tests provenance, not an LLM's synthesis ability. A valid citation does not establish that the excerpt answers the question.

Python

all_chunks = [c for doc in data["documents"]
              for c in chunk_document(doc["document_id"],doc["text"],chunk_size=12,overlap=3)]
chunk_docs = {c.chunk_id:asdict(c) for c in all_chunks}
chunk_ids = list(chunk_docs)
chunk_texts = [chunk_docs[i]["text"] for i in chunk_ids]
lex = BM25Index(chunk_ids,chunk_texts)
vec = LSAIndex(chunk_ids,chunk_texts,dimension=4)
query = queries[6]["text"]
# Two search calls, one bounded rerank, one pack; metrics never repeat them.
branches = {"lex":[r.document_id for r in lex.search(query,top_k=8)],
            "lsa":[r.document_id for r in vec.search(query,top_k=8)]}
fused = reciprocal_rank_fusion(branches)
ranked = rerank_candidates(query,[r.document_id for r in fused[:8]],
                           dict(zip(chunk_ids,chunk_texts)))
context = pack_context([r.document_id for r in ranked],{r.document_id:r.score for r in ranked},
                       chunk_docs,max_tokens=42)
item = context.items[0]
answer = {"text":item.text, "citation":item.citation, "chunk_id":item.document_id}
record = chunk_docs[answer["chunk_id"]]
source = documents[record["source_id"]]["text"]
assert answer["text"] == source[record["start_char"]:record["end_char"]]
assert tokenize(answer["text"]) == tokenize(source)[record["start_token"]:record["end_token"]]
trace = {"search_calls":2,"rerank_calls":1,"scored_pairs":len(ranked),
         "answer":answer,"used_word_units":context.used_tokens}
assert json.loads(json.dumps(trace)) == trace
print(json.dumps(trace,ensure_ascii=False,indent=2))

8. Storage, ANN, and RAG versus long context

Division by 2**20 and 2**30 yields MiB/GiB. Raw vector payload excludes HNSW edges, index metadata, and allocator overhead. ANN and cascade times are transparent arithmetic scenarios at stated throughput, not engine benchmarks. KV uses 128 KiB per token; the 200K-token document and 3712-token RAG context differ by 53.88× in length, linearized prefill, and KV payload.

Python

def index_storage_gib(num_vectors: int, dimension: int, bits_per_component: float) -> float:
    """Return raw vector payload in GiB using powers-of-two units."""

    num_vectors = _positive_int("num_vectors", num_vectors)
    dimension = _positive_int("dimension", dimension)
    bits_per_component = _positive_finite("bits_per_component", bits_per_component)
    return num_vectors * dimension * bits_per_component / 8.0 / 2**30

def brute_ann_scenario(
    *,
    num_vectors: int,
    dimension: int,
    cpu_flops_per_second: float,
    visited_vectors: int,
) -> dict[str, float]:
    """Return a transparent exact-vs-ANN arithmetic scenario."""

    num_vectors = _positive_int("num_vectors", num_vectors)
    dimension = _positive_int("dimension", dimension)
    visited_vectors = _positive_int("visited_vectors", visited_vectors)
    cpu_flops_per_second = _positive_finite("cpu_flops_per_second", cpu_flops_per_second)
    exact_flops = 2.0 * num_vectors * dimension
    ann_flops = 2.0 * visited_vectors * dimension
    exact_seconds = exact_flops / cpu_flops_per_second
    ann_milliseconds = ann_flops / cpu_flops_per_second * 1e3
    return {
        "exact_flops": exact_flops,
        "exact_seconds": exact_seconds,
        "ann_flops": ann_flops,
        "ann_milliseconds": ann_milliseconds,
        "speedup": exact_seconds * 1e3 / ann_milliseconds,
    }

def cascade_latency_scenario() -> dict[str, float]:
    """Return a scenario mixing 2PT compute estimates and an assumed decode rate."""

    peak = 990e12
    mfu = 0.5
    effective = peak * mfu
    encoder_parameters = 0.6e9
    query_tokens = 32
    candidates = 100
    rerank_tokens = 432
    generator_parameters = 8.03e9
    context_tokens = 512 + 8 * 400
    answer_tokens = 300
    decode_tokens_per_second = 209.4
    stages = {
        "embed_ms": 2 * encoder_parameters * query_tokens / effective * 1e3,
        "ann_ms": brute_ann_scenario(
            num_vectors=10_000_000,
            dimension=1024,
            cpu_flops_per_second=10e9,
            visited_vectors=2_000,
        )["ann_milliseconds"],
        "rerank_ms": candidates * 2 * encoder_parameters * rerank_tokens / effective * 1e3,
        "prefill_ms": 2 * generator_parameters * context_tokens / effective * 1e3,
        "decode_ms": answer_tokens / decode_tokens_per_second * 1e3,
    }
    total = sum(stages.values())
    return {
        **stages,
        "total_ms": total,
        "context_tokens": context_tokens,
        "decode_share": stages["decode_ms"] / total,
        "retrieval_before_prefill_share": (
            stages["embed_ms"] + stages["ann_ms"] + stages["rerank_ms"]
        )
        / total,
    }

def long_context_vs_rag(*, long_tokens: int = 200_000, rag_tokens: int = 3_712) -> dict[str, float]:
    """Compare prefill, KV payload, and token count under one explicit model geometry."""

    long_tokens = _positive_int("long_tokens", long_tokens)
    rag_tokens = _positive_int("rag_tokens", rag_tokens)
    if rag_tokens >= long_tokens:
        raise ValueError("rag_tokens must be smaller than long_tokens")
    generator_parameters = 8.03e9
    effective_flops = 990e12 * 0.5
    kv_bytes_per_token = 128 * 2**10
    long_prefill_seconds = 2 * generator_parameters * long_tokens / effective_flops
    rag_prefill_seconds = 2 * generator_parameters * rag_tokens / effective_flops
    ratio = long_tokens / rag_tokens
    return {
        "long_tokens": long_tokens,
        "rag_tokens": rag_tokens,
        "long_prefill_seconds": long_prefill_seconds,
        "rag_prefill_seconds": rag_prefill_seconds,
        "long_kv_gib": long_tokens * kv_bytes_per_token / 2**30,
        "rag_kv_gib": rag_tokens * kv_bytes_per_token / 2**30,
        "ratio": ratio,
        "kv_bytes_per_token": kv_bytes_per_token,
    }
storage = index_storage_gib(10_000_000,1024,1)
assert storage == 1_280_000_000 / 2**30
assert analytical_chunk_count(20_000_000,400,64) == 59_524
cascade = cascade_latency_scenario()
assert np.isclose(cascade["total_ms"],1658.3129827096177)
assert np.isclose(long_context_vs_rag()["ratio"],200_000/3712)
print(json.dumps(cascade,indent=2))

9. Negative tests

Expected exceptions count as successful checks: blank queries or documents, invalid overlap, an out-of-vocabulary dense query, duplicate IDs, malformed qrels, an impossible context budget, and a comparison in which RAG does not reduce the input.

Python

def rejects(fn):
    try:
        fn()
    except (ValueError,TypeError,KeyError):
        return
    raise AssertionError("invalid input was accepted")
rejects(lambda: bm25.search("zzzzunseen"))
rejects(lambda: dense.search("zzzzunseen"))
rejects(lambda: chunk_document("x","text",chunk_size=4,overlap=4))
rejects(lambda: reciprocal_rank_fusion({"x":["a","a"]}))
rejects(lambda: ndcg_at_k(["a","a"],{"a":3},2))
rejects(lambda: ndcg_at_k(["a"],{"a":1.9},1))
rejects(lambda: evaluate_rankings({},[],k=5))
rejects(lambda: BM25Index([1,"1"],["a","b"]))
rejects(lambda: rerank_candidates("a",[1],{"1":"a"}))
rejects(lambda: pack_context([ids[0]],{ids[0]:float("nan")},documents,max_tokens=50))
rejects(lambda: pack_context([ids[0]],{ids[0]:1},documents,max_tokens=1))
assert reciprocal_rank([],{"a":1}) == 0
assert recall_at_k(["a"],{"a":1,"b":1},1) == 0.5
assert ndcg_at_k(["a","b"],{"a":3,"b":1},2) == 1
assert analytical_chunk_count(10**20+1,4,1) == (10**20+1-4+2)//3+1
print("PASS")

Exercises

[auto] Add min_chunk_tokens and a trace for a discarded short tail.
[auto] Use the existing weighted RRF implementation to find a query whose top result changes. Explain the change from its stored branch contributions.
[auto] Add precision_at_k and average_precision; contrast AP with MRR on a query containing two relevant documents.
[auto] Permit partial packing only with explicit allow_truncation=True and retain the original token span.
[project] Replace LSA with an open bi-encoder while preserving the same interface and tests; measure embedding latency and Recall@10 separately.
[project] Add a real ANN index and plot Recall@10 against latency relative to exact cosine search.

The mandatory route runs on CPU and needs no repository files. It is a complete educational retrieval/RAG cascade, not a deployment-ready search service.

Module 17 project-exercise demo — Candidates, reranking work, and late chunking

Project demonstration

Download .ipynb module_17_project_demo.ipynb · 29.4 KB

Show notebook text

Module 17 project-exercise demo — Candidates, reranking work, and late chunking

Implementation demo, not a completed full-scale project. Demo 1 runs BM25, LSA, RRF, and reranking over four fictional records containing confusable error codes. Scores are stored; work is counted as actual scored pairs. Demo 2 computes attention before and after chunking on designed vectors to expose how neighboring context reaches a pooled chunk. All inputs are small, deterministic, offline, and CPU-friendly. The goal is to validate interfaces, metrics, invariants, and report structure before spending GPU or API budget.

Python

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

Only reranking depth changes: one, two, or four candidates. Locate error999 in each branch, then check whether each candidate set makes it available to the reranker. Scored pairs measure work, not seconds; pairs may be batched. All three settings reuse one diagnostic query, so this experiment says nothing about held-out performance.

Python

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437

import math, re, json
from collections import Counter
from dataclasses import dataclass
from typing import Mapping, Sequence
WORD_RE = re.compile(r"[^\W_]+", re.UNICODE)
def _positive_int(name: str, value: int) -> int:
    """Validate and return one strictly positive integer."""

    if not isinstance(value, (int, np.integer)) or isinstance(value, bool) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)

def _positive_finite(name: str, value: float) -> float:
    """Validate and return one strictly positive finite float."""

    value = float(value)
    if not math.isfinite(value) or value <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return value

def tokenize(text: str) -> tuple[str, ...]:
    """Find word spans in the original text, then casefold each word for search."""

    if not isinstance(text, str):
        raise TypeError("text must be a string")
    return tuple(match.group().casefold() for match in WORD_RE.finditer(text))

def _validate_ids(ids):
    """IDs remain nonempty strings; coercion must not merge 1 and '1'."""
    if any(not isinstance(i, str) or not i.strip() for i in ids):
        raise ValueError("document IDs must be nonempty strings")
    if len(ids) != len(set(ids)):
        raise ValueError("document IDs must be unique")

def _validate_judgments(ranking, relevance):
    """Reject duplicate ranks, fractional/negative grades and empty gold sets."""
    _validate_ids(ranking)
    _validate_ids(list(relevance))
    if any(isinstance(g, bool) or not isinstance(g, (int, np.integer))
           or not 0 <= g <= 30 for g in relevance.values()):
        raise ValueError("grades must be integers from 0 through 30")
    if not any(g > 0 for g in relevance.values()):
        raise ValueError("at least one positive judgment is required")

def stable_ranking(scores: Sequence[float], document_ids: Sequence[str]) -> tuple[str, ...]:
    """Sort scores descending with document ID as a deterministic tie-break."""

    score_array = np.asarray(scores, dtype=np.float64)
    if score_array.ndim != 1 or score_array.size != len(document_ids):
        raise ValueError("scores and document_ids must be aligned rank-1 sequences")
    if not np.all(np.isfinite(score_array)):
        raise ValueError("scores contain NaN or infinity")
    _validate_ids(document_ids)
    if len(set(document_ids)) != len(document_ids):
        raise ValueError("document_ids must be unique")
    return tuple(
        document_id
        for _, document_id in sorted(
            zip(score_array.tolist(), map(str, document_ids)),
            key=lambda pair: (-pair[0], pair[1]),
        )
    )

@dataclass(frozen=True)
class SearchResult:
    """One ranked document and its branch score."""

    document_id: str
    score: float
    rank: int

class BM25Index:
    """A small exact BM25 index with explicit empty-input policy."""

    def __init__(
        self,
        document_ids: Sequence[str],
        texts: Sequence[str],
        *,
        k1: float = 1.5,
        b: float = 0.75,
    ) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        self.tokens = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in self.tokens):
            raise ValueError("empty documents must be removed before indexing")
        self.k1 = _positive_finite("k1", k1)
        self.b = float(b)
        if not math.isfinite(self.b) or not 0.0 <= self.b <= 1.0:
            raise ValueError("b must lie in [0,1]")
        n_documents = len(self.tokens)
        self.average_length = sum(map(len, self.tokens)) / n_documents
        self.document_frequency = Counter(
            token for document_tokens in self.tokens for token in set(document_tokens)
        )
        self.idf = {
            token: math.log(
                (n_documents - frequency + 0.5) / (frequency + 0.5) + 1.0
            )
            for token, frequency in self.document_frequency.items()
        }

    def scores(self, query: str) -> np.ndarray:
        """Return one BM25 score per indexed document."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        if not any(token in self.idf for token in query_tokens):
            raise ValueError("query has no indexed lexical terms")
        scores = np.zeros(len(self.tokens), dtype=np.float64)
        for row, document_tokens in enumerate(self.tokens):
            frequencies = Counter(document_tokens)
            length_norm = 1.0 - self.b + self.b * len(document_tokens) / self.average_length
            for token in query_tokens:
                tf = frequencies.get(token, 0)
                if tf == 0 or token not in self.idf:
                    continue
                scores[row] += (
                    self.idf[token]
                    * tf
                    * (self.k1 + 1.0)
                    / (tf + self.k1 * length_norm)
                )
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank all documents or the requested top-k with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )

class LSAIndex:
    """Exact cosine search over dense TF-IDF + SVD document vectors."""

    def __init__(self, document_ids: Sequence[str], texts: Sequence[str], *, dimension: int = 4) -> None:
        if len(document_ids) != len(texts) or not document_ids:
            raise ValueError("document_ids and texts must be equally sized and non-empty")
        _validate_ids(document_ids)
        self.document_ids = tuple(map(str, document_ids))
        tokenized = tuple(tokenize(text) for text in texts)
        if any(not tokens for tokens in tokenized):
            raise ValueError("empty documents must be removed before indexing")
        vocabulary = sorted({token for tokens in tokenized for token in tokens})
        self.vocabulary = {token: index for index, token in enumerate(vocabulary)}
        counts = np.zeros((len(tokenized), len(vocabulary)), dtype=np.float64)
        document_frequency = np.zeros(len(vocabulary), dtype=np.float64)
        for row, tokens in enumerate(tokenized):
            frequencies = Counter(tokens)
            for token, count in frequencies.items():
                counts[row, self.vocabulary[token]] = count
                document_frequency[self.vocabulary[token]] += 1.0
        self.idf = np.log((1.0 + len(tokenized)) / (1.0 + document_frequency)) + 1.0
        tf = np.where(counts > 0.0, 1.0 + np.log(np.maximum(counts, 1.0)), 0.0)
        tfidf = tf * self.idf
        _, singular_values, vt = np.linalg.svd(tfidf, full_matrices=False)
        numerical_rank = int(np.sum(singular_values > 1e-10))
        dimension = min(_positive_int("dimension", dimension), numerical_rank)
        self.projection = vt[:dimension].T
        dense = tfidf @ self.projection
        norms = np.linalg.norm(dense, axis=1, keepdims=True)
        if np.any(norms <= 0.0):
            raise ValueError("at least one document has a zero dense representation")
        self.document_vectors = dense / norms

    def encode_query(self, query: str) -> np.ndarray:
        """Map one non-empty in-vocabulary query to a unit dense vector."""

        query_tokens = tokenize(query)
        if not query_tokens:
            raise ValueError("query must contain at least one word token")
        counts = Counter(query_tokens)
        sparse = np.zeros(len(self.vocabulary), dtype=np.float64)
        for token, count in counts.items():
            position = self.vocabulary.get(token)
            if position is not None:
                sparse[position] = (1.0 + math.log(count)) * self.idf[position]
        dense = sparse @ self.projection
        norm = float(np.linalg.norm(dense))
        if norm <= 0.0:
            raise ValueError("query contains no terms represented by the dense index")
        return dense / norm

    def scores(self, query: str) -> np.ndarray:
        """Return exact cosine scores against every document vector."""

        scores = self.document_vectors @ self.encode_query(query)
        if not np.all(np.isfinite(scores)):
            raise FloatingPointError("dense search produced NaN or infinity")
        return scores

    def search(self, query: str, *, top_k: int | None = None) -> tuple[SearchResult, ...]:
        """Rank exact dense cosine scores with stable ties."""

        scores = self.scores(query)
        ranking = stable_ranking(scores, self.document_ids)
        if top_k is not None:
            top_k = _positive_int("top_k", top_k)
            ranking = ranking[:top_k]
        score_by_id = dict(zip(self.document_ids, scores.tolist()))
        return tuple(
            SearchResult(document_id, float(score_by_id[document_id]), rank)
            for rank, document_id in enumerate(ranking, start=1)
        )

@dataclass(frozen=True)
class FusedResult:
    """One RRF result with a per-branch contribution trace."""

    document_id: str
    score: float
    rank: int
    contributions: Mapping[str, float]
    source_ranks: Mapping[str, int]

def reciprocal_rank_fusion(
    rankings: Mapping[str, Sequence[str]],
    *,
    k: float = 60.0,
    weights: Mapping[str, float] | None = None,
    depth: int | None = None,
) -> tuple[FusedResult, ...]:
    """Accumulate ranks in linear time, then sort unique candidates.

    Ranks are one-based. Duplicate document IDs inside one branch are rejected because
    they would double-count evidence.  The implementation builds rank maps once and
    never calls ``list.index`` inside a loop.
    """

    k = _positive_finite("k", k)
    if not rankings:
        raise ValueError("at least one ranking is required")
    if depth is not None:
        depth = _positive_int("depth", depth)
    branch_weights = {name: 1.0 for name in rankings}
    if weights is not None:
        if set(weights) != set(rankings):
            raise ValueError("weights must contain exactly the ranking branch names")
        branch_weights = {
            name: _positive_finite(f"weight[{name}]", weights[name]) for name in rankings
        }
    scores: dict[str, float] = {}
    contributions: dict[str, dict[str, float]] = {}
    source_ranks: dict[str, dict[str, int]] = {}
    for branch_name, ranking in rankings.items():
        _validate_ids(ranking)
        sequence = tuple(ranking)
        if len(sequence) != len(set(sequence)):
            raise ValueError(f"ranking {branch_name!r} contains duplicate document IDs")
        if depth is not None:
            sequence = sequence[:depth]
        for rank, document_id in enumerate(sequence, start=1):
            contribution = branch_weights[branch_name] / (k + rank)
            scores[document_id] = scores.get(document_id, 0.0) + contribution
            contributions.setdefault(document_id, {})[branch_name] = contribution
            source_ranks.setdefault(document_id, {})[branch_name] = rank
    ordered = sorted(scores, key=lambda document_id: (-scores[document_id], document_id))
    return tuple(
        FusedResult(
            document_id=document_id,
            score=float(scores[document_id]),
            rank=rank,
            contributions=dict(sorted(contributions[document_id].items())),
            source_ranks=dict(sorted(source_ranks[document_id].items())),
        )
        for rank, document_id in enumerate(ordered, start=1)
    )

def _character_trigrams(token: str) -> frozenset[str]:
    """Return boundary-aware character trigrams for one token."""

    wrapped = f"^{token}$"
    return frozenset(wrapped[index : index + 3] for index in range(max(0, len(wrapped) - 2)))

def _token_similarity(left: str, right: str) -> float:
    """Return Jaccard similarity between two token trigram sets."""

    a = _character_trigrams(left)
    b = _character_trigrams(right)
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)

def token_interaction_score(query: str, document: str) -> float:
    """Score a query/document pair with a transparent MaxSim-style proxy.

    This is not a neural cross-encoder.  It deliberately exposes the per-query-token
    maximum interaction used by late-interaction retrieval families.
    """

    query_tokens = tokenize(query)
    document_tokens = tokenize(document)
    if not query_tokens:
        raise ValueError("query must contain at least one word token")
    if not document_tokens:
        raise ValueError("document must contain at least one word token")
    return float(
        sum(
            max(_token_similarity(query_token, document_token) for document_token in document_tokens)
            for query_token in query_tokens
        )
        / len(query_tokens)
    )

def rerank_candidates(
    query: str,
    candidate_ids: Sequence[str],
    documents: Mapping[str, str],
) -> tuple[SearchResult, ...]:
    """Rerank a candidate list with the token-interaction proxy."""

    _validate_ids(candidate_ids)
    candidate_ids = tuple(candidate_ids)
    if not candidate_ids:
        raise ValueError("candidate_ids must be non-empty")
    if len(candidate_ids) != len(set(candidate_ids)):
        raise ValueError("candidate_ids must be unique")
    missing = [document_id for document_id in candidate_ids if document_id not in documents]
    if missing:
        raise KeyError(f"missing documents: {missing}")
    scores = [token_interaction_score(query, documents[document_id]) for document_id in candidate_ids]
    ranking = stable_ranking(scores, candidate_ids)
    score_by_id = dict(zip(candidate_ids, scores))
    return tuple(
        SearchResult(document_id, float(score_by_id[document_id]), rank)
        for rank, document_id in enumerate(ranking, start=1)
    )

def reciprocal_rank(ranking: Sequence[str], relevance: Mapping[str, int]) -> float:
    """Return reciprocal rank of the first positively relevant document."""
    _validate_judgments(ranking, relevance)

    if not any(int(grade) > 0 for grade in relevance.values()):
        raise ValueError("relevance must contain at least one positive grade")
    for rank, document_id in enumerate(ranking, start=1):
        if int(relevance.get(document_id, 0)) > 0:
            return 1.0 / rank
    return 0.0

def recall_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return set recall among all positively relevant documents."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    relevant = {document_id for document_id, grade in relevance.items() if int(grade) > 0}
    if not relevant:
        raise ValueError("relevance must contain at least one positive grade")
    return len(set(ranking[:k]) & relevant) / len(relevant)

def ndcg_at_k(ranking: Sequence[str], relevance: Mapping[str, int], k: int) -> float:
    """Return nDCG@k with gains 2^grade-1 and log2 discount."""
    _validate_judgments(ranking, relevance)

    k = _positive_int("k", k)
    grades = [int(grade) for grade in relevance.values()]
    if any(grade < 0 for grade in grades):
        raise ValueError("relevance grades must be non-negative")
    if not any(grade > 0 for grade in grades):
        raise ValueError("relevance must contain at least one positive grade")
    dcg = sum(
        (2 ** int(relevance.get(document_id, 0)) - 1.0) / math.log2(rank + 1.0)
        for rank, document_id in enumerate(ranking[:k], start=1)
    )
    ideal = sorted(grades, reverse=True)[:k]
    idcg = sum((2**grade - 1.0) / math.log2(rank + 1.0) for rank, grade in enumerate(ideal, start=1))
    return dcg / idcg

def evaluate_rankings(
    rankings: Mapping[str, Sequence[str]],
    query_records: Sequence[Mapping[str, object]],
    *,
    k: int = 5,
) -> dict[str, float]:
    """Aggregate MRR, Recall@k, and nDCG@k over one frozen query set."""

    k = _positive_int("k", k)
    if not query_records or len({r["query_id"] for r in query_records}) != len(query_records):
        raise ValueError("queries must be nonempty with unique IDs")
    if len(rankings) != len(query_records):
        raise ValueError("rankings must contain exactly one entry per query")
    reciprocal_ranks: list[float] = []
    recalls: list[float] = []
    ndcgs: list[float] = []
    for record in query_records:
        query_id = str(record["query_id"])
        if query_id not in rankings:
            raise KeyError(f"missing ranking for {query_id}")
        relevance = dict(record["relevance"])
        ranking = tuple(rankings[query_id])
        if len(ranking) != len(set(ranking)):
            raise ValueError(f"ranking for {query_id} contains duplicates")
        reciprocal_ranks.append(reciprocal_rank(ranking, relevance))
        recalls.append(recall_at_k(ranking, relevance, k))
        ndcgs.append(ndcg_at_k(ranking, relevance, k))
    return {
        "mrr": float(np.mean(reciprocal_ranks)),
        f"recall_at_{k}": float(np.mean(recalls)),
        f"ndcg_at_{k}": float(np.mean(ndcgs)),
    }
texts = ["TS-999 update fails after version 4.7.3 restart",
         "TS-998 update fails after version 4.7.3 restart",
         "Restart clears the cache after an update",
         "An update changes the search index"]
ids = ["error999","error998","cache","index"]
query = "TS-999 update"
bm = BM25Index(ids,texts); ls = LSAIndex(ids,texts,dimension=2)
branches = {"bm25":[r.document_id for r in bm.search(query)],
            "lsa":[r.document_id for r in ls.search(query)]}
fused = reciprocal_rank_fusion(branches)
records = {}
for depth in (1,2,4):
    head = [r.document_id for r in fused[:depth]]
    ranked = rerank_candidates(query,head,dict(zip(ids,texts)))
    records[depth] = {"ranking":[r.document_id for r in ranked], "pairs_scored":len(ranked)}
qrels = {"error999":3}
for depth,row in records.items():
    row["recall_at_1"] = recall_at_k(row["ranking"],qrels,1)
assert sum(row["pairs_scored"] for row in records.values()) == 7
print(json.dumps({"branches":branches,"stored_results":records},indent=2))
# Work is counted as pairs, not invented milliseconds.
plt.rcParams.update({"font.family":"DejaVu Sans","font.size":12,
    "axes.titlesize":17,"axes.labelsize":13,"xtick.labelsize":11,"ytick.labelsize":11})
plt.figure(figsize=(11,5))
plt.plot(list(records),[r["recall_at_1"] for r in records.values()],"o-")
plt.xlabel("Оценённые пары" if DEMO_LANGUAGE=="ru" else "Scored pairs")
plt.ylabel("Recall@1");plt.ylim(-.05,1.1);plt.xticks([1,2,4]);plt.tight_layout();plt.show()

Compare stored rankings and Recall@1 against scored-pair count. A wider list helps only if it admits useful evidence and the scorer promotes it. A reranker cannot select a document it never received. A flat curve is informative too: extra work need not change the answer.

The second experiment gives two documents different first rows and identical remaining rows. Cutting off the first row before attention makes the inputs identical. Applying attention to the whole document before pooling the target rows allows the surrounding signal to reach the chunk. Follow the attention weights to identify that path.

Python

# An exact attention operation with designed features, not a pretrained language model.
# First row = antecedent; remaining two rows form the chunk to retrieve.
def contextualize(x):
    logits = x @ x.T / np.sqrt(x.shape[1])
    logits -= logits.max(axis=1,keepdims=True)
    weights = np.exp(logits)
    weights /= weights.sum(axis=1,keepdims=True)
    return weights @ x

river = np.array([[2.,0.],[0.,1.],[0.,1.]])
finance = np.array([[0.,2.],[0.,1.],[0.,1.]])
early = [contextualize(x[1:]).mean(axis=0) for x in (river,finance)]
late = [contextualize(x)[1:].mean(axis=0) for x in (river,finance)]
assert np.array_equal(early[0],early[1])
assert not np.allclose(late[0],late[1])
print({"early_vectors":[x.tolist() for x in early],"late_vectors":[x.tolist() for x in late]})
# Changing only the excluded antecedent changes the pooled chunk after full attention.
assert late[0][0] > late[1][0]

Early chunking produced identical vectors; late chunking produced different ones. Only a row outside the target chunk changed, isolating context transfer through attention. The hand-designed coordinates are not learned language representations, so this result establishes neither pronoun understanding nor better retrieval on real documents.

Scaling to the full project

One query is a diagnostic, not an evaluation distribution. For the project, separate development and held-out queries, choose or train the embedder without test-set tuning, and measure actual latency. The second experiment demonstrates an operator, not pronoun understanding: no language training occurred. Exercise E5 still requires a long-context encoder, real anaphora examples, and a separate ColPali versus OCR experiment over identical pages; this notebook runs neither and claims no visual-retrieval advantage.

Minimum report

Data, model, and code revisions, plus the random seed.
The primary metric, plus quality, cost, latency, and memory where applicable.
A baseline, a matched budget, and at least one negative control.
Failed runs and the limits of the supported conclusions.