diff --git a/nextcloud_mcp_server/config.py b/nextcloud_mcp_server/config.py index 97783d9c1..a8d7c3f29 100644 --- a/nextcloud_mcp_server/config.py +++ b/nextcloud_mcp_server/config.py @@ -1731,7 +1731,20 @@ def __post_init__(self): # rather than raise: the prefix set is the gateway's, not a closed # universe, and refusing to boot over a model name we cannot validate # would be worse than saying so. - if self.search_rerank_enabled and self.search_rerank_url: + # + # "Explicit URL" does NOT imply "not the gateway": pointing + # SEARCH_RERANK_URL at the gateway's own /v1/rerank is a legitimate + # configuration (it is how you pin the endpoint while still using the + # gateway), and the routing prefix is correct there. Only warn when the + # URL is somewhere OTHER than the configured gateway — otherwise the + # warning fires on a working setup and trains operators to ignore it. + _direct_rerank = bool(self.search_rerank_url) and not ( + self.embedding_gateway_url + and self.search_rerank_url.startswith( + self.embedding_gateway_url.rstrip("/") + ) + ) + if self.search_rerank_enabled and _direct_rerank: _prefix, _, _bare = self.search_rerank_model.partition("/") if _prefix in GATEWAY_MODEL_NAMESPACES: logger.warning( diff --git a/scripts/retrieval_eval/README.md b/scripts/retrieval_eval/README.md new file mode 100644 index 000000000..78b559766 --- /dev/null +++ b/scripts/retrieval_eval/README.md @@ -0,0 +1,63 @@ +# retrieval_eval — OHR-Bench retrieval / generation eval harness + +Self-contained, reproducible benchmark for the document-processing retrieval +pipeline. Reconstructs the lost `chunk_eval2.py` scratchpad (notes 390421 / +390460) as committed code. Where `scripts/chunk_density_sweep.py` measures only +**density** (chunks/MB, offline), this measures **quality** end to end: + +``` +extract (pypdfium2_fast) → chunk (PageAwareChunker) → embed (gateway/Ollama) + → index (networked Qdrant: dense + BM25 sparse) → hybrid retrieve (RRF/DBSF) + → optional rerank → score (OHR-Bench official LCS / page-hit / F1) +``` + +Used to benchmark three "new-model" levers across all 7 OHR-Bench domains: +**embedding model**, **reranker**, **generation model** (Deck board 12 card +#651; productionization design in note 390473). + +## Prerequisites +- `~/Downloads/OHR-Bench//*.pdf` (1,261 born-digital PDFs) and + `~/Software/OHR-Bench/data/qas_v2.json` (gold Q&A). +- A **networked** Qdrant (`QDRANT_URL` + `QDRANT_API_KEY`) — directory mode + throttles indexing throughput (note 390460); it is intentionally unsupported. +- Tailnet access to the embedding gateway (`EMBED_GATEWAY_URL`, default + `astrolabe-gateway-dev.tail148d5.ts.net`, no bearer) and, for the self-hosted + embedder, Ollama (`OLLAMA_HOST`). These are exported by the repo `.envrc`. + +## Metric fidelity +`metrics.py` copies OHR-Bench `src/metric/common.py` (`lcs_score`, `f1_score`, +`exact_match_score`, `normalize_answer`) verbatim, so numbers are +leaderboard-comparable. Page-gating mirrors `src/tasks/retrieval.py` with one +documented adaptation: packed chunks span a page RANGE, so a chunk counts for +every page it covers (the rule used for packed configs in notes 390421/390460). +Unit-tested in `tests/unit/test_retrieval_eval_metrics.py`. + +## Usage +```bash +OUT=./retrieval_eval_out # gitignored + +# 1. Sample a stratified plan (40 questions/domain, all 7 domains) +uv run python -m scripts.retrieval_eval plan --per-domain 40 --out $OUT/plan.jsonl + +# 2. Index one (embedder × strategy) cell (one Qdrant collection per cell) +uv run python -m scripts.retrieval_eval index \ + --embedder mistral-embed --strategy page-pack@4096 --recreate --out-dir $OUT + +# 3. Retrieval eval (± rerank), scored + sliced by domain +uv run python -m scripts.retrieval_eval eval \ + --embedder mistral-embed --strategy page-pack@4096 --plan $OUT/plan.jsonl \ + --rerank strong --out-dir $OUT # rerank: none | strong | cheap + +# 4. Generation eval (F1/EM) with an OpenRouter chat model +uv run python -m scripts.retrieval_eval generate \ + --embedder mistral-embed --strategy page-pack@4096 --plan $OUT/plan.jsonl \ + --chat-model openrouter/openai/gpt-4o-mini --out-dir $OUT +``` + +Registries (`__main__.py`): `EMBEDDERS` (mistral-embed, titan-v2, te3-small, +te3-large, arctic-110m), `STRATEGIES` (page-pack@4096 default + prod control), +`RERANKERS` (none / strong=`openrouter/cohere/rerank-v3.5` / +cheap=`Xenova/ms-marco-MiniLM-L-6-v2`), `CHAT_MODELS`. Smoke any cell fast with +`--domains finance --limit 2`. +``` +``` diff --git a/scripts/retrieval_eval/__init__.py b/scripts/retrieval_eval/__init__.py new file mode 100644 index 000000000..8883771d1 --- /dev/null +++ b/scripts/retrieval_eval/__init__.py @@ -0,0 +1,20 @@ +"""Self-contained OHR-Bench retrieval/generation eval harness (committed). + +Reconstructs and supersedes the lost ``chunk_eval2.py`` scratchpad that produced +Nextcloud notes 390421 / 390460. Where ``scripts/chunk_density_sweep.py`` covers +only the *density* half (chunks/MB, no network), this package covers the +*quality* half end to end: + + extract (pypdfium2_fast) -> chunk (PageAwareChunker) -> embed (gateway/Ollama) + -> index (networked Qdrant, dense + BM25 sparse) -> hybrid retrieve (RRF/DBSF) + -> optional rerank -> score (OHR-Bench official LCS / page-hit / F1) + +It is used to benchmark three "new-model" levers that prior notes left open: +embedding model, reranker, and generation model — across all 7 OHR-Bench domains. +See Deck board 12 card #651 and note 390473 (productionization design). + +Run ``uv run python -m scripts.retrieval_eval --help`` for the CLI. + +The metric functions in :mod:`scripts.retrieval_eval.metrics` are ported verbatim +from OHR-Bench (``src/metric/common.py``) so results are leaderboard-comparable. +""" diff --git a/scripts/retrieval_eval/__main__.py b/scripts/retrieval_eval/__main__.py new file mode 100644 index 000000000..d92e69f66 --- /dev/null +++ b/scripts/retrieval_eval/__main__.py @@ -0,0 +1,900 @@ +"""CLI for the OHR-Bench retrieval/generation eval harness. + +Env (see repo ``.envrc``): + EMBED_GATEWAY_URL default https://astrolabe-gateway-dev.tail148d5.ts.net + OLLAMA_HOST default https://ollama.internal.coutinho.io + QDRANT_URL default http://localhost:6333 + QDRANT_API_KEY Qdrant auth (required by the networked container) + +Typical run (one embedder × strategy cell): + uv run python -m scripts.retrieval_eval plan --per-domain 40 --out $OUT/plan.jsonl + uv run python -m scripts.retrieval_eval index --embedder mistral-embed --strategy page-pack@4096 + uv run python -m scripts.retrieval_eval eval --embedder mistral-embed --strategy page-pack@4096 \ + --plan $OUT/plan.jsonl --rerank none --out-dir $OUT +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import random +import time +from collections import defaultdict +from pathlib import Path + +import anyio +import httpx +from qdrant_client import AsyncQdrantClient + +from . import metrics +from .clients import ( + Bm25Encoder, + EmbedderSpec, + GatewayChat, + GatewayReranker, + LocalCrossEncoderReranker, + make_embedder, +) +from .ocr import MistralSyncOCR, SuryaBatchOCR +from .pipeline import ( + IndexStats, + clone_collection, + embed_and_upsert, + ensure_collection, + eur_per_gib_mo, + extract_and_chunk, + index_corpus, + iter_corpus, + records_from_ocr, + retrieve, +) + +logger = logging.getLogger("retrieval_eval") + +DEFAULT_CORPUS = Path("~/Downloads/OHR-Bench").expanduser() +DEFAULT_QAS = Path("~/Software/OHR-Bench/data/qas_v2.json").expanduser() +GATEWAY_URL = os.environ.get( + "EMBED_GATEWAY_URL", "https://astrolabe-gateway-dev.tail148d5.ts.net" +) +OLLAMA_URL = os.environ.get("OLLAMA_HOST") or os.environ.get( + "OLLAMA_BASE_URL", "https://ollama.internal.coutinho.io" +) + +# --- model registry (discovered 2026-07-12; Deck #651) --------------------- +EMBEDDERS: dict[str, EmbedderSpec] = { + "mistral-embed": EmbedderSpec( + "mistral-embed", "gateway", "mistral/mistral-embed", 1024 + ), + "titan-v2": EmbedderSpec( + "titan-v2", "gateway", "bedrock/amazon.titan-embed-text-v2:0", 1024 + ), + "te3-small": EmbedderSpec( + "te3-small", "gateway", "openrouter/openai/text-embedding-3-small", 1536 + ), + "te3-large": EmbedderSpec( + "te3-large", "gateway", "openrouter/openai/text-embedding-3-large", 3072 + ), + "arctic-110m": EmbedderSpec( + "arctic-110m", "ollama", "snowflake-arctic-embed:110m", 768 + ), + # Self-hosted open embedder. Never benchmarked before, and it is what a + # gateway-free self-hoster running Infinity/vLLM actually uses. + "bge-m3": EmbedderSpec("bge-m3", "gateway", "local/BAAI/bge-m3", 1024), + "titan-v1": EmbedderSpec( + "titan-v1", "gateway", "bedrock/amazon.titan-embed-text-v1", 1536 + ), +} + +# strategy name -> (chunk_size, pack_pages, collection slug) +STRATEGIES: dict[str, tuple[int, bool, str]] = { + "page-pack@4096": (4096, True, "pp4096"), + "page-aware@2048": (2048, False, "pa2048"), # production baseline / control + "page-pack@2048": (2048, True, "pp2048"), + "page-aware@4096": (4096, False, "pa4096"), + # Bracket the tested 2048/4096 pair so the chunk-size axis has a shape, not + # just two points. 512 is below a typical page, 8192 packs several. + "page-pack@1024": (1024, True, "pp1024"), + "page-pack@8192": (8192, True, "pp8192"), + "page-aware@1024": (1024, False, "pa1024"), +} + +# rerank presets -> (kind, model) ; kind in {"none","gateway","local"} +RERANKERS: dict[str, tuple[str, str]] = { + "none": ("none", ""), + # THE GAP this run exists to close: `local/BAAI/bge-reranker-v2-m3` is the + # SHIPPED default (SEARCH_RERANK_MODEL) and the only model ADR-034 fitted a + # calibration curve for, yet it has never been measured for RETRIEVAL + # quality. Prior work benchmarked cohere-v3.5 (+) and ms-marco (-) — neither + # is what we ship. + "bge": ("gateway", "local/BAAI/bge-reranker-v2-m3"), + # Strong reference point. Re-pointed: the gateway no longer serves + # `openrouter/cohere/rerank-v3.5`; the Bedrock-routed id is the live one. + "strong": ("gateway", "bedrock/cohere.rerank-v3-5:0"), + "amazon": ("gateway", "bedrock/amazon.rerank-v1:0"), + "cheap": ("local", "Xenova/ms-marco-MiniLM-L-6-v2"), +} + +CHAT_MODELS = [ + "openrouter/openai/gpt-4o-mini", + "openrouter/meta-llama/llama-3.3-70b-instruct", + "openrouter/qwen/qwen-2.5-72b-instruct", + "bedrock/eu.amazon.nova-lite-v1:0", + # Native mistral/ provider (same vendor as the production embedder) — routed + # by the gateway even though /v1/chat/models doesn't advertise them. + "mistral/ministral-8b-latest", + "mistral/ministral-3b-latest", + "mistral/mistral-small-latest", +] + + +def collection_name(embedder: str, strategy: str, suffix: str = "") -> str: + slug = STRATEGIES[strategy][2] + return f"reval_{EMBEDDERS[embedder].slug}_{slug}{suffix}" + + +# OCR engine -> (mode, model). surya=batch (triggers leaf.cloud GPU); mistral=sync (cloud). +OCR_ENGINES: dict[str, tuple[str, str]] = { + "surya": ("batch", "surya/surya-ocr-2"), + "mistral": ("sync", "mistral/mistral-ocr-4-0"), +} + + +def _gateway_client() -> httpx.AsyncClient: + return httpx.AsyncClient(base_url=GATEWAY_URL, timeout=httpx.Timeout(180.0)) + + +def _ollama_client() -> httpx.AsyncClient: + return httpx.AsyncClient(base_url=OLLAMA_URL, timeout=httpx.Timeout(180.0)) + + +def _qdrant() -> AsyncQdrantClient: + return AsyncQdrantClient( + url=os.environ.get("QDRANT_URL", "http://localhost:6333"), + api_key=os.environ.get("QDRANT_API_KEY"), + timeout=120, + ) + + +# --------------------------------------------------------------------------- +# plan +# --------------------------------------------------------------------------- +def cmd_plan(args: argparse.Namespace) -> None: + qas = json.loads(Path(args.qas).read_text()) + by_domain: dict[str, list[dict]] = defaultdict(list) + for q in qas: + by_domain[q["doc_type"]].append(q) + domains = args.domains or sorted(by_domain) + rng = random.Random(args.seed) + sample: list[dict] = [] + for domain in domains: + pool = by_domain.get(domain, []) + sample += rng.sample(pool, min(args.per_domain, len(pool))) + rng.shuffle(sample) + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("w", encoding="utf-8") as f: + for q in sample: + f.write( + json.dumps( + { + "ID": q["ID"], + "doc_name": q["doc_name"], + "doc_type": q["doc_type"], + "question": q["questions"], + "answer": q["answers"], + "evidence_source": q["evidence_source"], + "evidence_page_no": q["evidence_page_no"], + "evidence_context": q["evidence_context"], + # The Nextcloud content type the gold document lives in + # (note / deck_card / file / ...). Carried on every plan + # row so the eval can slice by it — see the BY SOURCE + # TYPE block in _report_eval for why that slice is + # mandatory on a mixed corpus. + # + # OHR-Bench is PDFs only, so a generator that does not + # set it yields None and the slice is omitted rather + # than fabricating a single-bucket comparison. + "source_type": q.get("source_type"), + }, + ensure_ascii=False, + ) + + "\n" + ) + counts = defaultdict(int) + for q in sample: + counts[q["doc_type"]] += 1 + print(f"Wrote {len(sample)} questions to {out}") + for d in sorted(counts): + print(f" {d:16} {counts[d]}") + + +def _load_plan(path: str) -> dict[str, dict]: + plans: dict[str, dict] = {} + with open(path, encoding="utf-8") as f: + for line in f: + p = json.loads(line) + plans[p["ID"]] = p + return plans + + +# --------------------------------------------------------------------------- +# index +# --------------------------------------------------------------------------- +async def _run_index(args: argparse.Namespace) -> None: + spec = EMBEDDERS[args.embedder] + chunk_size, pack_pages, _ = STRATEGIES[args.strategy] + coll = collection_name(args.embedder, args.strategy, args.suffix) + limiter = anyio.CapacityLimiter(args.embed_concurrency) + bm25 = Bm25Encoder() + + async with _gateway_client() as gw, _ollama_client() as oll: + qdrant = _qdrant() + try: + embedder = make_embedder(spec, gateway=gw, ollama=oll, limiter=limiter) + # confirm live dimension before creating the collection + probe, _ = await embedder.embed(["dimension probe"]) + dim = len(probe[0]) + if dim != spec.dim: + logger.warning( + "declared dim %d != live dim %d for %s", spec.dim, dim, spec.name + ) + spec.dim = dim + created = await ensure_collection(qdrant, coll, dim, recreate=args.recreate) + if not created: + print( + f"collection {coll} exists — skipping (use --recreate to rebuild)" + ) + return + doc_allow: set[str] | None = None + if args.from_plan: + doc_allow = {p["doc_name"] for p in _load_plan(args.from_plan).values()} + logger.info( + "indexing %d gold docs from %s", len(doc_allow), args.from_plan + ) + t0 = time.monotonic() + stats = await index_corpus( + qdrant=qdrant, + collection=coll, + embedder=embedder, + bm25=bm25, + spec=spec, + corpus=Path(args.corpus).expanduser(), + domains=args.domains, + limit=args.limit, + chunk_size=chunk_size, + pack_pages=pack_pages, + doc_allow=doc_allow, + ) + finally: + await qdrant.close() + + elapsed = time.monotonic() - t0 + report = { + "collection": coll, + "embedder": spec.name, + "model": spec.model, + "dim": stats.dim, + "strategy": args.strategy, + "docs": stats.docs, + "empty_docs": stats.empty_docs, # scanned/no-text PDFs, not indexed + "chunks": stats.total_chunks, + "embed_tokens": stats.embed_tokens, + "seconds": round(elapsed, 1), + "by_domain": { + d: { + "chunks": dd.chunks, + "chunks_per_mb": round(dd.chunks_per_mb, 1), + "eur_per_gib_mo": round(eur_per_gib_mo(dd.chunks_per_mb, stats.dim), 3), + } + for d, dd in sorted(stats.by_domain.items()) + }, + } + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / f"density_{coll}.json").write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2)) + + +# --------------------------------------------------------------------------- +# eval (retrieve + optional rerank + score) +# --------------------------------------------------------------------------- +def _make_reranker(name: str, gw: httpx.AsyncClient, limiter: anyio.CapacityLimiter): + kind, model = RERANKERS[name] + if kind == "none": + return None + if kind == "gateway": + return GatewayReranker(gw, model, limiter) + return LocalCrossEncoderReranker(model) + + +async def _run_eval(args: argparse.Namespace) -> None: + spec = EMBEDDERS[args.embedder] + coll = collection_name(args.embedder, args.strategy, args.suffix) + plans = _load_plan(args.plan) + limiter = anyio.CapacityLimiter(args.embed_concurrency) + bm25 = Bm25Encoder() + results: dict[str, list[dict]] = {} + q_limiter = anyio.CapacityLimiter(args.query_concurrency) + lock = anyio.Lock() + + async with _gateway_client() as gw, _ollama_client() as oll: + qdrant = _qdrant() + embedder = make_embedder(spec, gateway=gw, ollama=oll, limiter=limiter) + reranker = _make_reranker(args.rerank, gw, limiter) + + # A reranker reorders a broad pool into the top-k; hybrid-only fetches 2k. + # + # CONFOUND: with the defaults that is 50 vs 20, so a rerank-vs-none A/B + # compares TWO things at once — reordering skill AND a 2.5x deeper + # candidate pool. The pool alone lifts recall (more gold documents make + # the top-k at all), which is not what "does reranking help" is asking. + # `--fetch` forces both arms to the same depth so the only difference is + # the reordering. Note 390487's "gain grows with pool size + # (+0.026@20 -> +0.051@50)" is exactly the signature this produces. + fetch = args.fetch or (args.rerank_pool if reranker is not None else args.k * 2) + + async def handle(pid: str, plan: dict) -> None: + async with q_limiter: + cands = await retrieve( + qdrant=qdrant, + collection=coll, + query=plan["question"], + embedder=embedder, + bm25=bm25, + k=args.k, + fusion=args.fusion, + fetch=fetch, + rescore=args.rescore, + oversampling=args.oversampling, + ) + if reranker is not None and cands: + order = await reranker.rerank( + plan["question"], [c["text"] for c in cands], top_n=args.k + ) + cands = [cands[i] for i, _ in order] + else: + cands = cands[: args.k] + async with lock: + results[pid] = cands + + try: + async with anyio.create_task_group() as tg: + for pid, plan in plans.items(): + tg.start_soon(handle, pid, plan) + finally: + await qdrant.close() + + _report_retrieval(args, coll, plans, results) + + +def _report_retrieval( + args: argparse.Namespace, + coll: str, + plans: dict[str, dict], + results: dict[str, list[dict]], +) -> None: + offset = metrics.calibrate_offset(plans, results) + answered = [p for p in plans if p in results] + domains = sorted({plans[p]["doc_type"] for p in answered}) + sources = sorted({plans[p]["evidence_source"] for p in answered}) + + def fmt(label: str, m: dict | None) -> str: + """One line per slice, BOTH metric families side by side. + + The rank-sensitive columns are not optional extras. On the same runs the + two families have disagreed in magnitude by 14x and, across chunk sizes, + in sign — so printing only the page-gated ones lets the reader draw a + conclusion the data does not support. + """ + if not m: + return f" {label:22} (no data)" + return ( + f" {label:22} n={int(m['n']):<4} page_lcs={m['page_lcs']:.3f} " + f"page_hit={m['page_hit']:.3f} doc_hit={m['doc_hit']:.3f} " + f"doc_lcs={m['doc_lcs']:.3f} | S@1={m['success_at_1']:.3f} " + f"S@3={m['success_at_3']:.3f} MRR={m['mrr']:.3f} " + f"rank={m['mean_gold_rank']:.2f} found={int(m['found'])}" + ) + + print(f"\n=== {coll} rerank={args.rerank} fusion={args.fusion} k={args.k} ===") + print( + f"offset={offset} (gold = page_number - offset), answered {len(answered)}/{len(plans)}" + ) + print("OVERALL") + print(fmt("all", metrics.aggregate(plans, results, answered, offset))) + print("BY DOMAIN") + per_domain = {} + for d in domains: + m = metrics.aggregate( + plans, results, [p for p in answered if plans[p]["doc_type"] == d], offset + ) + per_domain[d] = m + print(fmt(d, m)) + print("BY EVIDENCE SOURCE") + for s in sources: + print( + fmt( + s, + metrics.aggregate( + plans, + results, + [p for p in answered if plans[p]["evidence_source"] == s], + offset, + ), + ) + ) + + # BY SOURCE TYPE — the Nextcloud content type a gold document lives in + # (note / deck_card / file / mail_message / ...), carried on the plan row as + # `source_type`. + # + # This slice is MANDATORY for any mixed-content corpus and must never be + # pooled away into a single headline number. The failure mode it exists to + # expose is LENGTH ASYMMETRY: a 20-word Deck card and an 800-token PDF chunk + # compete in the same Qdrant collection, and the density spread across our + # own doc types is roughly 47x (1.2 -> 57 chunks/MB). A system that is + # excellent on files and useless on cards posts a respectable aggregate + # score, and the aggregate is what gets quoted. + # + # Absent on the OHR-Bench corpus, which is PDFs only — the slice simply does + # not print there rather than inventing a single "file" bucket that would + # imply a comparison the corpus cannot support. + per_source_type = {} + source_types = sorted({st for p in answered if (st := plans[p].get("source_type"))}) + if source_types: + print("BY SOURCE TYPE") + for st in source_types: + m = metrics.aggregate( + plans, + results, + [p for p in answered if plans[p].get("source_type") == st], + offset, + ) + per_source_type[st] = m + print(fmt(st, m)) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + tag = f"{coll}_rr-{args.rerank}_{args.fusion}" + (out_dir / f"metrics_{tag}.json").write_text( + json.dumps( + { + "collection": coll, + "rerank": args.rerank, + "fusion": args.fusion, + "k": args.k, + "offset": offset, + "answered": len(answered), + "overall": metrics.aggregate(plans, results, answered, offset), + "by_domain": per_domain, + # Empty on a single-type corpus (OHR-Bench). Present and + # non-empty is the signal that per-type numbers exist and the + # overall figure must not be quoted on its own. + "by_source_type": per_source_type, + }, + indent=2, + ) + ) + (out_dir / f"results_{tag}.jsonl").write_text( + "\n".join(json.dumps({"id": pid, "results": results[pid]}) for pid in answered) + ) + + +# --------------------------------------------------------------------------- +# generate (retrieve context -> chat -> F1/EM) +# --------------------------------------------------------------------------- +_GEN_PROMPT = ( + "Answer the question using ONLY the context. Reply with the shortest exact answer " + "(a number, name, or short phrase); no explanation.\n\nContext:\n{context}\n\n" + "Question: {question}\nAnswer:" +) + + +async def _run_generate(args: argparse.Namespace) -> None: + spec = EMBEDDERS[args.embedder] + coll = collection_name(args.embedder, args.strategy, args.suffix) + plans = _load_plan(args.plan) + limiter = anyio.CapacityLimiter(args.embed_concurrency) + chat_limiter = anyio.CapacityLimiter(args.chat_concurrency) + bm25 = Bm25Encoder() + scored: dict[str, dict] = {} + lock = anyio.Lock() + + async with _gateway_client() as gw, _ollama_client() as oll: + qdrant = _qdrant() + embedder = make_embedder(spec, gateway=gw, ollama=oll, limiter=limiter) + reranker = _make_reranker(args.rerank, gw, limiter) + chat = GatewayChat(gw, args.chat_model, chat_limiter) + + async def handle(pid: str, plan: dict) -> None: + cands = await retrieve( + qdrant=qdrant, + collection=coll, + query=plan["question"], + embedder=embedder, + bm25=bm25, + k=args.k, + fusion=args.fusion, + ) + if reranker is not None and cands: + order = await reranker.rerank( + plan["question"], [c["text"] for c in cands], top_n=args.k + ) + cands = [cands[i] for i, _ in order] + else: + cands = cands[: args.k] + context = "\n\n".join(c["text"] for c in cands) + try: + answer = await chat.generate( + _GEN_PROMPT.format( + context=context[: args.max_context_chars], + question=plan["question"], + ), + max_tokens=args.max_tokens, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("gen failed %s: %s", pid, exc) + return + s = metrics.score_generation(answer, plan["answer"]) + async with lock: + scored[pid] = {**s, "answer": answer} + + try: + async with anyio.create_task_group() as tg: + for pid, plan in plans.items(): + tg.start_soon(handle, pid, plan) + finally: + await qdrant.close() + + _report_generation(args, coll, plans, scored) + + +def _report_generation( + args: argparse.Namespace, coll: str, plans: dict[str, dict], scored: dict[str, dict] +) -> None: + answered = [p for p in plans if p in scored] + by_domain: dict[str, list[str]] = defaultdict(list) + for p in answered: + by_domain[plans[p]["doc_type"]].append(p) + + def mean(pids: list[str], key: str) -> float: + return sum(scored[p][key] for p in pids) / len(pids) if pids else 0.0 + + print( + f"\n=== GENERATION {coll} chat={args.chat_model} rerank={args.rerank} k={args.k} ===" + ) + print(f"answered {len(answered)}/{len(plans)}") + print( + f" {'all':16} n={len(answered):<4} F1={mean(answered, 'f1'):.3f} EM={mean(answered, 'em'):.3f}" + ) + for d in sorted(by_domain): + pids = by_domain[d] + print( + f" {d:16} n={len(pids):<4} F1={mean(pids, 'f1'):.3f} EM={mean(pids, 'em'):.3f}" + ) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + chat_slug = args.chat_model.replace("/", "_") + (out_dir / f"gen_{coll}_{chat_slug}_rr-{args.rerank}.json").write_text( + json.dumps( + { + "collection": coll, + "chat_model": args.chat_model, + "rerank": args.rerank, + "overall": { + "f1": mean(answered, "f1"), + "em": mean(answered, "em"), + "n": len(answered), + }, + "by_domain": { + d: {"f1": mean(p, "f1"), "em": mean(p, "em"), "n": len(p)} + for d, p in sorted(by_domain.items()) + }, + }, + indent=2, + ) + ) + + +# --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# ocr — fill the born-digital blind spot: OCR the scanned gold docs and index +# --------------------------------------------------------------------------- +async def _find_pending( + corpus: Path, gold: set[str], chunk_size: int, pack_pages: bool +) -> list[tuple[str, str, Path]]: + """Gold docs that yield NO born-digital text (scanned) -> need OCR.""" + pending: list[tuple[str, str, Path]] = [] + for doc_key, domain, pdf in iter_corpus(corpus, None, None, gold): + try: + chunks = await extract_and_chunk( + pdf, chunk_size=chunk_size, pack_pages=pack_pages + ) + except Exception: # noqa: BLE001 - bad PDF => treat as OCR-pending + chunks = [] + if not chunks: + pending.append((doc_key, domain, pdf)) + return pending + + +async def _run_ocr(args: argparse.Namespace) -> None: + spec = EMBEDDERS[args.embedder] + chunk_size, pack_pages, _ = STRATEGIES[args.strategy] + coll = collection_name(args.embedder, args.strategy, args.suffix) + mode, model = OCR_ENGINES[args.engine] + gold = {p["doc_name"] for p in _load_plan(args.plan).values()} + corpus = Path(args.corpus).expanduser() + + pending = await _find_pending(corpus, gold, chunk_size, pack_pages) + if args.limit_docs: + pending = pending[: args.limit_docs] + logger.info( + "%d OCR-pending gold docs; engine=%s model=%s mode=%s", + len(pending), + args.engine, + model, + mode, + ) + dom_by = {dk: dom for dk, dom, _ in pending} + size_by = {dk: pdf.stat().st_size for dk, _, pdf in pending} + bm25 = Bm25Encoder() + t0 = time.monotonic() + + # --- Phase A: OCR on its own gateway client --- + results: dict[str, list[str]] = {} + async with _gateway_client() as gw: + if mode == "sync": + ocr = MistralSyncOCR(gw, model, anyio.CapacityLimiter(args.ocr_concurrency)) + lock = anyio.Lock() + + async def do_one(doc_key: str, pdf: Path) -> None: + try: + pages = await ocr.ocr_doc(pdf.read_bytes()) + except Exception as exc: # noqa: BLE001 - one bad doc must not abort + logger.warning("ocr failed %s: %s", doc_key, exc) + return + async with lock: + results[doc_key] = pages + + async with anyio.create_task_group() as tg: + for doc_key, _, pdf in pending: + tg.start_soon(do_one, doc_key, pdf) + else: # batch (surya) — triggers leaf.cloud GPU (unless --reuse-job) + batch = SuryaBatchOCR(gw, model, poll_seconds=args.poll_seconds) + docs = ( + [] + if args.reuse_job + else [(dk, pdf.read_bytes()) for dk, _, pdf in pending] + ) + results = await batch.ocr_docs(docs, reuse_job_id=args.reuse_job) + ocr_seconds = time.monotonic() - t0 + + items = [ + (dk, dom_by[dk], size_by[dk], results[dk]) for dk in results if dk in dom_by + ] + records = await records_from_ocr( + items, chunk_size=chunk_size, pack_pages=pack_pages + ) + + # --- Phase B: embed + upsert on a FRESH gateway client. The OCR poll can hold + # the client idle for many minutes; reusing that aged pool at high embed + # concurrency triggered a raw ssl.SSLError. A fresh client sidesteps it. + limiter = anyio.CapacityLimiter(args.embed_concurrency) + stats = IndexStats(dim=spec.dim) + async with _gateway_client() as gw2, _ollama_client() as oll: + qdrant = _qdrant() + try: + if not await qdrant.collection_exists(coll): + raise SystemExit( + f"collection {coll} missing — index born-digital first: " + f"index --embedder {args.embedder} --suffix {args.suffix} " + f"--from-plan {args.plan} --recreate" + ) + embedder = make_embedder(spec, gateway=gw2, ollama=oll, limiter=limiter) + await embed_and_upsert( + qdrant=qdrant, + collection=coll, + embedder=embedder, + bm25=bm25, + records=records, + stats=stats, + ) + finally: + await qdrant.close() + + report = { + "collection": coll, + "engine": args.engine, + "model": model, + "pending_docs": len(pending), + "ocr_docs_returned": len(results), + "indexed_docs": stats.docs, + "indexed_chunks": stats.total_chunks, + "ocr_seconds": round(ocr_seconds, 1), + "by_domain": {d: dd.chunks for d, dd in sorted(stats.by_domain.items())}, + } + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / f"ocr_{coll}_{args.engine}.json").write_text( + json.dumps(report, indent=2) + ) + print(json.dumps(report, indent=2)) + + +# --------------------------------------------------------------------------- +# quantize — clone a collection applying int8/binary quantization (no re-embed) +# --------------------------------------------------------------------------- +async def _run_quantize(args: argparse.Namespace) -> None: + source = collection_name(args.embedder, args.strategy, args.from_suffix) + target = collection_name(args.embedder, args.strategy, args.suffix) + qdrant = _qdrant() + try: + if not await qdrant.collection_exists(source): + raise SystemExit(f"source collection {source} missing") + info = await qdrant.get_collection(source) + dim = info.config.params.vectors["dense"].size # type: ignore[index,union-attr] + t0 = time.monotonic() + copied = await clone_collection( + qdrant, source, target, dim, quantization=args.mode + ) + finally: + await qdrant.close() + print( + json.dumps( + { + "source": source, + "target": target, + "mode": args.mode, + "points_copied": copied, + "seconds": round(time.monotonic() - t0, 1), + }, + indent=2, + ) + ) + + +def _add_common(p: argparse.ArgumentParser) -> None: + p.add_argument("--embedder", required=True, choices=sorted(EMBEDDERS)) + p.add_argument("--strategy", default="page-pack@4096", choices=sorted(STRATEGIES)) + p.add_argument( + "--suffix", + default="", + help="collection-name suffix (e.g. _surya, _mocr) to keep OCR-engine " + "variants separate from the born-digital baseline", + ) + p.add_argument("--out-dir", default="retrieval_eval_out") + # The gateway serves each request's texts serially (~1.6 s/text) but + # parallelizes across requests, scaling near-linearly to ~48 (measured, 0 + # errors). High concurrency is what makes the full sweep tractable. + p.add_argument("--embed-concurrency", type=int, default=48) + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" + ) + # Silence per-request HTTP chatter (httpx/qdrant/fastembed) so progress logs + # and the final report are readable; keep our own logger at INFO. + for noisy in ("httpx", "httpcore", "qdrant_client", "fastembed", "huggingface_hub"): + logging.getLogger(noisy).setLevel(logging.WARNING) + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + sub = ap.add_subparsers(dest="cmd", required=True) + + p = sub.add_parser("plan", help="stratified-sample questions -> plan.jsonl") + p.add_argument("--qas", default=str(DEFAULT_QAS)) + p.add_argument("--per-domain", type=int, default=40) + p.add_argument("--domains", nargs="*", default=None) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--out", default="retrieval_eval_out/plan.jsonl") + p.set_defaults(func=lambda a: cmd_plan(a)) + + p = sub.add_parser("index", help="index the corpus into Qdrant for one cell") + _add_common(p) + p.add_argument("--corpus", default=str(DEFAULT_CORPUS)) + p.add_argument("--domains", nargs="*", default=None) + p.add_argument( + "--limit", type=int, default=None, help="max PDFs per domain (smoke)" + ) + p.add_argument( + "--from-plan", + default=None, + help="index only the gold docs referenced by this plan.jsonl " + "(the gold-doc corpus of notes 390421/390460)", + ) + p.add_argument("--recreate", action="store_true", help="drop + rebuild if exists") + p.set_defaults(func=lambda a: anyio.run(_run_index, a)) + + p = sub.add_parser("eval", help="retrieve (+rerank) + score against a plan") + _add_common(p) + p.add_argument("--plan", required=True) + p.add_argument("--k", type=int, default=10) + p.add_argument("--rerank", default="none", choices=sorted(RERANKERS)) + p.add_argument( + "--rerank-pool", + type=int, + default=50, + help="candidates fetched for the reranker to reorder into top-k (rerank only)", + ) + p.add_argument( + "--fetch", + type=int, + default=None, + help=( + "force the candidate depth for BOTH arms (matched-pool A/B). " + "Without it the rerank arm fetches --rerank-pool while the " + "baseline fetches 2k, which confounds reordering with pool depth." + ), + ) + p.add_argument("--fusion", default="rrf", choices=["rrf", "dbsf"]) + p.add_argument("--query-concurrency", type=int, default=12) + p.add_argument( + "--rescore", + action=argparse.BooleanOptionalAction, + default=None, + help="quantization rescore (recall recovery on int8/binary collections)", + ) + p.add_argument( + "--oversampling", + type=float, + default=None, + help="quantization oversampling factor (widens the quantized shortlist; binary)", + ) + p.set_defaults(func=lambda a: anyio.run(_run_eval, a)) + + p = sub.add_parser( + "quantize", + help="clone a collection applying int8/binary quantization (no re-embed)", + ) + _add_common(p) + p.add_argument("--mode", required=True, choices=["none", "int8", "binary"]) + p.add_argument( + "--from-suffix", default="", help="source collection suffix (default: baseline)" + ) + p.set_defaults(func=lambda a: anyio.run(_run_quantize, a)) + + p = sub.add_parser("generate", help="retrieve context -> chat -> F1/EM") + _add_common(p) + p.add_argument("--plan", required=True) + p.add_argument("--chat-model", required=True) + p.add_argument("--k", type=int, default=5) + p.add_argument("--rerank", default="none", choices=sorted(RERANKERS)) + p.add_argument("--fusion", default="rrf", choices=["rrf", "dbsf"]) + p.add_argument("--max-tokens", type=int, default=64) + p.add_argument("--max-context-chars", type=int, default=12000) + p.add_argument("--chat-concurrency", type=int, default=6) + p.set_defaults(func=lambda a: anyio.run(_run_generate, a)) + + p = sub.add_parser( + "ocr", help="OCR the scanned gold docs and index into an existing collection" + ) + _add_common(p) + p.add_argument("--plan", required=True) + p.add_argument("--engine", required=True, choices=sorted(OCR_ENGINES)) + p.add_argument("--corpus", default=str(DEFAULT_CORPUS)) + p.add_argument( + "--ocr-concurrency", type=int, default=6, help="sync-OCR concurrency (mistral)" + ) + p.add_argument( + "--poll-seconds", type=float, default=10.0, help="batch poll interval (surya)" + ) + p.add_argument( + "--reuse-job", + default=None, + help="surya: re-index an already-completed batch job_id (no re-submit, no GPU)", + ) + p.add_argument( + "--limit-docs", type=int, default=None, help="cap pending docs (smoke)" + ) + p.set_defaults(func=lambda a: anyio.run(_run_ocr, a)) + + args = ap.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/retrieval_eval/clients.py b/scripts/retrieval_eval/clients.py new file mode 100644 index 000000000..74ccf21db --- /dev/null +++ b/scripts/retrieval_eval/clients.py @@ -0,0 +1,225 @@ +"""Thin async clients for the eval harness. + +The embedding gateway and Ollama are reached directly over the tailnet with the +OpenAI-compatible wire format the repo's providers use — but *without* the M2M +OIDC handshake that ``GatewayProvider`` enforces, because the dev gateway is +open on the tailnet (no bearer). This keeps the harness runnable without gateway +client-credentials. Production still uses ``GatewayProvider`` + M2M; the dense +vectors are byte-identical because the ``/v1/embeddings`` contract is the same. + +Rate limits (mistral-embed: 12 req/s, 20M tok/min) are respected by bounding +concurrency with an ``anyio.CapacityLimiter`` supplied by the caller. +""" + +from __future__ import annotations + +import logging +import ssl +from dataclasses import dataclass +from typing import Any + +import anyio +import httpx +from fastembed import SparseTextEmbedding +from fastembed.rerank.cross_encoder import TextCrossEncoder + +logger = logging.getLogger(__name__) + +# Gateway/Ollama calls are flaky at scale (a sweep makes thousands): transient +# 502/503/504, read timeouts, and — under high embed concurrency on an aged +# connection pool — raw ``ssl.SSLError`` ("passed invalid argument") that httpx +# does NOT wrap as a TransportError. Retry all of them so one blip doesn't abort +# a whole task group; the harness runs unattended for many minutes. +_RETRY_STATUS = {500, 502, 503, 504, 429} +_RETRY_EXC = (httpx.TransportError, ssl.SSLError) +_MAX_ATTEMPTS = 5 + + +async def _post_json(client: httpx.AsyncClient, url: str, payload: dict) -> dict: + """POST JSON with bounded exponential backoff on transient failures.""" + last: Exception | None = None + for attempt in range(_MAX_ATTEMPTS): + try: + resp = await client.post(url, json=payload) + resp.raise_for_status() + return resp.json() + except (*_RETRY_EXC, httpx.HTTPStatusError) as exc: + # Only retry transient transport/SSL errors or retryable status codes. + if isinstance(exc, httpx.HTTPStatusError) and ( + exc.response.status_code not in _RETRY_STATUS + ): + raise + last = exc + if attempt < _MAX_ATTEMPTS - 1: + await anyio.sleep(1.5 * (2**attempt)) + logger.warning("retry %s %s (attempt %d)", url, exc, attempt + 1) + assert last is not None + raise last + + +@dataclass +class EmbedderSpec: + """A dense embedding model under test.""" + + name: str # human label, e.g. "mistral-embed" + kind: str # "gateway" | "ollama" + model: str # provider model id, e.g. "mistral/mistral-embed" + dim: int # vector dimension (drives Qdrant density / €/GiB) + + @property + def slug(self) -> str: + """Filesystem/collection-safe identifier.""" + return self.name.replace("/", "_").replace(":", "-").replace(".", "-") + + +class GatewayEmbedder: + """Dense embeddings via the gateway ``POST /v1/embeddings`` (OpenAI format).""" + + def __init__( + self, client: httpx.AsyncClient, model: str, limiter: anyio.CapacityLimiter + ): + self._client = client + self._model = model + self._limiter = limiter + + async def embed(self, texts: list[str]) -> tuple[list[list[float]], int]: + async with self._limiter: + data = await _post_json( + self._client, "/v1/embeddings", {"model": self._model, "input": texts} + ) + vectors = [row["embedding"] for row in data["data"]] + tokens = int(data.get("usage", {}).get("total_tokens", 0)) + return vectors, tokens + + +class OllamaEmbedder: + """Dense embeddings via a local Ollama ``POST /api/embed`` (self-hosted).""" + + def __init__( + self, client: httpx.AsyncClient, model: str, limiter: anyio.CapacityLimiter + ): + self._client = client + self._model = model + self._limiter = limiter + + async def embed(self, texts: list[str]) -> tuple[list[list[float]], int]: + async with self._limiter: + data = await _post_json( + self._client, "/api/embed", {"model": self._model, "input": texts} + ) + vectors = data["embeddings"] + tokens = int(data.get("prompt_eval_count", 0)) + return vectors, tokens + + +class GatewayReranker: + """Cross-encoder reranking via the gateway ``POST /v1/rerank`` (Cohere shape).""" + + def __init__( + self, client: httpx.AsyncClient, model: str, limiter: anyio.CapacityLimiter + ): + self._client = client + self._model = model + self._limiter = limiter + self.label = model + + async def rerank( + self, query: str, documents: list[str], top_n: int + ) -> list[tuple[int, float]]: + """Return ``(original_index, relevance_score)`` pairs, best first.""" + async with self._limiter: + data = await _post_json( + self._client, + "/v1/rerank", + { + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_n, + }, + ) + return [(int(r["index"]), float(r["relevance_score"])) for r in data["results"]] + + +class LocalCrossEncoderReranker: + """Cheap CPU cross-encoder via fastembed (the note-390460 control). + + Default model ``Xenova/ms-marco-MiniLM-L-6-v2`` — the exact reranker that + *hurt* finance tables, kept as an A/B control against the strong gateway one. + """ + + def __init__(self, model: str = "Xenova/ms-marco-MiniLM-L-6-v2"): + self._encoder = TextCrossEncoder(model_name=model) + self.label = model + + async def rerank( + self, query: str, documents: list[str], top_n: int + ) -> list[tuple[int, float]]: + scores = await anyio.to_thread.run_sync( # type: ignore[attr-defined] + lambda: list(self._encoder.rerank(query, documents)) + ) + ranked = sorted(enumerate(scores), key=lambda t: t[1], reverse=True) + return [(idx, float(score)) for idx, score in ranked[:top_n]] + + +class Bm25Encoder: + """FastEmbed BM25 sparse vectors — the SAME model production uses (Qdrant/bm25).""" + + def __init__(self, model_name: str = "Qdrant/bm25"): + self._model = SparseTextEmbedding(model_name=model_name) + + async def encode(self, texts: list[str]) -> list[dict[str, Any]]: + """Return ``[{"indices": list[int], "values": list[float]}, ...]``.""" + embeddings = await anyio.to_thread.run_sync( # type: ignore[attr-defined] + lambda: list(self._model.embed(texts)) + ) + return [ + {"indices": e.indices.tolist(), "values": e.values.tolist()} + for e in embeddings + ] + + +class GatewayChat: + """Text generation via the gateway ``POST /v1/chat/completions``.""" + + def __init__( + self, client: httpx.AsyncClient, model: str, limiter: anyio.CapacityLimiter + ): + self._client = client + self._model = model + self._limiter = limiter + + async def generate(self, prompt: str, *, max_tokens: int = 128) -> str: + async with self._limiter: + data = await _post_json( + self._client, + "/v1/chat/completions", + { + "model": self._model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + }, + ) + return data["choices"][0]["message"]["content"] + + +def make_embedder( + spec: EmbedderSpec, + *, + gateway: httpx.AsyncClient, + ollama: httpx.AsyncClient, + limiter: anyio.CapacityLimiter, +) -> GatewayEmbedder | OllamaEmbedder: + if spec.kind == "gateway": + return GatewayEmbedder(gateway, spec.model, limiter) + if spec.kind == "ollama": + return OllamaEmbedder(ollama, spec.model, limiter) + raise ValueError(f"unknown embedder kind: {spec.kind!r}") + + +def resolve_dim(spec: EmbedderSpec, vectors: list[list[Any]]) -> int: + """Trust the live vector length over the declared dim (defends against drift).""" + if vectors and vectors[0]: + return len(vectors[0]) + return spec.dim diff --git a/scripts/retrieval_eval/metrics.py b/scripts/retrieval_eval/metrics.py new file mode 100644 index 000000000..67d179352 --- /dev/null +++ b/scripts/retrieval_eval/metrics.py @@ -0,0 +1,266 @@ +"""OHR-Bench metric port + retrieval/generation scoring. + +The core string metrics (:func:`normalize_answer`, :func:`lcs_score`, +:func:`f1_score`, :func:`exact_match_score`) are copied faithfully from OHR-Bench +``src/metric/common.py`` so our numbers are directly comparable to the published +leaderboard. The retrieval page-gating mirrors ``src/tasks/retrieval.py`` (a hit = +a retrieved chunk from the gold doc whose page matches the evidence page), with +one documented adaptation: our packed chunks span a page RANGE +(``page_number..page_end``), so a chunk counts for every page it covers — this is +the same range-coverage rule used for the packed configs in note 390421/390460. + +Pure module (no I/O, no network) so it is unit-tested in +``tests/unit/test_retrieval_eval_metrics.py``. +""" + +from __future__ import annotations + +import re +import string +from collections import Counter, defaultdict +from collections.abc import Iterable + +# jieba is only needed to tokenize Chinese for F1; OHR-Bench uses it. It is not a +# repo dependency, so fall back to character tokens when a CJK string appears and +# jieba is unavailable (English — the bulk of the corpus — never needs it). +try: # pragma: no cover - exercised only when the optional dep is installed + import jieba # type: ignore + + _HAVE_JIEBA = True +except ImportError: # pragma: no cover + jieba = None # type: ignore + _HAVE_JIEBA = False + +_ARTICLES = re.compile(r"\b(a|an|the)\b") +_PUNC = set(string.punctuation) + + +# --------------------------------------------------------------------------- +# String metrics — verbatim from OHR-Bench src/metric/common.py +# --------------------------------------------------------------------------- +def normalize_answer(s: str) -> str: + """Lowercase, strip punctuation + articles, collapse whitespace.""" + s = s.lower() + s = "".join(ch for ch in s if ch not in _PUNC) + s = _ARTICLES.sub(" ", s) + return " ".join(s.split()) + + +def lcs_score(prediction: str, ground_truth: str) -> float: + """Word-level LCS recall of the gold span. ``A`` = gold, ``B`` = prediction. + + Returns ``len(LCS)/len(gold_tokens)``; ``0.5`` when the gold span is empty + (matches OHR-Bench). This is OHR-Bench's official *retrieval* metric. + """ + a = normalize_answer(ground_truth).split() + b = normalize_answer(prediction).split() + if len(a) == 0: + return 0.5 + dp = [[0] * (len(b) + 1) for _ in range(len(a) + 1)] + for i in range(1, len(a) + 1): + for j in range(1, len(b) + 1): + if a[i - 1] == b[j - 1]: + dp[i][j] = dp[i - 1][j - 1] + 1 + else: + dp[i][j] = max(dp[i - 1][j], dp[i][j - 1]) + return dp[len(a)][len(b)] / len(a) + + +def _has_cjk(s: str) -> bool: + return any("一" <= ch <= "鿿" for ch in s) + + +def _tokens_for_f1(text: str) -> list[str]: + normalized = normalize_answer(text) + if _has_cjk(text): + if _HAVE_JIEBA: + return jieba.lcut(normalized) # type: ignore[no-any-return] + return list(normalized.replace(" ", "")) + return normalized.split() + + +def f1_score(prediction: str, ground_truth: str) -> float: + """Token-level F1 (OHR-Bench official *generation* metric).""" + norm_pred = normalize_answer(prediction) + norm_gt = normalize_answer(ground_truth) + + # yes/no/noanswer guard, verbatim from OHR-Bench. + for special in ("yes", "no", "noanswer"): + if norm_pred == special and norm_pred != norm_gt: + return 0.0 + if norm_gt == special and norm_pred != norm_gt: + return 0.0 + + pred_tokens = _tokens_for_f1(prediction) + gt_tokens = _tokens_for_f1(ground_truth) + if not pred_tokens or not gt_tokens: + return 0.0 + common = Counter(pred_tokens) & Counter(gt_tokens) + num_same = sum(common.values()) + if num_same == 0: + return 0.0 + precision = num_same / len(pred_tokens) + recall = num_same / len(gt_tokens) + return (2 * precision * recall) / (precision + recall) + + +def exact_match_score(prediction: str, ground_truth: str) -> int: + """1 if normalized prediction equals normalized gold, else 0.""" + return 1 if normalize_answer(prediction) == normalize_answer(ground_truth) else 0 + + +# --------------------------------------------------------------------------- +# Retrieval scoring (mirrors src/tasks/retrieval.py, range-aware for packing) +# --------------------------------------------------------------------------- +def gold_pages(evidence_page_no: int | list[int]) -> set[int]: + if isinstance(evidence_page_no, list): + return {int(p) for p in evidence_page_no} + return {int(evidence_page_no)} + + +def covered_pages( + page_number: int | None, page_end: int | None, offset: int +) -> set[int]: + """Gold-page indices a (possibly packed) chunk covers, at a given offset. + + ``gold_page == astrolabe_page_number - offset`` (0- vs 1-index calibration). + A packed chunk spans ``page_number..page_end`` inclusive. + """ + if page_number is None: + return set() + end = page_end if page_end is not None else page_number + lo, hi = min(page_number, end), max(page_number, end) + return {p - offset for p in range(lo, hi + 1)} + + +def _as_text(context: str | list[str]) -> str: + return "\n".join(context) if isinstance(context, list) else context + + +def score_retrieval(plan: dict, results: list[dict], offset: int) -> dict[str, float]: + """Score one question's retrieval results at a given page offset. + + ``results`` items are ``{"doc_key", "page_number", "page_end", "text"}``. + Returns ``page_lcs / page_hit / doc_hit / doc_lcs`` (OHR-Bench metrics). + """ + gkey = plan["doc_name"] + gpages = gold_pages(plan["evidence_page_no"]) + gctx = _as_text(plan["evidence_context"]) + + doc_chunks = [r for r in results if r.get("doc_key") == gkey] + page_chunks = [ + r + for r in doc_chunks + if covered_pages(r.get("page_number"), r.get("page_end"), offset) & gpages + ] + page_text = "\n\n".join(r.get("text", "") for r in page_chunks) + doc_text = "\n\n".join(r.get("text", "") for r in doc_chunks) + + return { + "page_lcs": lcs_score(page_text, gctx) if page_chunks else 0.0, + "page_hit": 1.0 if page_chunks else 0.0, + "doc_hit": 1.0 if doc_chunks else 0.0, + "doc_lcs": lcs_score(doc_text, gctx) if doc_chunks else 0.0, + } + + +def calibrate_offset(plans: dict[str, dict], results: dict[str, list[dict]]) -> int: + """Pick the page offset (0 or 1) that yields the most page-hits. + + OHR-Bench gold pages may be 0- or 1-indexed relative to our page numbers; + calibrate over the answered set exactly as the external harness does. + """ + answered = [pid for pid in plans if pid in results] + best_off, best_hits = 0, -1 + for off in (0, 1): + hits = sum( + score_retrieval(plans[pid], results[pid], off)["page_hit"] + for pid in answered + ) + if hits > best_hits: + best_off, best_hits = off, hits + return best_off + + +def gold_rank(plan: dict, results: list[dict], offset: int) -> int | None: + """1-based rank of the first result on a gold page of the gold document. + + ``score_retrieval``'s metrics cannot see this. ``page_hit``/``doc_hit`` are + pure set membership, and ``page_lcs`` JOINS every page-matching chunk before + computing LCS — so a reranker that lifts the gold document from rank 8 to + rank 1 moves none of them. Measured on OHR-Bench: enabling reranking changed + the ``page_lcs`` of 12 of 280 queries (+0.011, p=0.034) while moving + Success@1 by +0.140. The page-gated metric is structurally blind to + ordering, which is most of what retrieval tuning actually changes. + + Returns ``None`` when the gold page never appears in the result list. + """ + gkey = plan["doc_name"] + gpages = gold_pages(plan["evidence_page_no"]) + for i, r in enumerate(results, start=1): + if r.get("doc_key") != gkey: + continue + if covered_pages(r.get("page_number"), r.get("page_end"), offset) & gpages: + return i + return None + + +def aggregate( + plans: dict[str, dict], + results: dict[str, list[dict]], + pids: Iterable[str], + offset: int, +) -> dict[str, float] | None: + """Mean of each retrieval metric over ``pids``; ``None`` if empty. + + Reports BOTH families, always, because they answer different questions and + have been observed to disagree in sign on the same runs: + + * ``page_lcs`` / ``page_hit`` / ``doc_hit`` / ``doc_lcs`` — OHR-Bench's + official page-gated metrics. Set-membership: did the gold evidence reach + the top-k at all. Leaderboard-comparable, and blind to ordering. + * ``success_at_1`` / ``success_at_3`` / ``mrr`` / ``mean_gold_rank`` — + rank-sensitive. Did the right thing come FIRST. + + Emitting only the first family is how a 24.5% relative improvement in + Success@1 gets reported as a marginal +0.011 and dismissed. + """ + pids = [p for p in pids if p in results] + n = len(pids) + if n == 0: + return None + acc: dict[str, float] = defaultdict(float) + ranks: list[int] = [] + for pid in pids: + for k, v in score_retrieval(plans[pid], results[pid], offset).items(): + acc[k] += v + r = gold_rank(plans[pid], results[pid], offset) + if r is not None: + ranks.append(r) + return { + "n": float(n), + **{k: acc[k] / n for k in ("page_lcs", "page_hit", "doc_hit", "doc_lcs")}, + # Denominator is n, not len(ranks): a query whose gold never surfaced is + # a failure at rank 1, not an absent sample. Dividing by len(ranks) would + # let a system that retrieves fewer golds report a better Success@1. + "success_at_1": sum(1 for r in ranks if r == 1) / n, + "success_at_3": sum(1 for r in ranks if r <= 3) / n, + "mrr": sum(1.0 / r for r in ranks) / n, + # Mean over FOUND golds only — the average depth you must read to when + # the answer is there at all. Undefined when nothing was found. + "mean_gold_rank": (sum(ranks) / len(ranks)) if ranks else 0.0, + "found": float(len(ranks)), + } + + +# --------------------------------------------------------------------------- +# Generation scoring +# --------------------------------------------------------------------------- +def score_generation(prediction: str, answers: str | list[str]) -> dict[str, float]: + """Best F1/EM of ``prediction`` over one or more gold answers.""" + golds = answers if isinstance(answers, list) else [answers] + golds = [str(g) for g in golds] or [""] + return { + "f1": max(f1_score(prediction, g) for g in golds), + "em": float(max(exact_match_score(prediction, g) for g in golds)), + } diff --git a/scripts/retrieval_eval/ocr.py b/scripts/retrieval_eval/ocr.py new file mode 100644 index 000000000..e14462022 --- /dev/null +++ b/scripts/retrieval_eval/ocr.py @@ -0,0 +1,171 @@ +"""OCR tier for the eval harness — fills the born-digital blind spot. + +Some OHR-Bench gold docs are scanned/image-only (no text layer), so the tier-1 +extractor (`pypdfium2_fast._extract`) returns empty and they never get indexed — +their questions auto-miss. This module OCRs those docs through the gateway and +reconstructs the same ``(full_text, page_boundaries)`` contract the chunkers +consume, so the OCR text flows through the identical chunk -> embed -> index path. + +Two engines (per the gateway): +- **surya/surya-ocr-2** — self-hosted, in-cluster. Submitted via the BATCH API + (`POST /v1/ocr/batch` + poll `GET /v1/ocr/batch/{job_id}`); the batch submit + triggers leaf.cloud GPU provisioning, so poll patiently to completion. +- **mistral/mistral-ocr-*** — cloud API. Uses the SYNC endpoint (`POST /v1/ocr`), + one document per call; no GPU spin-up. + +Both return per-page ``markdown``; we join with no separator (offsets stay exact, +matching `_extract`'s contract) and build ``page_boundaries``. +""" + +from __future__ import annotations + +import base64 +import logging +from typing import Any + +import anyio +import httpx + +from .clients import _post_json + +logger = logging.getLogger(__name__) + + +def pages_to_extract(pages_markdown: list[str]) -> tuple[str, dict[str, Any]]: + """Build the ``_extract`` contract from per-page OCR markdown. + + ``full_text`` = pages joined with NO separator so ``page_boundaries`` offsets + index it exactly (same invariant as `pypdfium2_fast._extract`). + """ + boundaries: list[dict[str, Any]] = [] + offset = 0 + for n, md in enumerate(pages_markdown, start=1): + boundaries.append( + {"page": n, "start_offset": offset, "end_offset": offset + len(md)} + ) + offset += len(md) + full_text = "".join(pages_markdown) + return full_text, {"page_count": len(pages_markdown), "page_boundaries": boundaries} + + +def _pages_markdown(pages: list[dict] | None) -> list[str]: + """Per-page markdown, ordered by the page ``index`` (gateway may reorder).""" + ordered = sorted(pages or [], key=lambda pg: pg.get("index", 0)) + return [pg.get("markdown", "") or "" for pg in ordered] + + +class MistralSyncOCR: + """Cloud OCR via the sync ``POST /v1/ocr`` — one PDF per call, no GPU.""" + + def __init__( + self, client: httpx.AsyncClient, model: str, limiter: anyio.CapacityLimiter + ): + self._client = client + self._model = model + self._limiter = limiter + self.label = model + + async def ocr_doc(self, pdf_bytes: bytes) -> list[str]: + b64 = base64.b64encode(pdf_bytes).decode() + async with self._limiter: + data = await _post_json( + self._client, + "/v1/ocr", + { + "model": self._model, + "document_b64": b64, + "mime_type": "application/pdf", + }, + ) + # Sync response: {pages: [{index, markdown, blocks?}]} (production shape). + pages = data.get("pages") + if pages is None and isinstance(data.get("result"), dict): + pages = data["result"].get("pages") + return _pages_markdown(pages) + + +class SuryaBatchOCR: + """Self-hosted OCR via the BATCH API — triggers the leaf.cloud GPU. + + Submits all docs in one batch job, then polls to completion. The GPU is off + at submit time; provisioning takes minutes, so use a generous poll budget. + """ + + def __init__( + self, + client: httpx.AsyncClient, + model: str, + *, + poll_seconds: float = 10.0, + max_wait_seconds: float = 3600.0, + ): + self._client = client + self._model = model + self._poll = poll_seconds + self._max_wait = max_wait_seconds + self.label = model + + async def ocr_docs( + self, docs: list[tuple[str, bytes]], *, reuse_job_id: str | None = None + ) -> dict[str, list[str]]: + """OCR many docs in one batch job. Returns ``{doc_key: [page_markdown]}``. + + ``docs`` = ``[(doc_key, pdf_bytes), ...]``. custom_id = doc_key. Completed + job results stay re-fetchable, so ``reuse_job_id`` skips submission (and + the GPU trigger) and re-polls an existing job — used to recover after a + downstream crash without paying for OCR again. + """ + if reuse_job_id: + job_id = reuse_job_id + logger.info("surya batch %s REUSED (no re-submit, no GPU)", job_id) + else: + documents = [ + { + "custom_id": doc_key, + "mime_type": "application/pdf", + "document_b64": base64.b64encode(pdf_bytes).decode(), + } + for doc_key, pdf_bytes in docs + ] + submit = await _post_json( + self._client, + "/v1/ocr/batch", + {"model": self._model, "documents": documents}, + ) + job_id = submit["job_id"] + logger.info( + "surya batch %s submitted (%d docs) — GPU provisioning, polling...", + job_id, + len(documents), + ) + waited = 0.0 + terminal = {"succeeded", "failed", "completed", "error"} + while waited < self._max_wait: + await anyio.sleep(self._poll) + waited += self._poll + resp = await self._client.get(f"/v1/ocr/batch/{job_id}", timeout=60) + resp.raise_for_status() + job = resp.json() + status = job.get("status") + if status in terminal: + logger.info( + "surya batch %s -> %s (%s/%s docs) after %.0fs", + job_id, + status, + job.get("succeeded"), + job.get("total"), + waited, + ) + out: dict[str, list[str]] = {} + for item in job.get("results") or []: + if item.get("error"): + logger.warning( + "surya ocr error %s: %s", + item.get("custom_id"), + item["error"], + ) + continue + out[item["custom_id"]] = _pages_markdown(item.get("pages")) + return out + logger.info("surya batch %s status=%s (%.0fs)", job_id, status, waited) + raise TimeoutError(f"surya batch {job_id} did not finish in {self._max_wait}s") diff --git a/scripts/retrieval_eval/pipeline.py b/scripts/retrieval_eval/pipeline.py new file mode 100644 index 000000000..80b615b4d --- /dev/null +++ b/scripts/retrieval_eval/pipeline.py @@ -0,0 +1,527 @@ +"""Indexing + hybrid retrieval for the eval harness. + +Mirrors the production hybrid stack so results transfer: extraction via +``pypdfium2_fast._extract`` (the born-digital tier-1 path), chunking via +``PageAwareChunker`` (the shipped default), a networked Qdrant collection with the +same named vectors (``dense`` + ``sparse``) and BM25 model (``Qdrant/bm25``), and +Qdrant-native RRF/DBSF fusion — the exact query shape of +``search/bm25_hybrid.py``. Directory-mode Qdrant is deliberately NOT supported +here: it throttles indexing throughput (note 390460). +""" + +from __future__ import annotations + +import logging +import uuid +from collections import defaultdict +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import anyio +from qdrant_client import AsyncQdrantClient, models + +from nextcloud_mcp_server.document_processors.pypdfium2_fast import _extract +from nextcloud_mcp_server.vector.document_chunker import ( + DocumentChunker, + PageAwareChunker, +) + +from .clients import Bm25Encoder, EmbedderSpec, GatewayEmbedder, OllamaEmbedder +from .ocr import pages_to_extract + +logger = logging.getLogger(__name__) + +# €/GiB-mo dense-RAM carry: bytes/point = dim*4 (fp32) + ~2048 (sparse+payload+index +# overhead, calibrated so 1024-d == the note-390397 6,144 B/point), at €1.75/GB-mo. +_OVERHEAD_BYTES = 2048 +_EUR_PER_GB_MO = 1.75 + + +def eur_per_gib_mo(chunks_per_mb: float, dim: int) -> float: + bytes_per_point = dim * 4 + _OVERHEAD_BYTES + return chunks_per_mb * bytes_per_point * 1000 * _EUR_PER_GB_MO / 1e9 + + +@dataclass +class Chunk: + text: str + page_number: int | None + page_end: int | None + + +@dataclass +class DomainDensity: + chunks: int = 0 + source_bytes: int = 0 + + @property + def chunks_per_mb(self) -> float: + return ( + self.chunks / (self.source_bytes / 1_000_000) if self.source_bytes else 0.0 + ) + + +@dataclass +class IndexStats: + dim: int = 0 + docs: int = 0 + total_chunks: int = 0 + embed_tokens: int = 0 + # Docs that yielded no born-digital text (scanned / image-only PDFs) — NOT + # indexed, so their questions auto-miss. Surfaced (not silently dropped) so a + # coverage gap is visible: OHR-Bench includes scanned docs the tier-1 + # extractor can't read; that is the pipeline's real limitation, not a bug. + empty_docs: int = 0 + by_domain: dict[str, DomainDensity] = field( + default_factory=lambda: defaultdict(DomainDensity) + ) + + +def iter_corpus( + corpus: Path, + domains: list[str] | None, + limit: int | None, + doc_allow: set[str] | None = None, +) -> Iterator[tuple[str, str, Path]]: + """Yield ``(doc_key, domain, pdf_path)`` for ``/.pdf`` PDFs. + + ``doc_key`` is ``"/"`` (no extension) — the exact form of + OHR-Bench ``qas_v2.json`` ``doc_name``, so gold and indexed keys compare + directly with no domain/extension juggling. + + ``doc_allow`` restricts to that set of ``doc_key``s — used by ``--from-plan`` + to index only the gold docs referenced by the sampled questions (the + gold-doc corpus of notes 390421/390460), which is both faithful to prior + methodology and far cheaper than the full 1,261-doc corpus. + """ + dirs = domains or sorted(p.name for p in corpus.iterdir() if p.is_dir()) + for domain in dirs: + domain_dir = corpus / domain + if not domain_dir.is_dir(): + continue + pdfs = sorted(domain_dir.rglob("*.pdf")) + if limit is not None: + pdfs = pdfs[:limit] + for pdf in pdfs: + doc_key = pdf.relative_to(corpus).with_suffix("").as_posix() + if doc_allow is not None and doc_key not in doc_allow: + continue + yield doc_key, domain, pdf + + +async def extract_and_chunk( + pdf: Path, *, chunk_size: int, pack_pages: bool +) -> list[Chunk]: + """Extract born-digital text and chunk it with the page-aware chunker.""" + pdf_bytes = pdf.read_bytes() + text, metadata = await anyio.to_thread.run_sync(_extract, pdf_bytes) # type: ignore[attr-defined] + if not text: + return [] + boundaries = metadata.get("page_boundaries") or [] + if boundaries: + chunker = PageAwareChunker(chunk_size=chunk_size, pack_pages=pack_pages) + raw = await chunker.chunk_text(text, boundaries) + else: + raw = await DocumentChunker(chunk_size=chunk_size).chunk_text(text) + return [Chunk(c.text, c.page_number, c.page_end) for c in raw if c.text.strip()] + + +async def records_from_ocr( + items: list[tuple[str, str, int, list[str]]], + *, + chunk_size: int, + pack_pages: bool, +) -> list[_DocRecord]: + """Chunk OCR'd docs into ``_DocRecord``s via the same page-aware chunker. + + ``items`` = ``[(doc_key, domain, source_bytes, pages_markdown), ...]``. The OCR + per-page markdown is reconstructed into the ``(text, page_boundaries)`` contract + (`ocr.pages_to_extract`) so OCR text flows through the identical chunk path as + born-digital text — same page-pack@N behavior, same page-range citation. + """ + records: list[_DocRecord] = [] + for doc_key, domain, source_bytes, pages_md in items: + if not any(p.strip() for p in pages_md): + logger.warning("ocr produced no text: %s", doc_key) + continue + text, meta = pages_to_extract(pages_md) + boundaries = meta["page_boundaries"] + if boundaries: + raw = await PageAwareChunker( + chunk_size=chunk_size, pack_pages=pack_pages + ).chunk_text(text, boundaries) + else: + raw = await DocumentChunker(chunk_size=chunk_size).chunk_text(text) + chunks = [ + Chunk(c.text, c.page_number, c.page_end) for c in raw if c.text.strip() + ] + if chunks: + records.append(_DocRecord(doc_key, domain, source_bytes, chunks)) + return records + + +async def ensure_collection( + qdrant: AsyncQdrantClient, + name: str, + dim: int, + *, + recreate: bool, + quantization: str = "none", +) -> bool: + """Create the collection (dense + sparse). Returns True if (re)created. + + ``quantization`` in {none, int8, binary}: when set, the dense vector is + quantized (int8 scalar or 1-bit binary) with the quantized form pinned in RAM + (``always_ram``) and the original fp32 moved on-disk (``on_disk=True``) for + rescoring. This is the RAM-footprint lever the cost model doesn't yet price. + """ + exists = await qdrant.collection_exists(name) + if exists and not recreate: + return False + if exists: + await qdrant.delete_collection(name) + on_disk = quantization != "none" + await qdrant.create_collection( + collection_name=name, + vectors_config={ + "dense": models.VectorParams( + size=dim, distance=models.Distance.COSINE, on_disk=on_disk + ) + }, + sparse_vectors_config={"sparse": models.SparseVectorParams()}, + quantization_config=_quantization_config(quantization), + ) + return True + + +def _quantization_config(mode: str): + """Qdrant quantization_config for a mode; None for 'none'.""" + if mode == "int8": + return models.ScalarQuantization( + scalar=models.ScalarQuantizationConfig( + type=models.ScalarType.INT8, always_ram=True + ) + ) + if mode == "binary": + return models.BinaryQuantization( + binary=models.BinaryQuantizationConfig(always_ram=True) + ) + if mode == "none": + return None + raise ValueError(f"unknown quantization mode: {mode!r}") + + +async def clone_collection( + qdrant: AsyncQdrantClient, + source: str, + target: str, + dim: int, + *, + quantization: str, + batch: int = 512, +) -> int: + """Copy source → target applying ``quantization`` — reuses vectors, no re-embed. + + Streams points (dense + sparse + payload) from an existing collection into a + fresh one whose only difference is the quantization config. Lets us A/B RAM + modes on identical embeddings without paying the gateway again. + """ + await ensure_collection( + qdrant, target, dim, recreate=True, quantization=quantization + ) + copied = 0 + offset = None + while True: + points, offset = await qdrant.scroll( + collection_name=source, + limit=batch, + offset=offset, + with_payload=True, + with_vectors=True, + ) + if not points: + break + await qdrant.upsert( + collection_name=target, + points=[ + models.PointStruct( + id=p.id, + vector=p.vector, # type: ignore[arg-type] # scroll w/ vectors -> dict + payload=p.payload, + ) + for p in points + if p.vector is not None + ], + ) + copied += len(points) + if offset is None: + break + return copied + + +@dataclass +class _DocRecord: + doc_key: str + domain: str + source_bytes: int + chunks: list[Chunk] + start: int = 0 # offset into the flat chunk arrays + + +async def index_corpus( + *, + qdrant: AsyncQdrantClient, + collection: str, + embedder: GatewayEmbedder | OllamaEmbedder, + bm25: Bm25Encoder, + spec: EmbedderSpec, + corpus: Path, + domains: list[str] | None, + limit: int | None, + chunk_size: int, + pack_pages: bool, + doc_allow: set[str] | None = None, + extract_concurrency: int = 1, + embed_batch: int = 8, + sparse_batch: int = 512, + upsert_concurrency: int = 8, + upsert_batch: int = 256, +) -> IndexStats: + """Extract → chunk → embed (dense+sparse) → upsert. + + Embedding is the bottleneck (the gateway serves each request's texts serially + at ~1.6 s/text but parallelizes across requests — measured). So rather than + embed each doc's chunks sequentially (which makes the slowest single doc + dominate wall-clock), ALL chunks across ALL docs are flattened and embedded + through one global concurrent pool bounded by the embedder's own + ``CapacityLimiter``. Throughput then scales with that limiter (~29 chunks/s at + 48-way) regardless of per-doc size. + """ + stats = IndexStats(dim=spec.dim) + + # --- Phase 1: extract + chunk all docs --- + # pypdfium2 is NOT thread-safe: concurrent ``_extract`` calls corrupt the + # native heap (``free(): invalid size`` -> SIGABRT). So extraction defaults to + # serial (``extract_concurrency=1``). This costs nothing — extraction is fast + # and the slow phase (embedding) is parallelized separately below. + records: list[_DocRecord] = [] + ex_limiter = anyio.CapacityLimiter(extract_concurrency) + ex_lock = anyio.Lock() + + async def do_extract(doc_key: str, domain: str, pdf: Path) -> None: + async with ex_limiter: + try: + chunks = await extract_and_chunk( + pdf, chunk_size=chunk_size, pack_pages=pack_pages + ) + except Exception as exc: # noqa: BLE001 - a bad PDF must not abort the run + logger.warning("skip %s: %s", doc_key, exc) + return + if not chunks: + # No text layer (scanned/image PDF) — the tier-1 extractor's blind + # spot. Record it so coverage is auditable rather than silent. + async with ex_lock: + stats.empty_docs += 1 + logger.debug("no text layer (scanned?): %s", doc_key) + return + async with ex_lock: + records.append(_DocRecord(doc_key, domain, pdf.stat().st_size, chunks)) + + async with anyio.create_task_group() as tg: + for doc_key, domain, pdf in iter_corpus(corpus, domains, limit, doc_allow): + tg.start_soon(do_extract, doc_key, domain, pdf) + + logger.info( + "extracted %d docs (%d had no text layer); embedding (dim=%d)...", + len(records), + stats.empty_docs, + spec.dim, + ) + await embed_and_upsert( + qdrant=qdrant, + collection=collection, + embedder=embedder, + bm25=bm25, + records=records, + stats=stats, + embed_batch=embed_batch, + sparse_batch=sparse_batch, + upsert_concurrency=upsert_concurrency, + upsert_batch=upsert_batch, + ) + return stats + + +async def embed_and_upsert( + *, + qdrant: AsyncQdrantClient, + collection: str, + embedder: GatewayEmbedder | OllamaEmbedder, + bm25: Bm25Encoder, + records: list[_DocRecord], + stats: IndexStats, + embed_batch: int = 8, + sparse_batch: int = 512, + upsert_concurrency: int = 8, + upsert_batch: int = 256, +) -> None: + """Embed (dense+sparse) + upsert pre-chunked ``records`` into ``collection``. + + Shared by born-digital (`index_corpus`) and OCR (`index_ocr_records`) paths. + All chunks across all docs are flattened and dense-embedded through one global + concurrent pool (bounded by the embedder's ``CapacityLimiter``) so throughput + scales with concurrency, not per-doc size. Does NOT create the collection. + """ + flat_texts: list[str] = [] + for rec in records: + rec.start = len(flat_texts) + flat_texts.extend(c.text for c in rec.chunks) + if not flat_texts: + return + + # Dense: global concurrent pool. Untyped list so the None placeholders (filled + # below) don't fight PointStruct's vector type; every slot is set before upsert. + dense: list = [None] * len(flat_texts) # type: ignore[type-arg] + tok_lock = anyio.Lock() + done = [0] + + async def embed_span(start: int) -> None: + vecs, tok = await embedder.embed(flat_texts[start : start + embed_batch]) + for off, v in enumerate(vecs): + dense[start + off] = v + async with tok_lock: + stats.embed_tokens += tok + done[0] += len(vecs) + if done[0] % 2000 < len(vecs): + logger.info("embedded %d/%d chunks", done[0], len(flat_texts)) + + async with anyio.create_task_group() as tg: + for start in range(0, len(flat_texts), embed_batch): + tg.start_soon(embed_span, start) + + # Sparse BM25 (CPU, batched in-thread). + sparse: list[dict] = [] + for i in range(0, len(flat_texts), sparse_batch): + sparse.extend(await bm25.encode(flat_texts[i : i + sparse_batch])) + + up_limiter = anyio.CapacityLimiter(upsert_concurrency) + st_lock = anyio.Lock() + + def _vector(idx: int) -> dict[str, Any]: + return { + "dense": dense[idx], + "sparse": models.SparseVector( + indices=sparse[idx]["indices"], values=sparse[idx]["values"] + ), + } + + async def upsert_doc(rec: _DocRecord) -> None: + points = [ + models.PointStruct( + id=str(uuid.uuid4()), + vector=_vector(rec.start + j), + payload={ + "doc_key": rec.doc_key, + "domain": rec.domain, + "page_number": rec.chunks[j].page_number, + "page_end": rec.chunks[j].page_end, + "text": rec.chunks[j].text, + }, + ) + for j in range(len(rec.chunks)) + ] + async with up_limiter: + for i in range(0, len(points), upsert_batch): + await qdrant.upsert( + collection_name=collection, points=points[i : i + upsert_batch] + ) + async with st_lock: + stats.docs += 1 + stats.total_chunks += len(rec.chunks) + dd = stats.by_domain[rec.domain] + dd.chunks += len(rec.chunks) + dd.source_bytes += rec.source_bytes + + async with anyio.create_task_group() as tg: + for rec in records: + tg.start_soon(upsert_doc, rec) + + +async def retrieve( + *, + qdrant: AsyncQdrantClient, + collection: str, + query: str, + embedder: GatewayEmbedder | OllamaEmbedder, + bm25: Bm25Encoder, + k: int, + fusion: str, + fetch: int | None = None, + rescore: bool | None = None, + oversampling: float | None = None, +) -> list[dict]: + """Hybrid RRF/DBSF retrieval, returning up to ``fetch`` candidates. + + ``fetch`` defaults to ``2*k`` (the mild over-fetch the production tool uses for + dedup). A reranker passes a larger ``fetch`` (e.g. 50) so it has a broad + candidate pool to reorder into the top-``k`` — where a strong reranker earns + its keep. The dense/sparse prefetch and fusion all use ``fetch``. + + ``rescore``/``oversampling`` are Qdrant quantization search params applied to + the DENSE prefetch: rescore re-ranks the quantized candidates with the on-disk + fp32 vectors (recall recovery), oversampling widens the quantized shortlist + first (matters for binary). No-op on unquantized collections. + """ + limit = fetch if fetch is not None else k * 2 + dense, _ = await embedder.embed([query]) + sparse = (await bm25.encode([query]))[0] + fusion_enum = models.Fusion.RRF if fusion == "rrf" else models.Fusion.DBSF + dense_params = None + if rescore is not None or oversampling is not None: + dense_params = models.SearchParams( + quantization=models.QuantizationSearchParams( + rescore=rescore, oversampling=oversampling + ) + ) + resp = await qdrant.query_points( + collection_name=collection, + prefetch=[ + models.Prefetch( + query=dense[0], using="dense", limit=limit, params=dense_params + ), + models.Prefetch( + query=models.SparseVector( + indices=sparse["indices"], values=sparse["values"] + ), + using="sparse", + limit=limit, + ), + ], + query=models.FusionQuery(fusion=fusion_enum), + limit=limit, + with_payload=True, + with_vectors=False, + ) + out: list[dict] = [] + for point in resp.points: + p = point.payload or {} + out.append( + { + "doc_key": p.get("doc_key"), + "domain": p.get("domain"), + "page_number": p.get("page_number"), + "page_end": p.get("page_end"), + "text": p.get("text", ""), + "score": point.score, + } + ) + return out + + +async def stream_corpus_paths( + corpus: Path, domains: list[str] | None, limit: int | None +) -> AsyncIterator[tuple[str, str, Path]]: # pragma: no cover - convenience only + for item in iter_corpus(corpus, domains, limit): + yield item diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 80a4d8f9a..611ee5f4c 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -357,6 +357,26 @@ def test_direct_url_warns_about_the_gateway_namespaced_default_model(self, caplo assert "SEARCH_RERANK_MODEL" in caplog.text assert "BAAI/bge-reranker-v2-m3" in caplog.text + @patch.dict( + os.environ, + { + "SEARCH_RERANK_ENABLED": "true", + "EMBEDDING_GATEWAY_URL": "https://gw.example", + "SEARCH_RERANK_URL": "https://gw.example/v1/rerank", + }, + clear=True, + ) + def test_rerank_url_pointing_at_the_gateway_is_quiet(self, caplog): + """Pinning SEARCH_RERANK_URL to the gateway's OWN /v1/rerank is a valid + configuration -- the routing prefix is correct there. Warning on it fires + on a working setup, which is how operators learn to ignore warnings.""" + _reload_config() + with caplog.at_level(logging.WARNING, logger="nextcloud_mcp_server.config"): + settings = get_settings() + + assert settings.search_rerank_model == "local/BAAI/bge-reranker-v2-m3" + assert "SEARCH_RERANK_MODEL" not in caplog.text + @patch.dict( os.environ, { diff --git a/tests/unit/test_retrieval_eval_metrics.py b/tests/unit/test_retrieval_eval_metrics.py new file mode 100644 index 000000000..5055e8ba6 --- /dev/null +++ b/tests/unit/test_retrieval_eval_metrics.py @@ -0,0 +1,242 @@ +"""Fidelity tests for the OHR-Bench metric port in scripts.retrieval_eval. + +These pin the ported ``lcs_score`` / ``f1_score`` / page-gating to hand-worked +values so the committed harness provably matches OHR-Bench's official metrics +(the reason note 390434 flagged re-scoring on LCS in the first place). +""" + +import math + +import pytest + +from scripts.retrieval_eval import metrics + +pytestmark = pytest.mark.unit + + +def test_normalize_answer_strips_punct_articles_case(): + assert metrics.normalize_answer("The, QUICK fox!") == "quick fox" + + +def test_lcs_score_full_partial_empty(): + # gold fully covered by prediction -> recall 1.0 + assert metrics.lcs_score("the quick brown fox", "quick fox") == pytest.approx(1.0) + # only half the gold tokens present -> 0.5 + assert metrics.lcs_score("quick", "quick fox") == pytest.approx(0.5) + # empty gold -> 0.5 sentinel (OHR-Bench convention) + assert metrics.lcs_score("anything", "") == pytest.approx(0.5) + + +def test_lcs_is_subsequence_not_substring(): + # tokens must appear in order but not contiguously + assert metrics.lcs_score("a x b y c", "a b c") == pytest.approx(1.0) + + +def test_f1_score_identical_partial_none_and_yesno_guard(): + assert metrics.f1_score("total 842", "total 842") == pytest.approx(1.0) + assert metrics.f1_score("elephant", "842") == pytest.approx(0.0) + # pred=[cat,sat] gt=[cat]: p=1/2, r=1/1 -> F1=0.6667 + assert metrics.f1_score("the cat sat", "cat") == pytest.approx(2 / 3) + # yes/no/noanswer guard + assert metrics.f1_score("yes", "no") == 0.0 + + +def test_exact_match_normalizes(): + assert metrics.exact_match_score("The 842.", "842") == 1 + assert metrics.exact_match_score("843", "842") == 0 + + +def test_covered_pages_range_and_offset(): + assert metrics.covered_pages(4, 6, 0) == {4, 5, 6} + assert metrics.covered_pages(4, 6, 1) == {3, 4, 5} + assert metrics.covered_pages(7, None, 0) == {7} + assert metrics.covered_pages(None, None, 0) == set() + + +def _plan(**over): + base = { + "ID": "q1", + "doc_name": "finance/X", + "doc_type": "finance", + "evidence_page_no": 5, + "evidence_source": "table", + "evidence_context": "total 842", + } + base.update(over) + return base + + +def test_score_retrieval_packed_chunk_hits_page_range(): + plan = _plan() + results = [ + { + "doc_key": "finance/X", + "page_number": 4, + "page_end": 6, + "text": "row total 842 row", + }, + {"doc_key": "other/Y", "page_number": 5, "page_end": 5, "text": "noise"}, + ] + s = metrics.score_retrieval(plan, results, offset=0) + assert s["doc_hit"] == 1.0 + assert s["page_hit"] == 1.0 + assert s["page_lcs"] == pytest.approx( + 1.0 + ) # gold span fully present on the covered page + + +def test_score_retrieval_wrong_doc_is_miss(): + plan = _plan() + results = [ + {"doc_key": "finance/Z", "page_number": 5, "page_end": 5, "text": "total 842"} + ] + s = metrics.score_retrieval(plan, results, offset=0) + assert s["doc_hit"] == 0.0 + assert s["page_hit"] == 0.0 + assert s["page_lcs"] == 0.0 + + +def test_score_retrieval_doc_hit_without_page_hit(): + plan = _plan(evidence_page_no=99) + results = [ + {"doc_key": "finance/X", "page_number": 5, "page_end": 5, "text": "total 842"} + ] + s = metrics.score_retrieval(plan, results, offset=0) + assert s["doc_hit"] == 1.0 + assert s["page_hit"] == 0.0 + + +def test_calibrate_offset_picks_best(): + # gold page 5, retrieved page_number 6 -> only offset=1 makes it a hit + plans = {"q1": _plan(evidence_page_no=5)} + results = { + "q1": [ + { + "doc_key": "finance/X", + "page_number": 6, + "page_end": 6, + "text": "total 842", + } + ] + } + assert metrics.calibrate_offset(plans, results) == 1 + + +def test_aggregate_means_and_empty(): + plans = {"q1": _plan(), "q2": _plan(ID="q2")} + results = { + "q1": [ + { + "doc_key": "finance/X", + "page_number": 5, + "page_end": 5, + "text": "total 842", + } + ], + "q2": [{"doc_key": "nope/W", "page_number": 1, "page_end": 1, "text": "x"}], + } + m = metrics.aggregate(plans, results, ["q1", "q2"], offset=0) + assert m is not None + assert m["n"] == 2.0 + assert m["doc_hit"] == pytest.approx(0.5) + assert metrics.aggregate(plans, results, [], offset=0) is None + + +def test_score_generation_best_over_answers(): + s = metrics.score_generation("842", ["841", "842"]) + assert s["f1"] == pytest.approx(1.0) + assert s["em"] == 1.0 + assert not math.isnan(s["f1"]) + + +# --------------------------------------------------------------------------- +# Rank-sensitive metrics +# +# The page-gated family cannot see ordering: page_hit/doc_hit are set +# membership and page_lcs joins every matching chunk before scoring. These +# tests pin the family that CAN, because a reranker's entire job is invisible +# to the other one. +# --------------------------------------------------------------------------- +def _hit(page=5, text="total 842", doc="finance/X"): + return {"doc_key": doc, "page_number": page, "page_end": page, "text": text} + + +def _miss(doc="other/Y", page=1): + return {"doc_key": doc, "page_number": page, "page_end": page, "text": "x"} + + +def test_gold_rank_is_one_based_and_finds_first_gold_page(): + plan = _plan() + assert metrics.gold_rank(plan, [_hit()], offset=0) == 1 + assert metrics.gold_rank(plan, [_miss(), _hit()], offset=0) == 2 + assert metrics.gold_rank(plan, [_miss(), _miss("z/Z"), _hit()], offset=0) == 3 + + +def test_gold_rank_is_none_when_gold_never_appears(): + assert metrics.gold_rank(_plan(), [_miss(), _miss("z/Z")], offset=0) is None + + +def test_rank_metrics_see_reordering_that_page_lcs_cannot(): + """The whole reason this family exists. + + Same two result SETS, different ORDER. page_lcs/page_hit are identical + because they only ask whether the gold chunk is present; Success@1 and MRR + separate them. Measured on OHR-Bench this gap was 14x: +0.011 page_lcs + versus +0.140 Success@1 on the same runs. + """ + plans = {"q1": _plan()} + gold_first = {"q1": [_hit(), _miss()]} + gold_second = {"q1": [_miss(), _hit()]} + + a = metrics.aggregate(plans, gold_first, ["q1"], offset=0) + b = metrics.aggregate(plans, gold_second, ["q1"], offset=0) + assert a is not None and b is not None + + # Blind to the reorder... + assert a["page_lcs"] == pytest.approx(b["page_lcs"]) + assert a["page_hit"] == pytest.approx(b["page_hit"]) + # ...sensitive to it. + assert a["success_at_1"] == 1.0 + assert b["success_at_1"] == 0.0 + assert a["mrr"] == pytest.approx(1.0) + assert b["mrr"] == pytest.approx(0.5) + + +def test_success_denominator_is_all_queries_not_just_found_ones(): + """A query whose gold never surfaced is a failure at rank 1, not an absent + sample. Dividing by the found count would let a system that retrieves FEWER + golds report a better Success@1.""" + plans = {"q1": _plan(), "q2": _plan(ID="q2")} + results = {"q1": [_hit()], "q2": [_miss()]} + + m = metrics.aggregate(plans, results, ["q1", "q2"], offset=0) + assert m is not None + assert m["found"] == 1.0 + assert m["success_at_1"] == pytest.approx(0.5) # not 1.0 + assert m["mrr"] == pytest.approx(0.5) + + +def test_mean_gold_rank_averages_found_only_and_is_zero_when_none_found(): + plans = {"q1": _plan(), "q2": _plan(ID="q2")} + results = {"q1": [_miss(), _miss("z/Z"), _hit()], "q2": [_miss()]} + + m = metrics.aggregate(plans, results, ["q1", "q2"], offset=0) + assert m is not None + assert m["mean_gold_rank"] == pytest.approx(3.0) # q2 excluded, not scored 0 + + none_found = metrics.aggregate(plans, {"q1": [_miss()]}, ["q1"], offset=0) + assert none_found is not None + assert none_found["mean_gold_rank"] == 0.0 + assert none_found["found"] == 0.0 + + +def test_success_at_3_is_inclusive_of_rank_three(): + plans = {"q1": _plan()} + at3 = {"q1": [_miss(), _miss("z/Z"), _hit()]} + at4 = {"q1": [_miss(), _miss("z/Z"), _miss("w/W"), _hit()]} + + m3 = metrics.aggregate(plans, at3, ["q1"], offset=0) + m4 = metrics.aggregate(plans, at4, ["q1"], offset=0) + assert m3 is not None and m4 is not None + assert m3["success_at_3"] == 1.0 + assert m4["success_at_3"] == 0.0 diff --git a/third_party/astrolabe b/third_party/astrolabe index ccd2d180e..318fba63c 160000 --- a/third_party/astrolabe +++ b/third_party/astrolabe @@ -1 +1 @@ -Subproject commit ccd2d180e88d5ba71948c4617159cf654481a9c4 +Subproject commit 318fba63cdd85e61b8bee8b76877e6c52b0b7fe0