"""Retrieval evaluation: Accuracy@k and MRR across chunking configurations.""" import itertools import logging from dataclasses import dataclass from pathlib import Path from ingestion.embedder import Embedder from ingestion.pipeline import IngestionPipeline from retrieval.index import VectorIndex from retrieval.searcher import search logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Placeholder test cases — replace with real questions from your PDFs. # Each dict needs: question, expected_source (filename), expected_page (1-based). # --------------------------------------------------------------------------- PLACEHOLDER_TEST_CASES: list[dict] = [ { "question": "What two components does RAG combine to generate answers?", "expected_source": "original_rag_paper.pdf", "expected_page": 1, }, { "question": "Which dataset is used to evaluate open-domain question answering in the RAG paper?", "expected_source": "original_rag_paper.pdf", "expected_page": 6, }, { "question": "What is the role of the retriever in the RAG architecture?", "expected_source": "original_rag_paper.pdf", "expected_page": 2, }, ] # --------------------------------------------------------------------------- # Metrics dataclass # --------------------------------------------------------------------------- @dataclass class EvalMetrics: accuracy_at_1: float accuracy_at_3: float accuracy_at_5: float mrr: float # Mean Reciprocal Rank n_queries: int def __str__(self) -> str: return ( f"Acc@1={self.accuracy_at_1:.2f} " f"Acc@3={self.accuracy_at_3:.2f} " f"Acc@5={self.accuracy_at_5:.2f} " f"MRR={self.mrr:.2f} " f"(n={self.n_queries})" ) # --------------------------------------------------------------------------- # Core evaluation # --------------------------------------------------------------------------- def evaluate_retrieval( test_cases: list[dict], embedder: Embedder, index: VectorIndex, k: int = 5, ) -> EvalMetrics: """Compute Accuracy@1/3/5 and MRR for a set of labelled questions. Each test case must have: question (str) — the query to embed and search expected_source (str) — filename of the expected chunk (e.g. "paper.pdf") expected_page (int) — 1-based page number of the expected chunk A result is considered a hit when both source and page_num match the expected values. MRR is 0 for queries where the expected chunk does not appear in the top-k results. Args: test_cases: List of labelled query dicts. embedder: Embedder used at ingestion time (must be the same model). index: Populated VectorIndex to evaluate against. k: Maximum rank to consider (search retrieves this many results). Returns: EvalMetrics with aggregated scores. """ if not test_cases: raise ValueError("test_cases must be non-empty") hits_at_1 = hits_at_3 = hits_at_5 = 0 reciprocal_ranks: list[float] = [] for case in test_cases: results = search(case["question"], embedder, index, k=k).chunks rank = _find_rank( results, expected_source=case["expected_source"], expected_page=int(case["expected_page"]), ) if rank is not None: if rank <= 1: hits_at_1 += 1 if rank <= 3: hits_at_3 += 1 if rank <= 5: hits_at_5 += 1 reciprocal_ranks.append(1.0 / rank) else: reciprocal_ranks.append(0.0) n = len(test_cases) return EvalMetrics( accuracy_at_1=hits_at_1 / n, accuracy_at_3=hits_at_3 / n, accuracy_at_5=hits_at_5 / n, mrr=sum(reciprocal_ranks) / n, n_queries=n, ) # --------------------------------------------------------------------------- # Chunking strategy comparison # --------------------------------------------------------------------------- def compare_chunking_strategies( pdf_paths: list[str | Path], test_cases: list[dict] | None = None, chunk_sizes: list[int] | None = None, overlaps: list[int] | None = None, ) -> list[dict]: """Re-ingest PDFs under every (chunk_size, overlap) combination and compare retrieval metrics. Builds a fresh VectorIndex for each configuration so results are independent. The Embedder is created once and shared across all runs. Only the recursive_character strategy is evaluated — varying chunk size and overlap is the most practically useful axis for that splitter. Args: pdf_paths: List of PDF paths to ingest for each configuration. test_cases: Labelled queries (defaults to PLACEHOLDER_TEST_CASES). chunk_sizes: Character lengths to try (default: [300, 500, 800]). overlaps: Overlap values to try (default: [0, 50]). Returns: List of result dicts, each with keys: chunk_size, overlap, and the four metric fields. Also prints a formatted comparison table. """ if test_cases is None: test_cases = PLACEHOLDER_TEST_CASES if chunk_sizes is None: chunk_sizes = [300, 500, 800] if overlaps is None: overlaps = [0, 50] pdf_paths = [Path(p) for p in pdf_paths] missing = [p for p in pdf_paths if not p.exists()] if missing: raise FileNotFoundError(f"PDF(s) not found: {missing}") print("Loading embedder (shared across all runs)...") embedder = Embedder() configs = list(itertools.product(chunk_sizes, overlaps)) rows: list[dict] = [] for chunk_size, overlap in configs: label = f"chunk={chunk_size}, overlap={overlap}" print(f"\nIngesting with {label}...") index = VectorIndex(dimension=embedder.dimension) pipeline = IngestionPipeline( embedder=embedder, index=index, strategy="recursive_character", chunk_size=chunk_size, overlap=overlap, ) for pdf_path in pdf_paths: result = pipeline.ingest_pdf(pdf_path) if result.error: logger.warning("Skipped %s: %s", pdf_path.name, result.error) else: print(f" {pdf_path.name}: {result.chunks} chunks") metrics = evaluate_retrieval(test_cases, embedder, index, k=5) rows.append({ "chunk_size": chunk_size, "overlap": overlap, "acc@1": metrics.accuracy_at_1, "acc@3": metrics.accuracy_at_3, "acc@5": metrics.accuracy_at_5, "mrr": metrics.mrr, }) _print_table(rows) return rows # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _find_rank(results, expected_source: str, expected_page: int) -> int | None: """Return the 1-based rank of the first matching result, or None.""" for rank, result in enumerate(results, start=1): meta = result.metadata source_match = Path(meta.get("source", "")).name == Path(expected_source).name page_match = meta.get("page_num") == expected_page if source_match and page_match: return rank return None def _print_table(rows: list[dict]) -> None: col_w = [10, 9, 8, 8, 8, 8] headers = ["chunk_size", "overlap", "Acc@1", "Acc@3", "Acc@5", "MRR"] divider = "-" * sum(col_w) print(f"\n{'Chunking strategy comparison (recursive_character)':^{sum(col_w)}}") print(divider) print("".join(h.ljust(w) for h, w in zip(headers, col_w))) print(divider) for row in rows: values = [ str(row["chunk_size"]), str(row["overlap"]), f"{row['acc@1']:.2f}", f"{row['acc@3']:.2f}", f"{row['acc@5']:.2f}", f"{row['mrr']:.2f}", ] print("".join(v.ljust(w) for v, w in zip(values, col_w))) print(divider)