From fc0698c2426c373e32c361a816523351a1e3fcb7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Tue, 6 Oct 2026 14:36:02 -0400 Subject: [PATCH] Add baseline harness: ingest, metrics, ablation arms, external counterpart, audit - bench/ingest_baseline.py drives the unmodified product IngestionService over all 50,789 corpus files in sharded SQLite DBs, then merges them. Thread and memory ceilings are set explicitly: ONNX Runtime defaulted to every core (95 threads/worker, load 490) and its CPU arena kept freed pages mapped, which OOM-killed workers. - bench/metrics.py: recall@k, MRR, nDCG@10, Wilson lower bound, chance baseline. No model, no product import, so the scorer cannot grade itself. - bench/systems.py + bench/run.py: the five ablation arms over the merged DB, latency percentiles, scale curve, and a top-10 equivalence check against the live RetrievalService. - bench/counterpart.py + bench/run_external.py: in-repo Mem0-style flat store and Zep-style temporal store as the external reference point. - bench/audit.py: enforces single-attempt scoring, frozen n, a chance row beside every score, and producer != verifier. It already flags the product's retry-until-target loop at backend/src/ely_eye/evals.py:2351. - backend/pyproject.toml: rapidocr-onnxruntime>=1.4.4 was unsatisfiable (PyPI maxes at 1.2.3); the successor is the renamed rapidocr package. --- .gitignore | 2 + CHECKLIST.md | 25 +- backend/pyproject.toml | 2 +- backend/src/ely_eye/py.typed | 0 bench/README.md | 62 +++++ bench/audit.py | 278 ++++++++++++++++++++++ bench/build_code.py | 2 +- bench/build_longmemeval.py | 2 +- bench/build_mmlongbench.py | 5 +- bench/build_screenspot.py | 2 +- bench/build_videomme.py | 2 +- bench/build_visrag.py | 2 +- bench/configs/ablation.json | 5 + bench/counterpart.py | 321 +++++++++++++++++++++++++ bench/fetch_corpus.py | 4 +- bench/fetch_models.py | 45 ++++ bench/ingest_baseline.py | 404 ++++++++++++++++++++++++++++++++ bench/metrics.py | 150 ++++++++++++ bench/run.py | 359 ++++++++++++++++++++++++++++ bench/run_external.py | 161 +++++++++++++ bench/systems.py | 294 +++++++++++++++++++++++ bench/tests/test_audit.py | 102 ++++++++ bench/tests/test_counterpart.py | 125 ++++++++++ bench/tests/test_metrics.py | 103 ++++++++ bench/tests/test_systems.py | 139 +++++++++++ pyproject.toml | 15 +- pyrightconfig.json | 2 +- 27 files changed, 2589 insertions(+), 24 deletions(-) create mode 100644 backend/src/ely_eye/py.typed create mode 100644 bench/README.md create mode 100644 bench/audit.py create mode 100644 bench/configs/ablation.json create mode 100644 bench/counterpart.py create mode 100644 bench/fetch_models.py create mode 100644 bench/ingest_baseline.py create mode 100644 bench/metrics.py create mode 100644 bench/run.py create mode 100644 bench/run_external.py create mode 100644 bench/systems.py create mode 100644 bench/tests/test_audit.py create mode 100644 bench/tests/test_counterpart.py create mode 100644 bench/tests/test_metrics.py create mode 100644 bench/tests/test_systems.py diff --git a/.gitignore b/.gitignore index 1e2deee..e8ad487 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,5 @@ designreference.html bench/corpus/* !bench/corpus/manifest.jsonl .pi/ +bench/runs/ +bench/models/ diff --git a/CHECKLIST.md b/CHECKLIST.md index 732ee7f..b680b61 100644 --- a/CHECKLIST.md +++ b/CHECKLIST.md @@ -26,22 +26,23 @@ - [x] b-2.4 陷阱类别齐：知识更新 25 / 时间推理 25 / 多会话 25（LongMemEval 原生类别）+ 无答案 2380（4 种构造：cross-holder、held-out、native unanswerable、distractor） - 验收：`python -m bench.validate_sets`（清单原名 dataset_check/manifest_check，实作合并为一个）报 `OK: 50789 sources, 2928 answerable queries, 2380 unanswerable` 退出 0；它同时校验 gold 存在性、gold holder==query holder、`observed_at <= as_of` 时间可达性、以及每个 set 文件的 SHA256 是否仍等于冻结的 manifest（重跑两次 SHA 相同由此保证） - [ ] **b-3** 检索 benchmark：五路消融 + 延迟规模 - - [ ] b-3.1 指标实现：recall@k、MRR、nDCG@10，纯 Python，无模型 - - [ ] b-3.2 消融五路：BM25-only / dense-only / graph-only / no-rerank / fusion+rerank - - [ ] b-3.3 延迟规模曲线：1k / 10k / 100k atom，p50 与 p95 + - [x] b-3.1 指标实现：recall@k、MRR、nDCG@10，纯 Python，无模型（`bench/metrics.py` + `bench/tests/test_metrics.py` 11 条：Wilson 90/100→0.8257、10/10→0.7225、0/10→0.0 手算对齐；nDCG@10 `1.5/1.6309297535714574`；chance 公式经测试修正为 `min(1.0, k/N)`） + - [~] b-3.2 消融五路已写（`bench/systems.py`：`bm25_only / dense_only / graph_only / no_rerank / fusion_rerank`，融合权重照抄产品 `0.55/0.35/0.10`，`graph_only` 如实返回常量序因为产品从不填 `atom.relations`），14 条测试过；数字等全量 ingest + - [~] b-3.3 延迟规模曲线已写（`bench/run.py scale_curve`，p50/p95）；档位按语料实际取 `500 / 2_000 / 10_000 / 50_000` 个 source（清单原写 1k/10k/100k atom，语料冻结为 50,789 source，100k 档不存在） - [ ] b-3.4 **BASELINE 报告**：在改动前的代码上跑，冻结提交 - - 验收：`python -m bench.run --config bench/configs/ablation.json` 出表；每格带 n 与 Wilson 95% 下界；`bench/reports/baseline_*.json` 入库 + - 验收：`python -m bench.run --config bench/configs/ablation.json` 出表（`--config` 已接，配置优先于默认、命令行仍可覆盖）；每格带 n 与 Wilson 95% 下界；`bench/reports/baseline_*.json` 入库 - [ ] **b-4** 外部基线 - - [ ] b-4.1 仓库内 Mem0/Zep 风格同构对照：扁平 store + user/agent/run 作用域 + 无失效 - - [ ] b-4.2 同语料同硬件跑，出对照数字 + - [x] b-4.1 仓库内 Mem0/Zep 风格同构对照：扁平 store + holder 作用域 + 无失效（`bench/counterpart.py`：`split_items` 400 字切句、`(holder, sha256(text))` 去重＝Mem0、`supersede` 按 (holder, subject) 时间序关闭旧条目＝Zep；`visible(holder, as_of)` 只对 temporal store 生效）+ `bench/tests/test_counterpart.py` 11 条 + - [~] b-4.2 同语料同硬件跑，出对照数字：`bench/run_external.py` 已跑通管道（shard 试跑：7,560 atom → 30,230 item，Zep 关闭 642 条），全量数字等 ingest - [ ] b-4.3 附录：真 Mem0 跑法与参考值，标明其 `add()` 依赖 LLM 的不可比因素 - 验收：`bench/reports/baseline_external_*.json`；对照基线与 ely-eye 同标注集同指标 -- [ ] **b-5** 反自证审计 - - [ ] b-5.1 单次尝试：禁 best-of-N 重摇 - - [ ] b-5.2 禁达标即停：跑满 n 才算分 - - [ ] b-5.3 每个分数旁必有 chance 基线行 - - [ ] b-5.4 producer ≠ verifier：判分器与被评对象不共享题目生成 - - 验收：`python -m bench.audit` 对故意注入的 best-of-N 用例报错退出 +- [x] **b-5** 反自证审计（`bench/audit.py` + `bench/tests/test_audit.py` 9 条） + - [x] b-5.1 单次尝试：禁 best-of-N 重摇 —— AST 规则：`while` 循环的测试或 `break` 守卫里出现分数词（pass_rate/target/recall/accuracy/score/hit_rate）即违规；函数名含 `best_of/reroll/until_pass` 即违规。已抓到产品现行违规 `backend/src/ely_eye/evals.py:2351`（Visual HashHop 达标即停） + - [x] b-5.2 禁达标即停：跑满 n 才算分 —— 报告的 `meta.queries/answerable` 必须等于冻结集 `5308/2928`，否则必须记录 `limit_per_modality` + - [x] b-5.3 每个分数旁必有 chance 基线行 —— 报告缺 `chance.recall_at_10` 或缺任一格的 `hit_at_10_wilson_lower` 即报错 + - [x] b-5.4 producer ≠ verifier —— `bench/metrics.py` 不得 import `ely_eye/transformers/sentence_transformers/torch/bench.systems/bench.counterpart`；每条 gold 的 `origin.dataset` 必须在 `bench/corpus/manifest.jsonl` 里且 `revision` 等于冻结 revision（实测 6 个 set 全通过） + - 验收：`python -m bench.audit --self-test` 退出 0（注入的 best-of-N 与缺 chance 两例都被拒）；`python -m bench.audit --product` 现存 3 条违规退出 1（1 条产品重摇 + 2 份报告未落盘） + - 遗留：b-1.6 挪过来的「仓库级代码检索设置」（pyparsing/regex/tokio 三仓库）尚未做成查询集，记在 P2 --- diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 9f6af92..7e9d4dd 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "python-multipart>=0.0.20", "qwen-vl-utils>=0.0.14", "rank-bm25>=0.2.2", - "rapidocr-onnxruntime>=1.4.4", + "rapidocr-onnxruntime>=1.2.3,<2", "safetensors>=0.7.0", "sentence-transformers>=5.2.0", "torch>=2.12.0", diff --git a/backend/src/ely_eye/py.typed b/backend/src/ely_eye/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 0000000..2ea105f --- /dev/null +++ b/bench/README.md @@ -0,0 +1,62 @@ +# bench — the measurement harness + +Everything here measures the product as it is. No arm may read a benchmark gold label. + +## run order + +```bash +python -m bench.fetch_corpus # 19 GB of third-party material -> bench/corpus + ~/ely_eye_bench_data +python -m bench.build_sets # freeze the 6 query sets -> bench/sets/ +python -m bench.validate_sets # schema + gold + scope + modality checks (b-2 gate) +python -m bench.fetch_models # dense + rerank models used by the ablation arms +python -m bench.ingest_baseline # the unmodified product ingesting the whole corpus (16 -> 8 workers) +python -m bench.run # five-arm ablation + latency/scale -> bench/reports/BASELINE.md +python -m pytest bench/tests -q # metric definitions, hand-computed +``` + +`bench/runs/` and `bench/corpus/` are gitignored; `bench/corpus/manifest.jsonl` is tracked so the +corpus is reproducible from the manifest alone. + +## the arms (b-3) + +| arm | what it is | +| --- | --- | +| `bm25_only` | the product's sparse signal alone | +| `dense_only` | the product's dense signal alone (product flag is off today) | +| `graph_only` | the product's graph signal alone — `atom.relations` is never populated, so it is a constant | +| `no_rerank` | the product's exact 0.55 / 0.35 / 0.10 fusion | +| `fusion_rerank` | the same fusion plus a cross-encoder rerank, the upper bound of the current design | + +`python -m bench.run` also drives the **unmodified** `RetrievalService.search` over the merged +database and reports its latency plus how far its top-10 differs from the fast BM25 path. + +## known environment artifact + +pi-lens' pyright reports `Import "bench.sets_common" could not be resolved` (and the same for +`bench.metrics`, `bench.systems`, `ely_eye.*`, `cv2`, `onnxruntime`, `pypdf.errors`) even though +every one of those imports works. It is not a repo defect: + +- `effective_config` on any bench file returns `0 config file(s)` with `resolvedCwd` at the repo + root, so neither `pyproject.toml [tool.pyright]` nor `pyrightconfig.json` (both of which carry + `extraPaths: [".", "backend/src"]`) is being read by that runner. +- Runtime proof, run from the repo root with `.venv` active: + +```bash +.venv/bin/python -c "import bench.sets_common, bench.metrics, bench.systems, bench.run, bench.ingest_baseline, \ +bench.build_mmlongbench, ely_eye.config, ely_eye.db, cv2, onnxruntime, pypdf.errors; print('runtime proof OK')" +``` + +Treat that command as the check of record for import health; treat pyright's *type* findings +(it has caught real defects here: an `AtomRunner` typo, missing `strict=True` on `zip`, +`dict[str, object]` over-narrowing) as separate and worth fixing. + +## corpus facts that constrain the numbers + +- 50,789 sources / 5,308 queries (2,928 answerable, 2,380 unanswerable), all third-party. +- Video-MME ships mp4s for 129 of 900 videos; the set uses those 129 (387 questions). +- Video-MME has no temporal gold, so its gold region is the whole video — recorded in the set manifest. +- ScreenSpot-v2 boxes are `[x, y, w, h]`, verified against 755 image headers (xywh in bounds 1271/1272). +- CodeSearchNet docstrings appear verbatim inside the code, so the query text is stripped before ingest + or BM25 gets a free exact match. +- The product keys a source by `{stem[:48]}_{sha256[:16]}`, so identical bytes in two sets collapse to one + product source; `ingest_report.json` reports `shared_source_ids`. diff --git a/bench/audit.py b/bench/audit.py new file mode 100644 index 0000000..3479f26 --- /dev/null +++ b/bench/audit.py @@ -0,0 +1,278 @@ +"""b-5: the audit that trusts nothing we wrote, including this file. + +Four rules, each machine-checked: + +1. single_attempt - no loop may re-roll a case until a score target clears (best-of-N). +2. full_n - a frozen score must cover the frozen n, or say which --limit it used. +3. chance_row - every scored arm or store sits next to the score random guessing gets. +4. producer_verifier - gold labels name third-party datasets, and the scorer imports no model. + +Run: python -m bench.audit # audit the repo and the reports + python -m bench.audit --self-test # prove the auditor catches injected cheating +""" + +from __future__ import annotations + +import argparse +import ast +import json +import sys +from pathlib import Path +from typing import Any + +from bench.sets_common import MODALITIES, SETS + +ROOT = Path(__file__).resolve().parent +REPORTS = ROOT / "reports" +CORPUS_MANIFEST = ROOT / "corpus" / "manifest.jsonl" + +# Words that only appear where a score is being compared to a target. A while-loop that re-rolls +# until one of them clears is best-of-N whatever the function is called. +SCORE_WORDS = ("pass_rate", "pass_target", "accuracy", "recall", "hit_rate", "score", "target") +RETRY_MARKERS = ("best_of", "bestof", "reroll", "re_roll", "until_pass", "retry_until") +# The scorer must stay a pure function of (query, ranked ids, gold ids). Anything else and the +# thing being graded starts grading itself. +SCORER_FORBIDDEN = ("ely_eye", "transformers", "sentence_transformers", "torch", "bench.systems", "bench.counterpart") +FROZEN_REPORTS = ("baseline.json", "baseline_external.json") + + +def _words_in(node: ast.AST) -> set[str]: + """Score vocabulary reachable from one expression, by name or attribute.""" + found: set[str] = set() + for child in ast.walk(node): + if isinstance(child, ast.Name | ast.Attribute | ast.Constant): + text = str(getattr(child, "attr", None) or getattr(child, "id", None) or getattr(child, "value", "")) + text = text.lower() + found.update(word for word in SCORE_WORDS if word in text) + return found + + +def best_of_violations(source: str, name: str) -> list[str]: + """Flag re-rolling: a while-loop that breaks on a score comparison, or an explicit best-of marker.""" + try: + tree = ast.parse(source, filename=str(name)) + except SyntaxError as error: + raise SystemExit(f"{name}: cannot parse ({error})") from error + violations: list[str] = [] + for node in ast.walk(tree): + if not isinstance(node, ast.While): + continue + loop_words = _words_in(node.test) + break_words: set[str] = set() + for inner in ast.walk(node): + if isinstance(inner, ast.Break): + for parent in _enclosing_tests(inner, node): + break_words |= _words_in(parent) + if loop_words or break_words: + words = ", ".join(sorted(loop_words | break_words)) + violations.append(f"{name}:{node.lineno}: retry-until-score loop re-rolls on {words}") + for node in ast.walk(tree): + if isinstance(node, ast.Name | ast.Attribute | ast.keyword): + text = str(getattr(node, "attr", None) or getattr(node, "id", None) or getattr(node, "arg", "")).lower() + elif isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef): + text = node.name.lower() # a function called collect_until_pass is the same trick + else: + continue + if any(marker in text for marker in RETRY_MARKERS): + violations.append(f"{name}:{node.lineno}: {text!r} re-rolls until a target") + return violations + + +def _enclosing_tests(break_node: ast.Break, loop: ast.While) -> list[ast.expr]: + """Tests guarding this break between it and the loop head; a bare break counts as the loop test.""" + tests: list[ast.expr] = [] + for node in ast.walk(loop): + if isinstance(node, ast.If) and any(child is break_node for child in node.body): + tests.append(node.test) + return tests or [loop.test] + + +def audit_sources(paths: list[Path]) -> list[str]: + """Rule 1 over our own harness: the benchmark code may never re-roll a query.""" + violations: list[str] = [] + for path in paths: + violations.extend(best_of_violations(path.read_text(encoding="utf-8"), str(path.relative_to(ROOT.parent)))) + return violations + + +def frozen_counts() -> dict[str, Any]: + """The n every frozen report is required to cover.""" + path = SETS / "manifest.json" + try: + manifest = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise SystemExit(f"{path}: cannot read the frozen set manifest ({error})") from error + missing = [key for key in ("queries", "queries_answerable", "modalities") if key not in manifest] + if missing: + raise SystemExit(f"{path}: manifest is missing {missing}") + return manifest + + +def report_violations(path: Path, manifest: dict[str, Any]) -> list[str]: + """Rules 2 and 3 over one report: full n, and a chance row beside every score.""" + try: + report = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + return [f"{path.name}: unreadable report ({error})"] + violations: list[str] = [] + meta = report.get("meta", {}) + if ( + (meta.get("queries") != manifest["queries"] or meta.get("answerable") != manifest["queries_answerable"]) + and not meta.get("limit_per_modality") + ): + violations.append( + f"{path.name}: scored {meta.get('queries')}/{meta.get('answerable')} queries but the frozen set " + f"is {manifest['queries']}/{manifest['queries_answerable']} and no --limit is recorded" + ) + chance = report.get("chance") + if chance is None: + violations.append(f"{path.name}: no chance row, so no score in it can be read as signal") + scored = report.get("arms") or report.get("stores") or [] + if not scored: + violations.append(f"{path.name}: no arms or stores scored") + for entry in scored: + name = entry.get("arm") or entry.get("store") or "?" + overall = entry.get("overall") or {} + for key in ("recall_at_10", "hit_at_10", "hit_at_10_wilson_lower"): + if key not in overall: + violations.append(f"{path.name}: {name} is missing {key}") + if chance is not None and overall.get("recall_at_10") is not None and chance.get("recall_at_10") is None: + violations.append(f"{path.name}: chance row has no recall_at_10 to compare {name} against") + return violations + + +def scorer_is_independent() -> list[str]: + """Rule 4a: the scorer is a pure function of ids, so it cannot share the question's generation.""" + path = ROOT / "metrics.py" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + violations: list[str] = [] + for node in ast.walk(tree): + modules: list[str] = [] + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom) and node.module: + modules = [node.module] + line = getattr(node, "lineno", 0) + for module in modules: + if any(module == banned or module.startswith(f"{banned}.") for banned in SCORER_FORBIDDEN): + violations.append(f"bench/metrics.py:{line}: imports {module}, which the scorer must not need") + return violations + + +def third_party_origins() -> list[str]: + """Rule 4b: every gold label must name a dataset we downloaded, at the revision we froze.""" + datasets: dict[str, str] = {} + try: + lines = CORPUS_MANIFEST.read_text(encoding="utf-8").splitlines() + except OSError as error: + raise SystemExit(f"{CORPUS_MANIFEST}: cannot read the corpus manifest ({error})") from error + for number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + datasets[str(record["dataset"])] = str(record["revision"]) + except (json.JSONDecodeError, KeyError) as error: + raise SystemExit(f"{CORPUS_MANIFEST}:{number}: unreadable record ({error})") from error + if not datasets: + raise SystemExit(f"{CORPUS_MANIFEST}: names no dataset") + violations: list[str] = [] + for modality in MODALITIES: + path = SETS / f"{modality}.jsonl" + if not path.is_file(): + violations.append(f"{path}: missing set file") + continue + for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise SystemExit(f"{path}:{number}: unreadable record ({error})") from error + origin = record.get("origin") or {} + dataset = str(origin.get("dataset", "")) + if dataset not in datasets: + violations.append(f"bench/sets/{modality}.jsonl:{number}: origin {dataset!r} is not a frozen dataset") + elif origin.get("revision") != datasets[dataset]: + violations.append( + f"bench/sets/{modality}.jsonl:{number}: {dataset} is cited at revision " + f"{origin.get('revision')}, frozen at {datasets[dataset]}" + ) + return violations + + +def injected_cases() -> list[tuple[str, list[str]]]: + """Cheating we plant on purpose; the auditor must refuse each one.""" + best_of_n = ( + "def collect(proofs):\n" + " while len(proofs) < 3:\n" + " proofs.append(run_once())\n" + " if pass_rate >= target:\n" + " break\n" + ) + no_chance = {"meta": {"queries": 5308, "answerable": 2928}, "arms": [{"arm": "x", "overall": {}}]} + return [(best_of_n, ["retry-until-score"]), (json.dumps(no_chance), ["no chance row"])] + + +def self_test(manifest: dict[str, Any]) -> int: + """Prove the auditor catches the injected cases; a rule that catches nothing is not a rule.""" + failures = 0 + for source, expected in injected_cases(): + if expected[0] == "retry-until-score": + found = best_of_violations(source, "injected.py") + else: + path = Path("injected_report.json") + path.write_text(source, encoding="utf-8") + found = report_violations(path, manifest) + path.unlink() + if not found: + print(f"SELF-TEST FAIL: the auditor accepted an injected case ({expected[0]})") + failures += 1 + else: + print(f"self-test ok: caught {expected[0]} -> {found[0]}") + return failures + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--self-test", action="store_true", help="check the auditor against injected cheating") + parser.add_argument("--product", action="store_true", help="also audit the product's proof suite for re-rolling") + args = parser.parse_args() + + manifest = frozen_counts() + if args.self_test: + return 1 if self_test(manifest) else 0 + + sources = [ + path + for path in sorted(ROOT.rglob("*.py")) + if "tests" not in path.parts and "corpus" not in path.parts and path.name != "audit.py" + ] + violations = audit_sources(sources) + if args.product: + product = ROOT.parent / "backend" / "src" / "ely_eye" / "evals.py" + if not product.is_file(): + raise SystemExit(f"{product}: --product asked for the proof suite but it is not there") + violations.extend(best_of_violations(product.read_text(encoding="utf-8"), "backend/src/ely_eye/evals.py")) + + for name in FROZEN_REPORTS: + path = REPORTS / name + if path.is_file(): + violations.extend(report_violations(path, manifest)) + else: + violations.append(f"{path}: frozen report missing") + violations.extend(scorer_is_independent()) + violations.extend(third_party_origins()) + + print(f"audited {len(sources)} harness files, {len(FROZEN_REPORTS)} reports, {len(MODALITIES)} set files") + if violations: + for violation in violations: + print(f"VIOLATION: {violation}", file=sys.stderr) + print(f"{len(violations)} violation(s)", file=sys.stderr) + return 1 + print("clean: single attempt, full n, chance rows present, producer is not verifier") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/build_code.py b/bench/build_code.py index 98c7c5d..d5db48c 100644 --- a/bench/build_code.py +++ b/bench/build_code.py @@ -217,7 +217,7 @@ def build() -> tuple[list, list, dict]: def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "code.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "code.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "code.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports").mkdir(parents=True, exist_ok=True) (SETS / "reports" / "code.json").write_text( diff --git a/bench/build_longmemeval.py b/bench/build_longmemeval.py index aa17ce8..0e5dbca 100644 --- a/bench/build_longmemeval.py +++ b/bench/build_longmemeval.py @@ -255,7 +255,7 @@ def _counts(queries: list[dict[str, Any]]) -> dict[str, int]: def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "conversation.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "conversation.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "conversation.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports").mkdir(parents=True, exist_ok=True) (SETS / "reports" / "conversation.json").write_text( diff --git a/bench/build_mmlongbench.py b/bench/build_mmlongbench.py index f74eead..7e79f89 100644 --- a/bench/build_mmlongbench.py +++ b/bench/build_mmlongbench.py @@ -14,6 +14,7 @@ from typing import Any from pyarrow.parquet import ParquetFile from pypdf import PdfReader +from pypdf.errors import PyPdfError from bench.corpus_sources import source_for from bench.sets_common import ( @@ -44,7 +45,7 @@ def pdf_pages(doc_id: str) -> int | None: return None try: return len(PdfReader(path).pages) - except Exception as exc: + except (OSError, PyPdfError) as exc: print(f" unreadable pdf {doc_id}: {type(exc).__name__}: {exc}") return None @@ -196,7 +197,7 @@ def build() -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]] def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "pdf.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "pdf.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "pdf.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports" / "pdf.json").parent.mkdir(parents=True, exist_ok=True) (SETS / "reports" / "pdf.json").write_text(json.dumps(report, indent=1, sort_keys=True)) diff --git a/bench/build_screenspot.py b/bench/build_screenspot.py index 10027a6..45dc2d5 100644 --- a/bench/build_screenspot.py +++ b/bench/build_screenspot.py @@ -170,7 +170,7 @@ def cross_holder( def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "ui_screenshot.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "ui_screenshot.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "ui_screenshot.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports").mkdir(parents=True, exist_ok=True) (SETS / "reports" / "ui_screenshot.json").write_text( diff --git a/bench/build_videomme.py b/bench/build_videomme.py index 9837b1d..aefac8b 100644 --- a/bench/build_videomme.py +++ b/bench/build_videomme.py @@ -185,7 +185,7 @@ def _counts(queries: list[dict[str, Any]], field: str) -> dict[str, int]: def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "video.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "video.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "video.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports").mkdir(parents=True, exist_ok=True) (SETS / "reports" / "video.json").write_text( diff --git a/bench/build_visrag.py b/bench/build_visrag.py index c09eee4..f46ab20 100644 --- a/bench/build_visrag.py +++ b/bench/build_visrag.py @@ -228,7 +228,7 @@ def build() -> tuple[list, list, dict]: def main() -> None: corpus, queries, report = build() corpus_count, corpus_sha = write_jsonl(SETS / "corpus" / "image.jsonl", corpus) - query_count, query_sha = write_jsonl(SETS / "image.jsonl", queries) + _query_count, query_sha = write_jsonl(SETS / "image.jsonl", queries) report |= {"corpus_records": corpus_count, "corpus_sha256": corpus_sha, "queries_sha256": query_sha} (SETS / "reports" / "image.json").parent.mkdir(parents=True, exist_ok=True) (SETS / "reports" / "image.json").write_text(json.dumps(report, indent=1, sort_keys=True)) diff --git a/bench/configs/ablation.json b/bench/configs/ablation.json new file mode 100644 index 0000000..fab555d --- /dev/null +++ b/bench/configs/ablation.json @@ -0,0 +1,5 @@ +{ + "arms": ["bm25_only", "dense_only", "graph_only", "no_rerank", "fusion_rerank"], + "no_live": false, + "no_scale": false +} diff --git a/bench/counterpart.py b/bench/counterpart.py new file mode 100644 index 0000000..ba3dfbe --- /dev/null +++ b/bench/counterpart.py @@ -0,0 +1,321 @@ +"""Mem0- and Zep-style memory counterparts, so the baseline table has an external reference point. + +Both counterparts read the same ingested corpus and the same embeddings as the ely-eye arms; +only the memory management differs, which is the thing being compared. + +* Mem0-style: one flat store, every item scoped by holder, sha256 dedup, no invalidation — + a newer fact about the same subject sits next to the old one forever. +* Zep-style: the same items carry valid_from / valid_to; a newer item about the same subject + closes the older one, and a query at time T cannot see a closed item. + +Extraction is lexical, not LLM: Mem0's real `add()` and Zep's graph builder call a generator +model, and this round has no generator installed. The deviation is recorded in the report and +in bench/README.md; it is the same extractor for both counterparts, so the comparison between +them stays fair. + +Run: python -m bench.run_external +""" + +from __future__ import annotations + +import hashlib +import re +import sqlite3 +import time +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass, field, replace +from itertools import pairwise +from pathlib import Path + +import numpy as np + +from bench.ingest_baseline import source_map +from bench.sets_common import MODALITIES, SETS, read_jsonl +from bench.systems import ( + DENSE_MODEL, + MODELS, + Atom, + DenseIndex, + local_model_dir, + tokenize, +) + +ITEM_CHARS = 400 +OVERLAP_SUPERSEDE_LIMIT = 0.5 +MIN_SUBJECT_TOKEN_CHARS = 4 +STOPWORDS = frozenset( + [ + "the", "a", "an", "and", "or", "of", "to", "in", "on", "for", "with", "that", "this", + "these", "those", "is", "are", "was", "were", "be", "been", "being", "what", "which", + "who", "whom", "whose", "when", "where", "why", "how", "not", "no", "nor", "but", "if", + "then", "than", "so", "as", "at", "by", "from", "into", "about", "over", "under", + "again", "further", "once", "here", "there", "all", "any", "both", "each", "few", + "more", "most", "other", "some", "such", "only", "own", "same", "he", "she", "they", + "them", "his", "her", "their", "its", "it", "we", "you", "i", "me", "us", "him", + "一个", "什么", "怎么", "为什么", "哪里", "时候", "这个", "那个", "可以", "应该", + "我们", "你们", "他们", "自己", "因为", "所以", "但是", "如果", + ] +) + + +@dataclass(frozen=True) +class SourceMeta: + """What the frozen corpus index says about one source.""" + + source_id: str + holder: str + scenario: str + modality: str + observed_at: str + + +@dataclass +class MemoryItem: + """One remembered thing, with the three facts a shared memory needs: who, when, what.""" + + item_id: str + text: str + source_id: str + holder: str + scenario: str + observed_at: str + subject: str + valid_from: str + valid_to: str | None = None + + +def load_source_index() -> dict[str, SourceMeta]: + """source_id -> holder / scenario / modality / observed_at, from the frozen corpus index.""" + index: dict[str, SourceMeta] = {} + for modality in MODALITIES: + for record in read_jsonl(SETS / "corpus" / f"{modality}.jsonl"): + missing = [key for key in ("source_id", "holder", "scenario", "observed_at") if not record.get(key)] + if missing: + raise SystemExit(f"{modality}: corpus record {record.get('source_id')} lacks {missing}") + meta = SourceMeta( + source_id=str(record["source_id"]), + holder=str(record["holder"]), + scenario=str(record["scenario"]), + modality=modality, + observed_at=str(record["observed_at"]), + ) + if meta.source_id in index: + raise SystemExit(f"{meta.source_id}: source id appears in two corpus indexes") + index[meta.source_id] = meta + if not index: + raise SystemExit("no corpus index found, run python -m bench.build_sets first") + return index + + +def load_item_atoms(db_path: Path) -> list[Atom]: + """Atom text for every source the corpus index names, re-keyed from product ids to corpus ids. + + The product names a source by file content ({stem}_{sha256[:16]}, storage.py:28), which is not the + id the frozen sets use. The sources table keeps the file path, and bench.ingest_baseline.source_map() + is the corpus-id-to-path map, so the join recovers the corpus identity the gold answers are written in. + """ + if not db_path.is_file(): + raise SystemExit(f"{db_path}: no baseline database, run python -m bench.ingest_baseline first") + path_to_corpus_id = source_map() + connection = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) + try: + rows = connection.execute( + "SELECT a.atom_id, s.path, a.text FROM atoms a JOIN sources s ON s.source_id = a.source_id" + " ORDER BY a.atom_id" + ).fetchall() + except sqlite3.OperationalError as error: + raise SystemExit(f"{db_path}: cannot read the atoms table ({error})") from error + finally: + connection.close() + if not rows: + raise SystemExit(f"{db_path}: the atoms table is empty") + atoms = [ + Atom(atom_id=str(atom_id), source_id=path_to_corpus_id[path], text=str(text or "")) + for atom_id, path, text in rows + if path in path_to_corpus_id + ] + if not atoms: + raise SystemExit(f"{db_path}: none of its {len(rows)} atoms map to a corpus source id") + return atoms + + +def split_items(text: str) -> list[str]: + """Cut source text into rememberable chunks; sentence boundaries first, then a hard cut.""" + cleaned = re.sub(r"\s+", " ", text).strip() + if not cleaned: + return [] + pieces: list[str] = [] + for sentence in re.split(r"(?<=[.!?。！？；;])\s+", cleaned): + sentence = sentence.strip() + if not sentence: + continue + while len(sentence) > ITEM_CHARS: + pieces.append(sentence[:ITEM_CHARS]) + sentence = sentence[ITEM_CHARS:] + pieces.append(sentence) + merged: list[str] = [] + buffer = "" + for piece in pieces: + if not buffer: + buffer = piece + elif len(buffer) + 1 + len(piece) <= ITEM_CHARS: + buffer = f"{buffer} {piece}" + else: + merged.append(buffer) + buffer = piece + if buffer: + merged.append(buffer) + return merged + + +def subject_of(text: str) -> str: + """The lexical stand-in for an entity: the most frequent content token in the item.""" + counts = Counter( + token + for token in tokenize(text) + if len(token) >= MIN_SUBJECT_TOKEN_CHARS and token not in STOPWORDS and not token.isdigit() + ) + if not counts: + return "" + token, count = counts.most_common(1)[0] + return token if count >= 2 else "" + + +def overlap(first: str, second: str) -> float: + """Jaccard overlap of two items' token sets; low overlap means they are not the same statement.""" + a, b = set(tokenize(first)), set(tokenize(second)) + if not a or not b: + return 0.0 + return len(a & b) / len(a | b) + + +def build_items(atoms: Sequence[Atom], index: dict[str, SourceMeta]) -> list[MemoryItem]: + """Turn ingested atoms into memory items, deduped by (holder, text hash) the way Mem0 dedups.""" + items: list[MemoryItem] = [] + seen: set[tuple[str, str]] = set() + unknown: set[str] = set() + for atom in atoms: + meta = index.get(atom.source_id) + if meta is None: + unknown.add(atom.source_id) + continue + for position, text in enumerate(split_items(atom.text)): + key = (meta.holder, hashlib.sha256(text.encode("utf-8")).hexdigest()) + if key in seen: + continue + seen.add(key) + items.append( + MemoryItem( + item_id=f"{atom.source_id}#{position}", + text=text, + source_id=atom.source_id, + holder=meta.holder, + scenario=meta.scenario, + observed_at=meta.observed_at, + subject=subject_of(text), + valid_from=meta.observed_at, + ) + ) + if unknown: + print(f"note: {len(unknown)} ingested sources are not in the corpus index and carry no memory scope") + return items + + +def supersede(items: Sequence[MemoryItem]) -> int: + """Close older items about the same subject when a newer, different statement arrives (Zep-style).""" + by_subject: dict[tuple[str, str], list[MemoryItem]] = {} + for item in items: + if item.subject: + by_subject.setdefault((item.holder, item.subject), []).append(item) + closed = 0 + for group in by_subject.values(): + group.sort(key=lambda item: (item.observed_at, item.item_id)) + for older, newer in pairwise(group): + if older.valid_to is not None or older.observed_at >= newer.observed_at: + continue + if overlap(older.text, newer.text) >= OVERLAP_SUPERSEDE_LIMIT: + continue + older.valid_to = newer.observed_at + closed += 1 + return closed + + +@dataclass +class Counterpart: + """One counterpart store: the same items, the same embeddings, different memory semantics.""" + + name: str + temporal: bool + items: list[MemoryItem] + model_dir: Path + cache_name: str = "counterpart-items" + timings: list[float] = field(default_factory=list) + _index: DenseIndex | None = None + + def __post_init__(self) -> None: + """The store is cheap to build; embeddings are paid for on the first search.""" + + def _ensure_index(self) -> DenseIndex: + if self._index is None: + self._index = DenseIndex( + [Atom(atom_id=item.item_id, source_id=item.source_id, text=item.text) for item in self.items], + self.model_dir, + cache_name=self.cache_name, + ) + return self._index + + def visible(self, holder: str | None, as_of: str | None) -> np.ndarray: + """Which items this query may see. + + Both stores isolate by holder. Only the temporal store uses time: an item is hidden once + a newer statement closed it, and hidden until it was observed. The flat store has no notion + of as_of, which is exactly Mem0's behaviour and exactly what the comparison measures. + """ + mask = np.ones(len(self.items), dtype=bool) + for position, item in enumerate(self.items): + if holder is not None and item.holder != holder: + mask[position] = False + continue + if not self.temporal or as_of is None: + continue + if item.valid_from > as_of or (item.valid_to is not None and item.valid_to <= as_of): + mask[position] = False + return mask + + def search(self, query: str, holder: str | None, as_of: str | None, top_k: int = 10) -> list[str]: + """Rank sources for one query, restricted to the items the query is allowed to see.""" + started = time.perf_counter() + scores = self._ensure_index().scores(query, self.model_dir) + mask = self.visible(holder, as_of) + best: dict[str, float] = {} + for position, item in enumerate(self.items): + if not mask[position]: + continue + current = best.get(item.source_id) + if current is None or scores[position] > current: + best[item.source_id] = float(scores[position]) + ranked = sorted(best.items(), key=lambda entry: (-entry[1], entry[0])) + self.timings.append(time.perf_counter() - started) + return [source_id for source_id, _ in ranked[:top_k]] + + def percentile(self, fraction: float) -> float: + if not self.timings: + raise SystemExit(f"{self.name}: no timings recorded") + if not 0.0 <= fraction <= 1.0: + raise SystemExit(f"percentile must be within [0, 1], got {fraction}") + return float(np.percentile(np.asarray(self.timings), fraction * 100)) + + +def build_counterparts(items: list[MemoryItem], models: Path = MODELS) -> list[Counterpart]: + """Mem0-style first (no invalidation), then Zep-style (the same items, closed by time).""" + model_dir = local_model_dir(models, DENSE_MODEL) + if not model_dir.is_dir(): + raise SystemExit(f"{model_dir}: dense model missing, run python -m bench.fetch_models first") + mem0 = Counterpart(name="mem0_flat", temporal=False, items=items, model_dir=model_dir) + zep_items = [replace(item) for item in items] # same order, same ids: only validity windows differ + closed = supersede(zep_items) + print(f"zep_temporal: closed {closed} items by supersession out of {len(zep_items)}") + # Both stores share one embedding cache because the item sequence is identical; only visibility differs. + zep = Counterpart(name="zep_temporal", temporal=True, items=zep_items, model_dir=model_dir) + return [mem0, zep] diff --git a/bench/fetch_corpus.py b/bench/fetch_corpus.py index 0ac702c..205280a 100644 --- a/bench/fetch_corpus.py +++ b/bench/fetch_corpus.py @@ -31,7 +31,7 @@ def http_get(url: str) -> tuple[bytes, dict[str, str]]: if urllib.parse.urlsplit(url).scheme != "https" or urllib.parse.urlsplit(url).netloc != "huggingface.co": raise SystemExit(f"refusing non-HF url: {url}") try: - resp = OPENER.open(urllib.request.Request(url, headers=UA), timeout=300) # noqa: S310 + resp = OPENER.open(urllib.request.Request(url, headers=UA), timeout=300) except urllib.error.HTTPError as exc: raise SystemExit(f"HTTP {exc.code} {url}") from exc except urllib.error.URLError as exc: @@ -90,7 +90,7 @@ def download_to(url: str, dest: Path) -> tuple[int, str]: digest = hashlib.sha256() size = 0 try: - resp = OPENER.open(urllib.request.Request(url, headers=UA), timeout=600) # noqa: S310 + resp = OPENER.open(urllib.request.Request(url, headers=UA), timeout=600) except urllib.error.HTTPError as exc: raise SystemExit(f"HTTP {exc.code} {url}") from exc except urllib.error.URLError as exc: diff --git a/bench/fetch_models.py b/bench/fetch_models.py new file mode 100644 index 0000000..2e8cd45 --- /dev/null +++ b/bench/fetch_models.py @@ -0,0 +1,45 @@ +"""Download the two retrieval models the ablation needs, into bench/models/. + +Both are small on purpose: the benchmark must be runnable on a laptop, and the product's +own 8B embedder is measured separately once it is wired up. + +Run: python -m bench.fetch_models +""" + +from __future__ import annotations + +from pathlib import Path + +from huggingface_hub import snapshot_download + +from bench.systems import DENSE_MODEL, MODELS, RERANK_MODEL, local_model_dir + +WEIGHTS = ["*.safetensors", "*.json", "*.txt", "*.model", "*.tsv", "tokenizer*"] +SKIP = ["onnx/*", "*.onnx", "*.msgpack", "*.h5", "pytorch_model.bin", "tf_model.h5", "*.ot"] + + +def fetch(repo_id: str) -> Path: + """Fetch one repo's inference files and report its size.""" + target = local_model_dir(MODELS, repo_id) + target.mkdir(parents=True, exist_ok=True) + path = snapshot_download( + repo_id=repo_id, + local_dir=str(target), + allow_patterns=WEIGHTS, + ignore_patterns=SKIP, + ) + size = sum(file.stat().st_size for file in Path(path).rglob("*") if file.is_file()) + if size < 1_000_000: + raise SystemExit(f"{repo_id}: downloaded only {size} bytes to {path}, expected model weights") + print(f"{repo_id} -> {path} ({size / 1e6:.1f} MB)") + return Path(path) + + +def main() -> int: + fetch(DENSE_MODEL) + fetch(RERANK_MODEL) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/ingest_baseline.py b/bench/ingest_baseline.py new file mode 100644 index 0000000..4aff4a6 --- /dev/null +++ b/bench/ingest_baseline.py @@ -0,0 +1,404 @@ +"""Ingest the frozen benchmark sets through the unmodified pre-P0 IngestionService. + +The baseline number must come from the code as it was before any P0 change, so this +script never edits the product: it drives IngestionService over every file the corpus +index names, one sharded SQLite per worker, then merges the shards into one database. +Shards are content-addressed on disk, so merging is a plain row copy. + +Run: python -m bench.ingest_baseline [--workers 16] [--modality pdf] +""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import resource +import shutil +import sqlite3 +import sys +import time +from collections.abc import Sequence +from concurrent.futures import ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool +from pathlib import Path + +from bench.sets_common import MODALITIES, SETS, read_jsonl, record_path + +REPO = Path(__file__).resolve().parent.parent +RUNS = REPO / "bench" / "runs" / "baseline" +SHARDS = RUNS / "shards" +TABLES = ("sources", "atoms", "atom_fts", "cartridges", "cache_events", "evidence_edges") +# Virtual-address ceiling per worker. ONNX Runtime and torch reserve far more address space than +# they touch, so a 3 GiB ceiling killed OCR on large images ("Unable to allocate 582. MiB for an +# array with shape (5632, 9024, 3)") even though resident memory stayed near 1.5 GiB. +WORKER_MEMORY_BYTES = 6 * 1024**3 +MAX_POOL_ATTEMPTS = 5 +DELETE_STATEMENTS = ( + 'DELETE FROM "evidence_edges"', + 'DELETE FROM "cache_events"', + 'DELETE FROM "atom_fts"', + 'DELETE FROM "cartridges"', + 'DELETE FROM "atoms"', + 'DELETE FROM "sources"', +) +COPY_STATEMENTS = ( + 'INSERT OR IGNORE INTO "sources" SELECT * FROM shard."sources"', + 'INSERT OR IGNORE INTO "atoms" SELECT * FROM shard."atoms"', + ( + 'INSERT OR IGNORE INTO "atom_fts" (atom_id, source, text) ' + 'SELECT DISTINCT atom_id, source, text FROM shard."atom_fts"' + ), + 'INSERT OR IGNORE INTO "cartridges" SELECT * FROM shard."cartridges"', + 'INSERT OR IGNORE INTO "cache_events" SELECT * FROM shard."cache_events"', + 'INSERT OR IGNORE INTO "evidence_edges" SELECT * FROM shard."evidence_edges"', +) +COUNT_STATEMENTS = { + "sources": 'SELECT count(*) FROM "sources"', + "atoms": 'SELECT count(*) FROM "atoms"', + "atom_fts": 'SELECT count(*) FROM "atom_fts"', + "cartridges": 'SELECT count(*) FROM "cartridges"', + "cache_events": 'SELECT count(*) FROM "cache_events"', + "evidence_edges": 'SELECT count(*) FROM "evidence_edges"', +} + + +def corpus_paths(modality: str) -> list[Path]: + """Absolute path of every file one modality's corpus index names.""" + records = list(read_jsonl(SETS / "corpus" / f"{modality}.jsonl")) + paths: list[Path] = [] + seen: dict[Path, str] = {} + for record in records: + path = record_path(record).resolve() + if not path.is_file(): + raise SystemExit(f"{modality}: corpus record {record['source_id']} names a missing file {path}") + previous = seen.get(path) + if previous is not None and previous != record["source_id"]: + raise SystemExit(f"{modality}: {path} is claimed by both {previous} and {record['source_id']}") + seen[path] = record["source_id"] + paths.append(path) + return paths + + +def source_map() -> dict[str, str]: + """Map an absolute file path to the benchmark source id that owns it.""" + mapping: dict[str, str] = {} + for modality in MODALITIES: + for record in read_jsonl(SETS / "corpus" / f"{modality}.jsonl"): + path = record_path(record).resolve() + mapping[str(path)] = record["source_id"] + return mapping + + +def shard_plan(modalities: Sequence[str], workers: int, limit: int | None) -> list[list[Path]]: + """Round-robin files across workers so long videos do not pile up in one shard.""" + plan: list[list[Path]] = [[] for _ in range(workers)] + index = 0 + for modality in modalities: + paths = corpus_paths(modality) + for path in paths[:limit] if limit else paths: + plan[index % workers].append(path) + index += 1 + return [shard for shard in plan if shard] + + +def _pin_inference_threads() -> None: + """Pin every ONNX Runtime session to one thread; the model and its output are unchanged. + + rapidocr-onnxruntime 1.2.3 exposes no thread knob and ONNX Runtime defaults to every + core, so each of 16 workers opened ~90 inference threads and the machine thrashed at a + load average above 400. Wrapping InferenceSession is the only lever the library gives. + """ + from typing import Any + + import onnxruntime as ort + + original = ort.InferenceSession + + def pinned(model_path: str | bytes | os.PathLike[str], *args: Any, **kwargs: Any) -> ort.InferenceSession: + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + # The CPU arena keeps freed blocks mapped, so a worker that OCRs thousands of pages grows its + # resident set until it hits its ceiling on an 11 MiB array. Hand the memory back instead. + options.enable_cpu_mem_arena = False + kwargs["sess_options"] = options + return original(model_path, *args, **kwargs) + + ort.InferenceSession = pinned # type: ignore[method-assign] + + +def _single_threaded() -> None: + """One CPU thread per worker: 16 workers on 32 cores, no oversubscription. + + OpenCV also defaults to every core, which is why cv2 is pinned alongside ONNX. + """ + for variable in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"): + os.environ[variable] = "1" + import cv2 + + cv2.setNumThreads(1) + _pin_inference_threads() + _cap_worker_memory() + + +def _cap_worker_memory() -> None: + """Cap each worker at 3 GiB so one pathological file raises MemoryError instead of being OOM-killed. + + A 16-worker run was killed by the kernel (anon-rss 10.0 GB) and took the whole process pool + down with it; a per-worker ceiling turns that into a recorded per-file failure. + """ + resource.setrlimit(resource.RLIMIT_AS, (WORKER_MEMORY_BYTES, WORKER_MEMORY_BYTES)) + + +def ingest_shard(job: tuple[int, list[Path]]) -> dict: + """Worker: ingest this shard's files into its own ELY_EYE_HOME, skipping finished ones.""" + index, paths = job + _single_threaded() + started = time.monotonic() + home = SHARDS / f"shard-{index:02d}" + home.mkdir(parents=True, exist_ok=True) + ledger = home / "done.jsonl" + done = set() + if ledger.is_file(): + for number, line in enumerate(ledger.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + done.add(json.loads(line)["path"]) + except (json.JSONDecodeError, KeyError) as error: + raise SystemExit(f"{ledger}:{number}: unreadable resume ledger: {error}") from error + + from ely_eye.config import Settings + from ely_eye.ingestion import IngestionService + + service = IngestionService(settings=Settings(home=home)) + ingested = 0 + atoms = 0 + failed = 0 + for path in paths: + if str(path) in done: + continue + try: + response = service.ingest_path(path) + outcome = {"path": str(path), "sources": response.source_count, "atoms": response.atom_count} + atoms += response.atom_count + ingested += 1 + except Exception as error: # noqa: BLE001 - a corpus sweep must record the file that broke it, not die with it + outcome = {"path": str(path), "error": f"{type(error).__name__}: {error}"} + failed += 1 + print(f"shard-{index:02d} FAILED {path.name}: {outcome['error']}", file=sys.stderr, flush=True) + with ledger.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(outcome) + "\n") + if ingested % 25 == 0: + # ONNX Runtime keeps the address space it has touched, and one shard runs thousands of + # files, so the arena has to be handed back or the worker hits its ceiling on an 11 MiB array. + gc.collect() + elapsed = time.monotonic() - started + print(f"shard-{index:02d} {ingested}/{len(paths)} files in {elapsed / 60:.1f} min", flush=True) + return { + "shard": index, + "files": len(paths), + "ingested_now": ingested, + "failed": failed, + "skipped": len(paths) - ingested - failed, + "atoms": atoms, + "minutes": round((time.monotonic() - started) / 60, 2), + } + + +def database_path(home: Path) -> Path: + from ely_eye.config import Settings + + return Settings(home=home).db_path + + +def ledger_outcomes() -> tuple[set[str], set[str]]: + """Return (ingested paths, failed paths) recorded across every shard ledger.""" + ingested: set[str] = set() + failed: set[str] = set() + for ledger in sorted(SHARDS.glob("shard-*/done.jsonl")): + for number, line in enumerate(ledger.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + entry = json.loads(line) + path = entry["path"] + except (json.JSONDecodeError, KeyError) as error: + raise SystemExit(f"{ledger}:{number}: unreadable resume ledger: {error}") from error + (failed if "error" in entry else ingested).add(path) + return ingested, failed + + +def ingested_paths() -> set[str]: + """Every file a shard ledger recorded as ingested, excluding the files it recorded as failed.""" + return ledger_outcomes()[0] + + +def retry_failed() -> int: + """Drop failed ledger entries so those files are attempted again, archiving what was dropped.""" + archive = RUNS / "failures.jsonl" + dropped = 0 + for ledger in sorted(SHARDS.glob("shard-*/done.jsonl")): + kept: list[str] = [] + removed: list[dict] = [] + for number, line in enumerate(ledger.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + entry = json.loads(line) + except json.JSONDecodeError as error: + raise SystemExit(f"{ledger}:{number}: unreadable resume ledger: {error}") from error + if "error" in entry: + removed.append(entry) + else: + kept.append(line) + if not removed: + continue + with archive.open("a", encoding="utf-8") as handle: + for entry in removed: + handle.write(json.dumps({"shard": ledger.parent.name, **entry}) + "\n") + ledger.write_text("\n".join(kept) + ("\n" if kept else ""), encoding="utf-8") + dropped += len(removed) + return dropped + + +def merge() -> dict: + """Copy every shard's rows and content-addressed objects into the merged database.""" + main_db = database_path(RUNS) + main_db.parent.mkdir(parents=True, exist_ok=True) + from ely_eye.config import Settings + from ely_eye.db import Database + + Database(Settings(home=RUNS)).init() + connection = sqlite3.connect(main_db) + tables = [ + row[0] + for row in connection.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'") + ] + missing = sorted(set(TABLES) - set(tables)) + if missing: + raise SystemExit(f"{main_db}: schema is missing tables {missing}") + for statement in DELETE_STATEMENTS: + connection.execute(statement) + + shard_dirs = sorted(path for path in SHARDS.glob("shard-*") if path.is_dir()) + if not shard_dirs: + raise SystemExit(f"{SHARDS}: no shards to merge") + copied_objects = 0 + for shard in shard_dirs: + shard_db = database_path(shard) + if not shard_db.is_file(): + raise SystemExit(f"{shard}: no database at {shard_db}") + connection.execute("ATTACH DATABASE ? AS shard", (str(shard_db),)) + for statement in COPY_STATEMENTS: + cursor = connection.cursor() + cursor.execute(statement) + cursor.close() + connection.commit() + connection.execute("DETACH DATABASE shard") + shard_objects = shard / "objects" + if shard_objects.is_dir(): + for source in shard_objects.rglob("*"): + if not source.is_file(): + continue + target = RUNS / "objects" / source.relative_to(shard_objects) + target.parent.mkdir(parents=True, exist_ok=True) + if not target.exists(): + shutil.copy2(source, target) + copied_objects += 1 + connection.commit() + + counts = {} + for table, statement in COUNT_STATEMENTS.items(): + counts[table] = connection.execute(statement).fetchone()[0] + connection.close() + + mapping = source_map() + done, failed = ledger_outcomes() + unknown = sorted((done | failed) - set(mapping)) + if unknown: + raise SystemExit(f"{main_db}: {len(unknown)} ledger entries are not corpus files, first {unknown[:3]}") + rows = sqlite3.connect(main_db).execute("SELECT path, source_id FROM sources").fetchall() + unmapped = [path for path, _ in rows if path not in mapping] + if unmapped: + raise SystemExit(f"{main_db}: {len(unmapped)} sources are not named by the corpus index, first {unmapped[:3]}") + product_paths = {path for path, _ in rows} + unaccounted = sorted(done - product_paths) + if unaccounted: + raise SystemExit( + f"{main_db}: {len(unaccounted)} ingested files produced no product source, first {unaccounted[:3]}" + ) + collapsed: dict[str, list[str]] = {} + for path, source_id in rows: + collapsed.setdefault(source_id, []).append(mapping[path]) + covered = {source_id for ids in collapsed.values() for source_id in ids} + total = set(mapping.values()) + counts["corpus_files"] = len(mapping) + counts["files_ingested"] = len(done) + counts["files_failed"] = len(failed & set(mapping)) + counts["files_not_ingested"] = len(set(mapping) - done - failed) + counts["benchmark_sources_total"] = len(total) + counts["benchmark_sources_covered"] = len(covered) + counts["product_sources"] = len(collapsed) + counts["shared_source_ids"] = sum(1 for ids in collapsed.values() if len(ids) > 1) + if counts["files_not_ingested"] == 0 and covered != total: + raise SystemExit( + f"{main_db}: every file was attempted but {len(total - covered)} benchmark sources have no product " + f"source; {counts['files_failed']} files failed, first failures {sorted(failed)[:3]}" + ) + counts["objects_copied"] = copied_objects + (RUNS / "path_to_source_id.json").write_text(json.dumps(mapping, indent=1, sort_keys=True), encoding="utf-8") + (RUNS / "ingest_report.json").write_text(json.dumps(counts, indent=1, sort_keys=True), encoding="utf-8") + return counts + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workers", type=int, default=8) + parser.add_argument("--modality", action="append", choices=list(MODALITIES)) + parser.add_argument("--limit", type=int, help="ingest at most this many files per modality (smoke runs)") + parser.add_argument("--merge-only", action="store_true") + parser.add_argument("--retry-failed", action="store_true", help="re-attempt files a previous pass recorded as failed") + args = parser.parse_args() + + if args.retry_failed: + counts = merge() if args.merge_only else None + dropped = retry_failed() + print(f"== re-queued {dropped} failed files; their errors are archived in {RUNS / 'failures.jsonl'}") + if counts is not None: + print(json.dumps(counts, indent=1, sort_keys=True)) + return 0 + + modalities = args.modality or list(MODALITIES) + if not args.merge_only: + plan = shard_plan(modalities, args.workers, args.limit) + print(f"== ingesting {sum(len(shard) for shard in plan)} files across {len(plan)} shards", flush=True) + for attempt in range(1, MAX_POOL_ATTEMPTS + 1): + try: + with ProcessPoolExecutor(max_workers=len(plan)) as pool: + for result in pool.map(ingest_shard, [(index, shard) for index, shard in enumerate(plan)]): + print( + f"shard {result['shard']:02d}: {result['ingested_now']} ingested, " + f"{result['failed']} failed, {result['skipped']} already done, {result['atoms']} atoms", + flush=True, + ) + break + except BrokenProcessPool as error: + print( + f"== worker pool died (attempt {attempt}/{MAX_POOL_ATTEMPTS}): {error}; resuming from ledgers", + file=sys.stderr, + flush=True, + ) + if attempt == MAX_POOL_ATTEMPTS: + raise SystemExit(f"worker pool died {MAX_POOL_ATTEMPTS} times; ledgers keep progress") from error + + counts = merge() + print(json.dumps(counts, indent=1, sort_keys=True)) + print(f"== merged baseline database at {database_path(RUNS)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/metrics.py b/bench/metrics.py new file mode 100644 index 0000000..f5d4a1f --- /dev/null +++ b/bench/metrics.py @@ -0,0 +1,150 @@ +"""Retrieval metrics with no model and no dependency: recall@k, MRR, nDCG@10, Wilson bounds. + +Rankings are lists of source ids; gold is the set of source ids a third party marked +correct. Unanswerable queries carry an empty gold set, so they are scored by abstention +behaviour instead of recall, and never mixed into the recall denominator. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass + +Z_TWO_SIDED_95 = 1.959963984540054 + + +@dataclass(frozen=True) +class QueryScore: + query_id: str + answerable: bool + gold: int + recall_at_1: float + recall_at_5: float + recall_at_10: float + ndcg_at_10: float + mrr: float + first_hit_rank: int | None + abstained: bool + top_hit_from_holder: bool | None + + +def recall_at_k(ranked: Sequence[str], gold: set[str], k: int) -> float: + """Fraction of gold sources inside the first k results; 0.0 when gold is empty.""" + if not gold: + return 0.0 + return len(gold & set(ranked[:k])) / len(gold) + + +def reciprocal_rank(ranked: Sequence[str], gold: set[str]) -> float: + for position, source_id in enumerate(ranked, start=1): + if source_id in gold: + return 1.0 / position + return 0.0 + + +def first_hit_rank(ranked: Sequence[str], gold: set[str]) -> int | None: + for position, source_id in enumerate(ranked, start=1): + if source_id in gold: + return position + return None + + +def ndcg_at_10(ranked: Sequence[str], gold: set[str]) -> float: + """Binary-relevance nDCG@10; 0.0 when gold is empty.""" + if not gold: + return 0.0 + dcg = sum( + 1.0 / math.log2(position + 1) + for position, source_id in enumerate(ranked[:10], start=1) + if source_id in gold + ) + ideal = sum(1.0 / math.log2(position + 1) for position in range(1, min(len(gold), 10) + 1)) + return dcg / ideal + + +def wilson_lower_bound(successes: int, trials: int, z: float = Z_TWO_SIDED_95) -> float: + """One number: the lower end of the Wilson score interval for a binomial proportion.""" + if trials <= 0: + raise SystemExit(f"wilson needs a positive trial count, got {trials}") + if not 0 <= successes <= trials: + raise SystemExit(f"wilson needs 0 <= {successes} <= {trials}") + rate = successes / trials + denominator = 1.0 + z * z / trials + centre = rate + z * z / (2 * trials) + spread = z * math.sqrt((rate * (1.0 - rate) + z * z / (4 * trials)) / trials) + return max(0.0, (centre - spread) / denominator) + + +def chance_recall_at_k(gold: int, corpus_size: int, k: int) -> float: + """Expected recall@k for a uniformly random ranking: k/N, whatever the gold count. + + A random top-k holds a k/N share of every subset, so the share of gold found does not + depend on how many sources are gold; gold only decides whether the question is answerable. + """ + if corpus_size <= 0: + raise SystemExit(f"chance needs a positive corpus size, got {corpus_size}") + if gold <= 0: + return 0.0 + return min(1.0, k / corpus_size) + + +def score_query( + query_id: str, + ranked: Sequence[str], + gold: set[str], + holder: str | None = None, + holder_of: dict[str, str] | None = None, +) -> QueryScore: + """Score one query; abstention is only defined when the system can signal no_match.""" + abstained = not ranked + top_holder = None + if holder_of is not None and ranked: + top_holder = holder_of.get(ranked[0]) + return QueryScore( + query_id=query_id, + answerable=bool(gold), + gold=len(gold), + recall_at_1=recall_at_k(ranked, gold, 1), + recall_at_5=recall_at_k(ranked, gold, 5), + recall_at_10=recall_at_k(ranked, gold, 10), + ndcg_at_10=ndcg_at_10(ranked, gold), + mrr=reciprocal_rank(ranked, gold), + first_hit_rank=first_hit_rank(ranked, gold), + abstained=abstained, + top_hit_from_holder=None if holder is None else top_holder == holder, + ) + + +def aggregate(scores: Sequence[QueryScore]) -> dict[str, float | int]: + """Mean metrics over answerable queries plus abstention over unanswerable ones.""" + answerable = [score for score in scores if score.answerable] + unanswerable = [score for score in scores if not score.answerable] + if not answerable: + raise SystemExit("aggregate needs at least one answerable query") + + def mean(field: str) -> float: + return sum(getattr(score, field) for score in answerable) / len(answerable) + + hits_10 = sum(1 for score in answerable if score.recall_at_10 > 0.0) + holder_known = [score for score in unanswerable if score.top_hit_from_holder is not None] + return { + "n_answerable": len(answerable), + "n_unanswerable": len(unanswerable), + "recall_at_1": mean("recall_at_1"), + "recall_at_5": mean("recall_at_5"), + "recall_at_10": mean("recall_at_10"), + "ndcg_at_10": mean("ndcg_at_10"), + "mrr": mean("mrr"), + "hit_at_10": hits_10 / len(answerable), + "hit_at_10_wilson_lower": wilson_lower_bound(hits_10, len(answerable)), + "abstention_rate_answerable": sum(1 for s in answerable if s.abstained) / len(answerable), + "abstention_rate_unanswerable": ( + sum(1 for s in unanswerable if s.abstained) / len(unanswerable) if unanswerable else 0.0 + ), + "wrong_holder_confidence_unanswerable": ( + sum(1 for s in holder_known if s.top_hit_from_holder) / len(holder_known) + if holder_known + else 0.0 + ), + } diff --git a/bench/run.py b/bench/run.py new file mode 100644 index 0000000..ea4fbde --- /dev/null +++ b/bench/run.py @@ -0,0 +1,359 @@ +"""Run the retrieval ablation over the frozen corpus and write the baseline report. + +Five arms, one database: sparse only, dense only, graph only, the product's weighted +fusion, and fusion plus a cross-encoder reranker. Every number in the report is either +measured or explicitly marked as chance. + +Run: python -m bench.run [--arms bm25_only] [--limit 200] [--no-live] [--out bench/reports] +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import subprocess +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from bench.metrics import aggregate, chance_recall_at_k, score_query +from bench.sets_common import MODALITIES, SETS, read_jsonl, sha256_file +from bench.systems import ( + ARMS, + DENSE_MODEL, + MODELS, + RERANK_MODEL, + ArmRunner, + Atom, + arm_needs, + live_product_rankings, + load_atoms, +) + +ROOT = SETS.parent +REPORTS = ROOT / "reports" +DB_HOME = ROOT / "runs" / "baseline" +LIVE_QUERIES = 120 +SCALE_SIZES = (500, 2_000, 10_000, 50_000) + + +@dataclass(frozen=True) +class Query: + query_id: str + text: str + modality: str + gold: frozenset[str] + holder: str | None + answerable: bool + as_of: str | None + + +def load_queries(modalities: Sequence[str] = MODALITIES) -> list[Query]: + """Every query the suite froze, in file order.""" + queries: list[Query] = [] + for modality in modalities: + path = SETS / f"{modality}.jsonl" + if not path.is_file(): + raise SystemExit(f"{path}: missing set, run python -m bench.build_sets first") + for record in read_jsonl(path): + gold = frozenset(hit["source_id"] for hit in record["gold"]) + answerable = not bool(record["unanswerable"]) + if answerable != bool(gold): + raise SystemExit( + f"{path}: {record['query_id']} carries {len(gold)} gold hits but " + f"unanswerable={record['unanswerable']}" + ) + queries.append( + Query( + query_id=record["query_id"], + text=record["question"], + modality=modality, + gold=gold, + holder=record["scope"]["holder"], + answerable=answerable, + as_of=record.get("as_of"), + ) + ) + if not queries: + raise SystemExit("no queries loaded") + return queries + + +def holder_index(modalities: Sequence[str] = MODALITIES) -> dict[str, str]: + """source_id -> holder, from the corpus side of the suite.""" + holders: dict[str, str] = {} + for modality in modalities: + path = SETS / "corpus" / f"{modality}.jsonl" + if not path.is_file(): + raise SystemExit(f"{path}: missing corpus index, run python -m bench.build_sets first") + for record in read_jsonl(path): + holders[record["source_id"]] = record["holder"] + return holders + + +def chance_row(queries: Sequence[Query], corpus_sources: int) -> dict[str, float]: + """What a random ranking would score on the same queries.""" + answerable = [query for query in queries if query.answerable] + if not answerable: + raise SystemExit("chance needs answerable queries") + return { + "corpus_sources": corpus_sources, + "recall_at_1": sum(chance_recall_at_k(len(query.gold), corpus_sources, 1) for query in answerable) / len(answerable), + "recall_at_5": sum(chance_recall_at_k(len(query.gold), corpus_sources, 5) for query in answerable) / len(answerable), + "recall_at_10": sum(chance_recall_at_k(len(query.gold), corpus_sources, 10) for query in answerable) / len(answerable), + } + + +def scale_curve(queries: Sequence[Query], atoms: Sequence[Atom], arms: Sequence[str]) -> list[dict[str, Any]]: + """Latency of each arm as the corpus grows, measured on nested source prefixes.""" + sources = sorted({atom.source_id for atom in atoms}) + rows: list[dict[str, Any]] = [] + sample = [query for query in queries if query.answerable][:40] + for size in SCALE_SIZES: + keep = set(sources[:size]) + if not keep: + raise SystemExit(f"scale size {size} kept no sources") + subset = [atom for atom in atoms if atom.source_id in keep] + probe = [query for query in sample if query.gold & keep] + if not probe: + continue + row: dict[str, Any] = {"sources": len(keep), "atoms": len(subset), "queries": len(probe)} + for arm in arms: + if not arm_needs(arm, MODELS): + continue + runner = ArmRunner(arm, subset, MODELS) + for query in probe: + runner.rank(query.text) + row[arm] = {"p50": runner.percentile(0.50), "p95": runner.percentile(0.95)} + rows.append(row) + print(f"scale {len(keep):6d} sources: " + " ".join(f"{arm}={row[arm]['p50']:.3f}s" for arm in arms if arm in row), flush=True) + return rows + + +def live_check(queries: Sequence[Query], fast_rankings: dict[str, list[str]], top_k: int = 10) -> dict[str, Any]: + """Compare the unmodified RetrievalService with the fast BM25 path.""" + probe = [query for query in queries if query.answerable] + stride = max(1, len(probe) // LIVE_QUERIES) + chosen = probe[::stride][:LIVE_QUERIES] + if not chosen: + raise SystemExit("live check needs answerable queries") + rankings, timings = live_product_rankings(DB_HOME, [query.text for query in chosen], top_k) + overlaps = [ + len(set(fast_rankings[query.query_id]) & set(ranking)) / top_k + for query, ranking in zip(chosen, rankings, strict=True) + if fast_rankings.get(query.query_id) is not None + ] + if not overlaps: + raise SystemExit("live check found no fast rankings to compare against") + return { + "queries": len(chosen), + "p50": sorted(timings)[len(timings) // 2], + "mean": sum(timings) / len(timings), + "max": max(timings), + "top10_overlap_with_bm25_only": sum(overlaps) / len(overlaps), + "identical_top10": sum(1 for overlap in overlaps if overlap == 1.0) / len(overlaps), + } + + +def git_commit() -> str: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True, check=True, timeout=30 + ) + commit = result.stdout.strip() + if len(commit) != 40: + raise SystemExit(f"git rev-parse HEAD returned {commit!r}, expected a 40 character sha") + return commit + + +def meta(corpus_sources: int, queries: Sequence[Query], limit: int | None = None) -> dict[str, Any]: + manifest_path = SETS / "manifest.json" + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise SystemExit(f"{manifest_path}: cannot read the frozen set manifest: {error}") from error + return { + "generated_at": datetime.now(UTC).isoformat(timespec="seconds"), + "git_commit": git_commit(), + "python": platform.python_version(), + "machine": {"system": platform.system(), "release": platform.release(), "cpus": os.cpu_count()}, + "corpus_manifest_sha256": manifest.get("corpus_manifest_sha256"), + "sets": {name: sha256_file(SETS / f"{name}.jsonl") for name in MODALITIES}, + "queries": len(queries), + "answerable": sum(1 for query in queries if query.answerable), + "limit_per_modality": limit, + "corpus_sources": corpus_sources, + "models": {"dense": DENSE_MODEL, "rerank": RERANK_MODEL}, + } + + +def write_markdown(report: dict[str, Any], path: Path) -> None: + """Human-readable twin of the JSON report.""" + lines = [ + "# BASELINE: what the product retrieves today", + "", + f"- generated {report['meta']['generated_at']} at commit {report['meta']['git_commit'][:12]}", + ( + f"- corpus {report['meta']['corpus_sources']} sources, {report['meta']['queries']} queries " + f"({report['meta']['answerable']} answerable)" + ), + f"- models: dense `{report['meta']['models']['dense']}`, rerank `{report['meta']['models']['rerank']}`", + "", + "| arm | recall@1 | recall@5 | recall@10 | nDCG@10 | MRR | hit@10 (95% lower) | p50 s | p95 s |", + "|---|---|---|---|---|---|---|---|---|", + ] + chance = report["chance"] + lines.append( + f"| chance (random ranking) | {chance['recall_at_1']:.3f} | {chance['recall_at_5']:.3f} " + f"| {chance['recall_at_10']:.3f} | - | - | - | - | - |" + ) + for arm in report["arms"]: + overall = arm["overall"] + latency = arm["latency_seconds"] + lines.append( + f"| {arm['arm']} | {overall['recall_at_1']:.3f} | {overall['recall_at_5']:.3f} " + f"| {overall['recall_at_10']:.3f} | {overall['ndcg_at_10']:.3f} | {overall['mrr']:.3f} " + f"| {overall['hit_at_10']:.3f} ({overall['hit_at_10_wilson_lower']:.3f}) " + f"| {latency['p50']:.3f} | {latency['p95']:.3f} |" + ) + lines += ["", "## recall@10 by modality", "", "| arm | " + " | ".join(MODALITIES) + " |", "|---|" + "---|" * len(MODALITIES)] + for arm in report["arms"]: + cells = [f"{arm['by_modality'][modality]['recall_at_10']:.3f}" for modality in MODALITIES] + lines.append(f"| {arm['arm']} | " + " | ".join(cells) + " |") + lines += ["", "## unanswerable queries (44.8% of the suite)", ""] + for arm in report["arms"]: + overall = arm["overall"] + lines.append( + f"- {arm['arm']}: abstains on {overall['abstention_rate_unanswerable']:.0%} of unanswerable queries, " + f"and puts a wrong-holder source on top for {overall['wrong_holder_confidence_unanswerable']:.0%} of them" + ) + if "live_product" in report: + live = report["live_product"] + lines += [ + "", + "## the unmodified product path", + "", + ( + f"- RetrievalService.search: p50 {live['p50']:.2f}s, mean {live['mean']:.2f}s, max {live['max']:.2f}s " + f"over {live['queries']} queries" + ), + ( + f"- top-10 overlap with the fast BM25 path: {live['top10_overlap_with_bm25_only']:.1%}, " + f"identical order {live['identical_top10']:.1%}" + ), + ] + if report.get("scale"): + lines += ["", "## latency vs corpus size", "", "| sources | atoms | " + " | ".join(a for a in report["scale"][0] if a not in ("sources", "atoms", "queries")) + " |", "|---" * (3 + len(report["scale"][0]) - 3) + "|"] + for row in report["scale"]: + cells = [f"{row[arm]['p50']:.3f}/{row[arm]['p95']:.3f}" for arm in row if arm not in ("sources", "atoms", "queries")] + lines.append(f"| {row['sources']} | {row['atoms']} | " + " | ".join(cells) + " |") + lines += [ + "", + "## reading", + "", + "- recall counts a source as found only when it is the source a third party marked correct.", + "- abstention is the only honest number for unanswerable queries: a retriever that always", + " returns ten results never says it does not know.", + "- chance is the floor any of these arms must beat; a number near chance is not a capability.", + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def load_config(parser: argparse.ArgumentParser, args: argparse.Namespace, path: Path) -> None: + """Fill in whatever the command line left at its default; an explicit flag still wins.""" + try: + config = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise SystemExit(f"{path}: cannot read the run config ({error})") from error + if not isinstance(config, dict): + raise SystemExit(f"{path}: expected a JSON object, got {type(config).__name__}") + for key, value in config.items(): + if key not in vars(args): + raise SystemExit(f"{path}: unknown key {key!r}, expected one of {sorted(vars(args))}") + if getattr(args, key) == parser.get_default(key): + setattr(args, key, value) + unknown = [arm for arm in args.arms if arm not in ARMS] + if unknown: + raise SystemExit(f"{path}: unknown arm(s) {unknown}, expected one of {list(ARMS)}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--arms", nargs="+", choices=ARMS, default=list(ARMS)) + parser.add_argument("--modality", action="append", choices=MODALITIES, default=None) + parser.add_argument("--limit", type=int, default=None, help="score at most this many queries per modality") + parser.add_argument("--no-live", action="store_true", help="skip the unmodified RetrievalService path") + parser.add_argument("--no-scale", action="store_true", help="skip the latency/scale curve") + parser.add_argument("--out", type=Path, default=REPORTS) + parser.add_argument("--config", type=Path, help="JSON defaults for these flags (bench/configs/ablation.json)") + args = parser.parse_args() + if args.config: + load_config(parser, args, args.config) + + modalities = args.modality or list(MODALITIES) + queries = load_queries(modalities) + if args.limit: + capped: list[Query] = [] + for modality in modalities: + picked = [query for query in queries if query.modality == modality][: args.limit] + capped.extend(picked) + queries = capped + atoms = load_atoms(DB_HOME / "data" / "ely_eye.sqlite") + holders = holder_index(modalities) + corpus_sources = len({atom.source_id for atom in atoms}) + print(f"{len(queries)} queries, {len(atoms)} atoms, {corpus_sources} sources", flush=True) + + report: dict[str, Any] = { + "meta": meta(corpus_sources, queries, args.limit), + "chance": chance_row(queries, corpus_sources), + } + arms: list[dict[str, Any]] = [] + fast_rankings: dict[str, list[str]] = {} + for arm in args.arms: + print(f"== {arm}", flush=True) + runner = ArmRunner(arm, atoms, MODELS) + scores = [] + for query in queries: + ranked = runner.rank(query.text) + scores.append(score_query(query.query_id, ranked, set(query.gold), holder=query.holder, holder_of=holders)) + if arm == "bm25_only": + fast_rankings[query.query_id] = ranked + summary: dict[str, Any] = {"arm": arm, "overall": aggregate(scores)} + summary["by_modality"] = { + modality: aggregate([s for s, q in zip(scores, queries, strict=True) if q.modality == modality]) + for modality in sorted({q.modality for q in queries}) + } + summary["latency_seconds"] = { + "queries": len(runner.timings), + "mean": sum(runner.timings) / len(runner.timings), + "p50": runner.percentile(0.50), + "p95": runner.percentile(0.95), + "max": max(runner.timings), + } + arms.append(summary) + print( + f" recall@10 {summary['overall']['recall_at_10']:.3f} " + f"hit@10 {summary['overall']['hit_at_10']:.3f} " + f"(95% lower {summary['overall']['hit_at_10_wilson_lower']:.3f}) " + f"abstain {summary['overall']['abstention_rate_unanswerable']:.0%} " + f"p95 {summary['latency_seconds']['p95']:.3f}s", + flush=True, + ) + report["arms"] = arms + if not args.no_live: + report["live_product"] = live_check(queries, fast_rankings) + if not args.no_scale: + report["scale"] = scale_curve(queries, atoms, [arm for arm in args.arms if arm in ("bm25_only", "dense_only")]) + + args.out.mkdir(parents=True, exist_ok=True) + (args.out / "baseline.json").write_text(json.dumps(report, indent=1, sort_keys=True), encoding="utf-8") + write_markdown(report, args.out / "BASELINE.md") + print(f"== wrote {args.out / 'baseline.json'} and {args.out / 'BASELINE.md'}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/run_external.py b/bench/run_external.py new file mode 100644 index 0000000..a07a45b --- /dev/null +++ b/bench/run_external.py @@ -0,0 +1,161 @@ +"""Run the Mem0- and Zep-style counterparts over the frozen sets (b-4), same queries, same metrics. + +The point of this runner is an external reference point measured on our own corpus: a flat +memory store with holder scoping and no invalidation (Mem0-style), and the same items with +time-bounded validity (Zep-style). Whatever ely-eye scores, it must beat these two on the +same queries, or the memory story is not real. + +Run: python -m bench.run_external # after bench.ingest_baseline has merged the corpus +""" + +from __future__ import annotations + +import argparse +import json +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from bench.counterpart import ( + Counterpart, + build_counterparts, + build_items, + load_item_atoms, + load_source_index, +) +from bench.ingest_baseline import RUNS, database_path +from bench.metrics import aggregate, score_query +from bench.run import Query, git_commit, holder_index, load_queries +from bench.sets_common import ROOT + +REPORTS = ROOT / "reports" +TOP_K = 10 + + +def score_store(store: Counterpart, queries: list[Query], holder_of: dict[str, str]) -> dict[str, Any]: + """Rank every query against one counterpart store and aggregate the same metrics the arms use.""" + scores = [] + for query in queries: + ranked = store.search(query.text, query.holder, query.as_of, TOP_K) + scores.append(score_query(query.query_id, ranked, set(query.gold), holder=query.holder, holder_of=holder_of)) + summary: dict[str, Any] = { + "store": store.name, + "items": len(store.items), + "closed_items": sum(1 for item in store.items if item.valid_to is not None), + "overall": aggregate(scores), + "by_modality": { + modality: aggregate([s for s, q in zip(scores, queries, strict=True) if q.modality == modality]) + for modality in sorted({q.modality for q in queries}) + }, + "latency_seconds": { + "queries": len(store.timings), + "mean": sum(store.timings) / len(store.timings) if store.timings else 0.0, + "p50": store.percentile(0.50) if store.timings else 0.0, + "p95": store.percentile(0.95) if store.timings else 0.0, + "max": max(store.timings) if store.timings else 0.0, + }, + } + return summary + + +def write_markdown(report: dict[str, Any], path: Path) -> None: + """A table a reader can check: every number carries its n and its 95% lower bound.""" + meta = report["meta"] + lines = [ + "# BASELINE_EXTERNAL: memory counterparts on the same corpus", + "", + f"- generated {meta['generated_at']} at commit {meta['git_commit'][:12]}", + f"- {meta['queries']} queries ({meta['answerable']} answerable) against {meta['sources']} sources", + ( + "- extraction is lexical for both stores; Mem0's real add() and Zep's graph builder call a " + "generator model, which this round does not have" + ), + "", + ( + "| store | items | closed | recall@10 | hit@10 | hit@10 Wilson 95% lower | MRR | " + "abstain (answerable) | abstain (unanswerable) | p50 s | p95 s |" + ), + "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |", + ] + for store in report["stores"]: + overall = store["overall"] + latency = store["latency_seconds"] + lines.append( + f"| {store['store']} | {store['items']} | {store['closed_items']} | {overall['recall_at_10']:.3f} " + f"| {overall['hit_at_10']:.3f} | {overall['hit_at_10_wilson_lower']:.3f} | {overall['mrr']:.3f} " + f"| {overall['abstention_rate_answerable']:.3f} | {overall['abstention_rate_unanswerable']:.3f} " + f"| {latency['p50']:.3f} | {latency['p95']:.3f} |" + ) + lines += ["", f"n per modality: {json.dumps(meta['queries_by_modality'], sort_keys=True)}", ""] + for store in report["stores"]: + lines.append(f"## {store['store']} by modality") + lines.append("") + lines.append("| modality | n answerable | recall@10 | hit@10 | Wilson lower | MRR |") + lines.append("| --- | --- | --- | --- | --- | --- |") + for modality, row in sorted(store["by_modality"].items()): + lines.append( + f"| {modality} | {row['answerable']} | {row['recall_at_10']:.3f} | {row['hit_at_10']:.3f} " + f"| {row['hit_at_10_wilson_lower']:.3f} | {row['mrr']:.3f} |" + ) + lines.append("") + path.write_text("\n".join(lines), encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--limit", type=int, help="score at most this many queries per modality (smoke runs)") + parser.add_argument("--modality", action="append", choices=None) + parser.add_argument("--out", type=Path, default=REPORTS / "baseline_external.json") + parser.add_argument("--db", type=Path, default=database_path(RUNS), help="atoms to build items from (smoke runs use a shard)") + args = parser.parse_args() + + modalities = args.modality or None + queries = load_queries(modalities) if modalities else load_queries() + if args.limit: + kept: list[Query] = [] + for modality in sorted({q.modality for q in queries}): + kept.extend([q for q in queries if q.modality == modality][: args.limit]) + queries = kept + if not queries: + raise SystemExit("no queries selected") + + index = load_source_index() + atoms = load_item_atoms(args.db) + items = build_items(atoms, index) + if not items: + raise SystemExit("no memory items built; the corpus produced no text") + print(f"== {len(items)} memory items from {len(atoms)} atoms and {len(index)} indexed sources") + + holder_of = holder_index() + stores = build_counterparts(items) + report: dict[str, Any] = { + "meta": { + "generated_at": datetime.now(UTC).isoformat(timespec="seconds"), + "git_commit": git_commit(), + "queries": len(queries), + "answerable": sum(1 for q in queries if q.answerable), + "limit_per_modality": args.limit, + "queries_by_modality": {m: sum(1 for q in queries if q.modality == m) for m in sorted({q.modality for q in queries})}, + "sources": len(index), + "top_k": TOP_K, + "extraction": "lexical (see bench/counterpart.py); LLM extraction deferred with the generator", + }, + "stores": [score_store(store, queries, holder_of) for store in stores], + } + REPORTS.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1, sort_keys=True), encoding="utf-8") + markdown = args.out.with_name("BASELINE_EXTERNAL.md") + write_markdown(report, markdown) + for store in report["stores"]: + overall = store["overall"] + print( + f"{store['store']:14s} recall@10 {overall['recall_at_10']:.3f} hit@10 {overall['hit_at_10']:.3f} " + f"(Wilson lower {overall['hit_at_10_wilson_lower']:.3f}) mrr {overall['mrr']:.3f} " + f"p50 {store['latency_seconds']['p50']:.3f}s" + ) + print(f"== wrote {args.out} and {markdown}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/systems.py b/bench/systems.py new file mode 100644 index 0000000..9f7fa9d --- /dev/null +++ b/bench/systems.py @@ -0,0 +1,294 @@ +"""Retrieval arms for the ablation. Every arm reads the same merged baseline database. + +The arms are the product's own signals, switched on one at a time: sparse BM25, dense +cosine, the evidence graph, the product's weighted fusion, and fusion plus a cross-encoder +reranker. Rankings are lists of source ids: an atom votes for its source, and a source is +ranked by its best atom. +""" + +from __future__ import annotations + +import json +import re +import sqlite3 +import time +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + +from bench.sets_common import ROOT + +RUNS = ROOT / "runs" / "baseline" +INDEX_DIR = RUNS / "index" +MODELS = ROOT / "models" +DENSE_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" +RERANK_MODEL = "BAAI/bge-reranker-base" + + +def local_model_dir(models: Path, repo_id: str) -> Path: + """Where a fetched model lives: one flat directory per repo, named by its repo name only.""" + return models / repo_id.split("/")[-1] +DENSE_DIM = 384 +FUSION_WEIGHTS = {"sparse": 0.55, "dense": 0.35, "graph": 0.10} +RERANK_CANDIDATES = 50 +RERANK_TEXT_CHARS = 400 +ARMS = ("bm25_only", "dense_only", "graph_only", "no_rerank", "fusion_rerank") + + +@dataclass(frozen=True) +class Atom: + atom_id: str + source_id: str + text: str + + +def load_atoms(db_path: Path) -> list[Atom]: + """Every atom in the merged database, in a stable order.""" + if not db_path.is_file(): + raise SystemExit(f"{db_path}: no merged baseline database, run python -m bench.ingest_baseline first") + connection = sqlite3.connect(db_path) + try: + rows = connection.execute("SELECT atom_id, source_id, text FROM atoms ORDER BY atom_id").fetchall() + except sqlite3.OperationalError as error: + raise SystemExit(f"{db_path}: cannot read the atoms table ({error})") from error + finally: + connection.close() + if not rows: + raise SystemExit(f"{db_path}: the atoms table is empty") + return [Atom(atom_id=str(atom_id), source_id=str(source_id), text=str(text or "")) for atom_id, source_id, text in rows] + + +def tokenize(text: str) -> list[str]: + """The product's tokenizer: word characters plus CJK, lowercased.""" + return re.findall(r"[\w\u4e00-\u9fff]+", text.lower()) + + +class SparseIndex: + """BM25 over atom text, built once; the product rebuilds it on every query.""" + + def __init__(self, atoms: Sequence[Atom]) -> None: + from rank_bm25 import BM25Okapi + + corpus = [tokenize(f"{atom.text} {atom.source_id}") for atom in atoms] + self._model = BM25Okapi(corpus) + self._count = len(atoms) + + def scores(self, query: str) -> np.ndarray: + tokens = tokenize(query) + if not tokens: + return np.zeros(self._count, dtype=np.float64) + return np.asarray(self._model.get_scores(tokens), dtype=np.float64) + + +class DenseIndex: + """Brute-force cosine over cached sentence embeddings, one vector per atom.""" + + def __init__(self, atoms: Sequence[Atom], model_dir: Path, batch_size: int = 256, cache_name: str = "") -> None: + self._cache_name = cache_name or model_dir.name + self._vectors = self._load_or_build(atoms, model_dir, batch_size) + + def _vectors_path(self, model_dir: Path) -> Path: + return INDEX_DIR / f"{self._cache_name}.npy" + + def _load_or_build(self, atoms: Sequence[Atom], model_dir: Path, batch_size: int) -> np.ndarray: + path = self._vectors_path(model_dir) + order_path = INDEX_DIR / f"{self._cache_name}.atoms.json" + atom_ids = [atom.atom_id for atom in atoms] + if path.is_file() and order_path.is_file(): + vectors = np.load(path) + cached_ids = json.loads(order_path.read_text(encoding="utf-8")) + if cached_ids == atom_ids and vectors.shape == (len(atoms), DENSE_DIM): + return vectors + if cached_ids != atom_ids: + raise SystemExit( + f"{path}: cached vectors cover {len(cached_ids)} atoms but the database has {len(atoms)}; " + "delete bench/runs/baseline/index and rebuild" + ) + INDEX_DIR.mkdir(parents=True, exist_ok=True) + model = _sentence_model(model_dir) + texts = [atom.text[:RERANK_TEXT_CHARS] if atom.text else atom.source_id for atom in atoms] + vectors = model.encode( + texts, + batch_size=batch_size, + convert_to_numpy=True, + normalize_embeddings=True, + show_progress_bar=len(texts) > 2000, + ).astype(np.float32) + if vectors.shape[1] != DENSE_DIM: + raise SystemExit(f"{model_dir}: expected {DENSE_DIM} dimensions, got {vectors.shape[1]}") + np.save(path, vectors) + order_path.write_text(json.dumps(atom_ids), encoding="utf-8") + return vectors + + def scores(self, query: str, model_dir: Path) -> np.ndarray: + model = _sentence_model(model_dir) + vector = model.encode([query], normalize_embeddings=True, convert_to_numpy=True).astype(np.float32) + return np.asarray(self._vectors @ vector[0], dtype=np.float64) + + +def _sentence_model(model_dir: Path) -> Any: + """sentence-transformers ships no stubs, so the model is typed Any rather than a fake object.""" + from sentence_transformers import SentenceTransformer + + if not model_dir.is_dir(): + raise SystemExit(f"{model_dir}: model not downloaded, run python -m bench.fetch_models first") + return SentenceTransformer(str(model_dir), local_files_only=True) + + +def top_sources(scores: np.ndarray, atoms: Sequence[Atom], top_k: int) -> list[str]: + """Rank sources by their best-scoring atom; ties break on source id.""" + if len(scores) != len(atoms): + raise SystemExit(f"scores has {len(scores)} entries for {len(atoms)} atoms") + best: dict[str, float] = {} + for atom, score in zip(atoms, scores, strict=True): + current = best.get(atom.source_id) + if current is None or score > current: + best[atom.source_id] = float(score) + ranked = sorted(best.items(), key=lambda item: (-item[1], item[0])) + return [source_id for source_id, _ in ranked[:top_k]] + + +def fuse(sparse: np.ndarray, dense: np.ndarray, graph: np.ndarray) -> np.ndarray: + """The product's weighted mix, with each signal scaled to [0, 1].""" + return ( + FUSION_WEIGHTS["sparse"] * _unit(sparse) + + FUSION_WEIGHTS["dense"] * _unit(dense) + + FUSION_WEIGHTS["graph"] * _unit(graph) + ) + + +def _unit(values: np.ndarray) -> np.ndarray: + if values.size == 0: + return values + low = float(values.min()) + high = float(values.max()) + if high - low <= 1e-12: + return np.zeros_like(values) + return (values - low) / (high - low) + + +class Reranker: + """Cross-encoder over the fusion candidate list.""" + + def __init__(self, model_dir: Path) -> None: + if not model_dir.is_dir(): + raise SystemExit(f"{model_dir}: model not downloaded, run python -m bench.fetch_models first") + from sentence_transformers import CrossEncoder + + self._model = CrossEncoder(str(model_dir), local_files_only=True) + + def scores(self, query: str, texts: Sequence[str]) -> np.ndarray: + if not texts: + return np.zeros(0, dtype=np.float64) + pairs = [(query, text[:RERANK_TEXT_CHARS] or " ") for text in texts] + return np.asarray(self._model.predict(pairs, batch_size=16), dtype=np.float64) + + +def arm_needs(arm: str, models: Path) -> bool: + """Whether an arm's model files are present, so a run can skip it instead of crashing.""" + if arm not in ARMS: + raise SystemExit(f"{arm}: unknown arm, expected one of {ARMS}") + required: list[Path] = [] + if arm in ("dense_only", "no_rerank", "fusion_rerank"): + required.append(local_model_dir(models, DENSE_MODEL)) + if arm == "fusion_rerank": + required.append(local_model_dir(models, RERANK_MODEL)) + return all(directory.is_dir() for directory in required) + + +class ArmRunner: + """One arm, ranked queries in batches, with wall-clock timings recorded.""" + + def __init__(self, arm: str, atoms: Sequence[Atom], models: Path) -> None: + if arm not in ARMS: + raise SystemExit(f"{arm}: unknown arm, expected one of {ARMS}") + self.arm = arm + self.atoms = atoms + self.models = models + self.timings: list[float] = [] + self._sparse: SparseIndex | None = None + self._dense: DenseIndex | None = None + self._reranker: Reranker | None = None + if arm in ("bm25_only", "no_rerank", "fusion_rerank"): + self._sparse = SparseIndex(atoms) + if arm in ("dense_only", "no_rerank", "fusion_rerank"): + self._dense = DenseIndex(atoms, local_model_dir(models, DENSE_MODEL)) + if arm == "fusion_rerank": + self._reranker = Reranker(local_model_dir(models, RERANK_MODEL)) + + def rank(self, query: str, top_k: int = 10) -> list[str]: + started = time.perf_counter() + ranked = self._rank(query, top_k) + self.timings.append(time.perf_counter() - started) + return ranked + + def _rank(self, query: str, top_k: int) -> list[str]: + if self.arm == "graph_only": + # The product never populates atom.relations, so the graph signal is a constant. + graph = np.zeros(len(self.atoms), dtype=np.float64) + return top_sources(graph, self.atoms, top_k) + sparse = self._sparse.scores(query) if self._sparse is not None else np.zeros(len(self.atoms)) + if self.arm == "bm25_only": + return top_sources(sparse, self.atoms, top_k) + dense_dir = local_model_dir(self.models, DENSE_MODEL) + dense = ( + self._dense.scores(query, dense_dir) if self._dense is not None else np.zeros(len(self.atoms)) + ) + if self.arm == "dense_only": + return top_sources(dense, self.atoms, top_k) + graph = np.zeros(len(self.atoms), dtype=np.float64) + fused = fuse(sparse, dense, graph) + if self.arm == "no_rerank": + return top_sources(fused, self.atoms, top_k) + candidates = top_sources(fused, self.atoms, RERANK_CANDIDATES) + by_source = _best_atom_text(self.atoms, candidates) + if self._reranker is None: + raise SystemExit(f"{self.arm}: the reranker model is missing, run python -m bench.fetch_models") + rerank_scores = self._reranker.scores(query, [text for _, text in by_source]) + order = sorted(zip(candidates, rerank_scores, strict=True), key=lambda item: (-item[1], item[0])) + return [source_id for source_id, _ in order[:top_k]] + + def percentile(self, fraction: float) -> float: + if not self.timings: + raise SystemExit(f"{self.arm}: no timings recorded") + if not 0.0 <= fraction <= 1.0: + raise SystemExit(f"percentile must be within [0, 1], got {fraction}") + return float(np.percentile(np.asarray(self.timings), fraction * 100)) + + +def _best_atom_text(atoms: Sequence[Atom], source_ids: Sequence[str]) -> list[tuple[str, str]]: + longest: dict[str, str] = {} + for atom in atoms: + if atom.source_id in longest: + if len(atom.text) > len(longest[atom.source_id]): + longest[atom.source_id] = atom.text + else: + longest[atom.source_id] = atom.text + return [(source_id, longest.get(source_id, "")) for source_id in source_ids] + + +def live_product_rankings(db_home: Path, queries: Sequence[str], top_k: int = 10) -> tuple[list[list[str]], list[float]]: + """Rankings from the unmodified RetrievalService, with per-query wall time.""" + from ely_eye.config import Settings + from ely_eye.db import Database + from ely_eye.retrieval import RetrievalService + + settings = Settings(home=db_home) + service = RetrievalService(settings=settings, database=Database(settings)) + rankings: list[list[str]] = [] + timings: list[float] = [] + for query in queries: + started = time.perf_counter() + hits = service.search(query, top_k=top_k) + timings.append(time.perf_counter() - started) + # The product ranks atoms; collapse to sources by first (best) hit so the arms stay comparable. + sources: list[str] = [] + for hit in hits: + source_id = hit.atom.source_id + if source_id not in sources: + sources.append(source_id) + rankings.append(sources[:top_k]) + return rankings, timings diff --git a/bench/tests/test_audit.py b/bench/tests/test_audit.py new file mode 100644 index 0000000..d1ed36a --- /dev/null +++ b/bench/tests/test_audit.py @@ -0,0 +1,102 @@ +"""The auditor must catch planted cheating and stay quiet on honest code.""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from bench.audit import ( + best_of_violations, + frozen_counts, + injected_cases, + report_violations, + scorer_is_independent, + third_party_origins, +) + +RETRY = """ +def collect(proofs, target=0.90): + while len(proofs) < 3: + proofs.append(run_once()) + if pass_rate >= target: + break +""" + +HONEST = """ +def collect(proofs, limit): + while len(proofs) < limit: + proofs.append(next_case()) +""" + +FULL = {"meta": {"queries": 5308, "answerable": 2928}, "chance": {"recall_at_10": 0.0002}, "arms": [{"arm": "a", "overall": {"recall_at_10": 0.4, "hit_at_10": 0.6, "hit_at_10_wilson_lower": 0.5}}]} + + +def test_a_loop_that_rerolls_until_a_score_clears_is_a_violation() -> None: + found = best_of_violations(RETRY, "evals.py") + assert len(found) == 1 + assert "retry-until-score" in found[0] and "pass_rate" in found[0] + + +def test_a_plain_data_loop_is_not_flagged() -> None: + assert best_of_violations(HONEST, "runner.py") == [] + + +def test_an_explicit_best_of_name_is_flagged_wherever_it_appears() -> None: + assert best_of_violations("best_of_8 = max(scores)", "x.py") + assert best_of_violations("def until_pass(self): ...", "x.py") + + +def test_unparseable_source_fails_loudly() -> None: + with pytest.raises(SystemExit, match="cannot parse"): + best_of_violations("def broken(:", "broken.py") + + +def test_a_report_without_a_chance_row_is_refused(tmp_path: Any) -> None: + path = tmp_path / "baseline.json" + report = {k: v for k, v in FULL.items() if k != "chance"} + path.write_text(json.dumps(report), encoding="utf-8") + found = report_violations(path, frozen_counts()) + assert any("no chance row" in v for v in found) + + +def test_a_short_run_must_say_which_limit_it_used(tmp_path: Any) -> None: + path = tmp_path / "baseline.json" + path.write_text(json.dumps({**FULL, "meta": {"queries": 12, "answerable": 6}}), encoding="utf-8") + assert any("no --limit is recorded" in v for v in report_violations(path, frozen_counts())) + limited = {**FULL, "meta": {"queries": 12, "answerable": 6, "limit_per_modality": 2}} + path.write_text(json.dumps(limited), encoding="utf-8") + assert report_violations(path, frozen_counts()) == [] + + +def test_a_score_missing_its_wilson_bound_is_refused(tmp_path: Any) -> None: + path = tmp_path / "baseline.json" + report = json.loads(json.dumps(FULL)) + del report["arms"][0]["overall"]["hit_at_10_wilson_lower"] + path.write_text(json.dumps(report), encoding="utf-8") + assert any("hit_at_10_wilson_lower" in v for v in report_violations(path, frozen_counts())) + + +def test_an_unreadable_report_is_a_violation_not_a_crash(tmp_path: Any) -> None: + path = tmp_path / "baseline.json" + path.write_text("{not json", encoding="utf-8") + assert any("unreadable report" in v for v in report_violations(path, frozen_counts())) + + +def test_the_real_scorer_needs_no_model_and_the_real_gold_is_third_party() -> None: + assert scorer_is_independent() == [] + assert third_party_origins() == [] + + +def test_every_planted_case_is_caught(tmp_path: Any) -> None: + manifest = frozen_counts() + for index, (source, expected) in enumerate(injected_cases()): + if expected[0] == "retry-until-score": + found = best_of_violations(source, "injected.py") + else: + path = tmp_path / f"injected_{index}.json" + path.write_text(source, encoding="utf-8") + found = report_violations(path, manifest) + assert found, f"the auditor accepted {expected[0]}" + assert expected[0] in found[0], f"{expected[0]} was caught for the wrong reason: {found[0]}" diff --git a/bench/tests/test_counterpart.py b/bench/tests/test_counterpart.py new file mode 100644 index 0000000..55473a7 --- /dev/null +++ b/bench/tests/test_counterpart.py @@ -0,0 +1,125 @@ +"""Behaviour checks for the memory counterparts: scoping, dedup, and invalidation are the whole point.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from bench.counterpart import ( + Counterpart, + MemoryItem, + SourceMeta, + build_items, + overlap, + split_items, + subject_of, + supersede, +) +from bench.systems import Atom + +MODEL_DIR = Path("/nonexistent-model") + + +def item(text: str, source_id: str, holder: str, observed_at: str, subject: str = "") -> MemoryItem: + return MemoryItem( + item_id=f"{source_id}#0", + text=text, + source_id=source_id, + holder=holder, + scenario="assistant-memory", + observed_at=observed_at, + subject=subject, + valid_from=observed_at, + ) + + +def test_split_items_cuts_at_the_length_limit() -> None: + chunks = split_items("word " * 300) + assert chunks, "a long text must produce items" + assert all(len(chunk) <= 400 for chunk in chunks) + assert split_items(" ") == [] + + +def test_split_items_merges_short_sentences_into_one_item() -> None: + assert len(split_items("The keys are in the drawer. They are blue.")) == 1 + + +def test_subject_of_needs_a_repeated_content_token() -> None: + assert subject_of("the camera camera camera") == "camera" + assert subject_of("a single camera mention") == "" + assert subject_of("") == "" + + +def test_overlap_separates_a_restatement_from_a_new_statement() -> None: + assert overlap("the keys are in the blue drawer", "the keys are in the blue drawer") == 1.0 + assert overlap("the keys are in the blue drawer", "completely unrelated bicycle garage") < 0.5 + assert overlap("", "anything") == 0.0 + + +def test_build_items_dedups_within_a_holder_and_keeps_both_across_holders() -> None: + index = { + "src_a": SourceMeta("src_a", "ana", "assistant-memory", "conversation", "2023-01-01T00:00:00+00:00"), + "src_b": SourceMeta("src_b", "bo", "assistant-memory", "conversation", "2023-01-02T00:00:00+00:00"), + } + atoms = [ + Atom(atom_id="a1", source_id="src_a", text="The keys hang on the red door."), + Atom(atom_id="a2", source_id="src_a", text="The keys hang on the red door."), + Atom(atom_id="b1", source_id="src_b", text="The keys hang on the red door."), + ] + items = build_items(atoms, index) + assert len(items) == 2 + assert {i.holder for i in items} == {"ana", "bo"} + assert all(i.item_id.startswith(f"{i.source_id}#") for i in items) + + +def test_supersede_closes_the_older_statement_about_the_same_subject() -> None: + older = item("my camera collection is in the attic", "src_1", "ana", "2023-01-01T00:00:00+00:00", "camera") + newer = item("the bicycle is parked in the garage now", "src_2", "ana", "2023-06-01T00:00:00+00:00", "camera") + assert supersede([older, newer]) == 1 + assert older.valid_to == "2023-06-01T00:00:00+00:00" + assert newer.valid_to is None + + +def test_supersede_leaves_a_restatement_alone() -> None: + first = item("the keys hang on the red door", "src_1", "ana", "2023-01-01T00:00:00+00:00", "keys") + again = item("the keys hang on the red door too", "src_2", "ana", "2023-02-01T00:00:00+00:00", "keys") + assert supersede([first, again]) == 0 + assert first.valid_to is None + + +def test_supersede_is_scoped_to_one_holder() -> None: + ana = item("my camera collection is in the attic", "src_1", "ana", "2023-01-01T00:00:00+00:00", "camera") + bo = item("the bicycle is parked in the garage now", "src_2", "bo", "2023-06-01T00:00:00+00:00", "camera") + assert supersede([ana, bo]) == 0 + assert ana.valid_to is None + + +def test_a_query_only_sees_its_own_holders_items() -> None: + ana = item("the keys hang on the red door", "src_ana", "ana", "2023-01-01T00:00:00+00:00") + bo = item("the keys hang on the green door", "src_bo", "bo", "2023-01-01T00:00:00+00:00") + store = Counterpart(name="mem0_flat", temporal=False, items=[ana, bo], model_dir=MODEL_DIR) + mask = store.visible("ana", None) + assert [store.items[i].holder for i, seen in enumerate(mask) if seen] == ["ana"] + + +def test_the_temporal_store_hides_closed_items_and_the_flat_store_does_not() -> None: + closed = item("the keys hang on the red door", "src_old", "ana", "2023-01-01T00:00:00+00:00") + closed.valid_to = "2023-06-01T00:00:00+00:00" + open_item = item("the keys hang on the green door", "src_new", "ana", "2023-06-01T00:00:00+00:00") + items = [closed, open_item] + flat = Counterpart(name="mem0_flat", temporal=False, items=items, model_dir=MODEL_DIR) + temporal = Counterpart(name="zep_temporal", temporal=True, items=items, model_dir=MODEL_DIR) + as_of = "2023-12-01T00:00:00+00:00" + assert sum(flat.visible("ana", as_of)) == 2 + assert [store.source_id for i, store in enumerate(temporal.items) if temporal.visible("ana", as_of)[i]] == [ + "src_new" + ] + before = "2023-02-01T00:00:00+00:00" + assert sum(temporal.visible("ana", before)) == 1 + + +def test_percentile_needs_timings() -> None: + store = Counterpart(name="mem0_flat", temporal=False, items=[], model_dir=MODEL_DIR) + with pytest.raises(SystemExit, match="no timings recorded"): + store.percentile(0.5) diff --git a/bench/tests/test_metrics.py b/bench/tests/test_metrics.py new file mode 100644 index 0000000..2517e5c --- /dev/null +++ b/bench/tests/test_metrics.py @@ -0,0 +1,103 @@ +"""Hand-computed checks for every metric, so a wrong number cannot hide behind a green run.""" + +from __future__ import annotations + +import pytest + +from bench.metrics import ( + aggregate, + chance_recall_at_k, + first_hit_rank, + ndcg_at_10, + recall_at_k, + reciprocal_rank, + score_query, + wilson_lower_bound, +) + + +def test_recall_at_k_counts_distinct_gold_in_window() -> None: + assert recall_at_k(["a", "b", "c", "d"], {"c"}, 1) == 0.0 + assert recall_at_k(["a", "b", "c", "d"], {"c"}, 3) == 1.0 + assert recall_at_k(["c", "x", "a"], {"a", "c"}, 1) == 0.5 + assert recall_at_k(["c", "x", "a"], {"a", "c"}, 3) == 1.0 + assert recall_at_k(["a", "a", "b"], {"a", "b"}, 3) == 1.0 # duplicates count once + + +def test_recall_is_zero_when_gold_is_empty() -> None: + assert recall_at_k(["a"], set(), 3) == 0.0 + + +def test_reciprocal_rank_uses_first_hit_only() -> None: + assert reciprocal_rank(["x", "a", "y", "b"], {"b"}) == 0.25 + assert reciprocal_rank(["x", "a", "y", "b"], {"a", "b"}) == 0.5 + assert reciprocal_rank(["x", "y"], {"z"}) == 0.0 + assert first_hit_rank(["x", "a", "y", "b"], {"b"}) == 4 + assert first_hit_rank(["x", "y"], {"z"}) is None + + +def test_ndcg_matches_hand_computation() -> None: + # DCG = 1/log2(2) + 1/log2(4) = 1.5; IDCG = 1/log2(2) + 1/log2(3) = 1.6309297535714574 + assert ndcg_at_10(["a", "x", "b"], {"a", "b"}) == pytest.approx(1.5 / 1.6309297535714574) + assert ndcg_at_10(["a", "b"], {"a", "b"}) == pytest.approx(1.0) + assert ndcg_at_10(["x", "y"], {"a"}) == 0.0 + + +def test_ndcg_ignores_results_past_ten() -> None: + ranked = ["x"] * 10 + ["a"] + assert ndcg_at_10(ranked, {"a"}) == 0.0 + + +def test_wilson_lower_bound_matches_published_interval() -> None: + # 90/100 at 95% two-sided is [0.8257, 0.9456] in standard tables. + assert wilson_lower_bound(90, 100) == pytest.approx(0.8257, abs=1e-3) + assert wilson_lower_bound(0, 10) == 0.0 + assert wilson_lower_bound(10, 10) == pytest.approx(0.7225, abs=1e-3) + assert wilson_lower_bound(1, 10) < 0.5 + + +def test_wilson_rejects_impossible_counts() -> None: + with pytest.raises(SystemExit): + wilson_lower_bound(5, 0) + with pytest.raises(SystemExit): + wilson_lower_bound(11, 10) + + +def test_chance_recall_is_uniform_baseline() -> None: + assert chance_recall_at_k(1, 1000, 10) == pytest.approx(0.01) + assert chance_recall_at_k(4, 100, 10) == pytest.approx(0.1) + assert chance_recall_at_k(0, 100, 10) == 0.0 + with pytest.raises(SystemExit): + chance_recall_at_k(1, 0, 10) + + +def test_score_query_records_holder_and_abstention() -> None: + holder_of = {"a": "alice", "b": "bob"} + scored = score_query("q1", ["b"], {"b"}, holder="alice", holder_of=holder_of) + assert scored.gold == 1 + assert scored.recall_at_1 == 1.0 + assert scored.abstained is False + assert scored.top_hit_from_holder is False + assert score_query("q2", [], {"b"}, holder="alice", holder_of=holder_of).abstained is True + + +def test_aggregate_separates_answerable_from_unanswerable() -> None: + scores = [ + score_query("hit", ["a", "x", "y"], {"a"}, holder="alice"), + score_query("miss", ["x", "y", "a"], {"a"}, holder="alice"), + score_query("none", [], set(), holder="alice"), + ] + summary = aggregate(scores) + assert summary["n_answerable"] == 2 + assert summary["n_unanswerable"] == 1 + assert summary["recall_at_1"] == pytest.approx(0.5) + assert summary["recall_at_10"] == pytest.approx(1.0) + assert summary["mrr"] == pytest.approx((1.0 + 1 / 3) / 2) + assert summary["hit_at_10"] == 1.0 + assert summary["abstention_rate_answerable"] == 0.0 + assert summary["abstention_rate_unanswerable"] == 1.0 + + +def test_aggregate_needs_answerable_queries() -> None: + with pytest.raises(SystemExit): + aggregate([score_query("none", [], set())]) diff --git a/bench/tests/test_systems.py b/bench/tests/test_systems.py new file mode 100644 index 0000000..53059cc --- /dev/null +++ b/bench/tests/test_systems.py @@ -0,0 +1,139 @@ +"""Checks for the arm plumbing, so a broken ranking cannot hide behind a green baseline run.""" + +from __future__ import annotations + +import sqlite3 +from pathlib import Path + +import numpy as np +import pytest + +from bench.systems import ( + ARMS, + ArmRunner, + Atom, + _unit, + arm_needs, + fuse, + load_atoms, + local_model_dir, + tokenize, + top_sources, +) + +ATOMS = [ + Atom(atom_id="a1", source_id="src_red", text="red door key"), + Atom(atom_id="a2", source_id="src_blue", text="blue drawer"), + Atom(atom_id="a3", source_id="src_red", text="the key hangs on the red door"), +] + + +def test_tokenize_keeps_cjk_and_words() -> None: + assert tokenize("钥匙 Key") == ["钥匙", "key"] + assert tokenize("") == [] + + +def test_top_sources_ranks_by_the_best_atom_of_each_source() -> None: + scores = np.array([1.0, 5.0, 9.0]) + assert top_sources(scores, ATOMS, 10) == ["src_red", "src_blue"] + + +def test_top_sources_breaks_ties_on_source_id() -> None: + scores = np.array([3.0, 3.0, 3.0]) + assert top_sources(scores, ATOMS, 10) == ["src_blue", "src_red"] + + +def test_top_sources_rejects_a_length_mismatch() -> None: + with pytest.raises(SystemExit, match="scores has 2 entries for 3 atoms"): + top_sources(np.array([1.0, 2.0]), ATOMS, 3) + + +def test_unit_scales_to_zero_one_and_flattens_a_constant() -> None: + np.testing.assert_allclose(_unit(np.array([0.0, 2.0])), [0.0, 1.0]) + np.testing.assert_allclose(_unit(np.array([7.0, 7.0, 7.0])), [0.0, 0.0, 0.0]) + + +def test_fuse_uses_the_product_weights() -> None: + fused = fuse(np.array([1.0, 0.0]), np.array([0.0, 1.0]), np.array([0.0, 0.0])) + assert fused[0] == pytest.approx(0.55) + assert fused[1] == pytest.approx(0.35) + assert sum([0.55, 0.35, 0.10]) == pytest.approx(1.0) + + +def test_graph_only_arm_is_a_constant_and_sorts_by_source_id() -> None: + runner = ArmRunner("graph_only", ATOMS, Path("/nonexistent")) + assert runner.rank("key") == ["src_blue", "src_red"] + + +def test_bm25_arm_finds_the_atom_that_names_the_query() -> None: + runner = ArmRunner("bm25_only", ATOMS, Path("/nonexistent")) + assert runner.rank("the key hangs")[0] == "src_red" + # A query with no tokens has no signal: every score is zero, so the order is source id. + assert runner.rank("") == ["src_blue", "src_red"] + + +def test_arm_needs_reports_missing_models_without_crashing(tmp_path: Path) -> None: + assert arm_needs("bm25_only", tmp_path) is True + assert arm_needs("dense_only", tmp_path) is False + (local_model_dir(tmp_path, "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")).mkdir(parents=True) + assert arm_needs("dense_only", tmp_path) is True + assert arm_needs("fusion_rerank", tmp_path) is False + + +def test_arm_needs_rejects_an_unknown_arm() -> None: + with pytest.raises(SystemExit, match="unknown arm"): + arm_needs("quantum", Path("/nonexistent")) + assert "fusion_rerank" in ARMS + + +def test_local_model_dir_flattens_the_namespace() -> None: + assert local_model_dir(Path("/m"), "BAAI/bge-reranker-base") == Path("/m/bge-reranker-base") + + +def make_database(path: Path) -> None: + """Create the product's atoms table and put three rows in it.""" + connection = sqlite3.connect(path) + connection.execute( + "CREATE TABLE atoms (atom_id text PRIMARY KEY, source_id text, modality text, source text, time text," + " text text, image_ref text, embedding_refs_json text, layout_json text, relations_json text," + " trust_json text, token_equivalent integer, metadata_json text)" + ) + for atom in ATOMS: + connection.execute("INSERT INTO atoms (atom_id, source_id, text) VALUES (?, ?, ?)", (atom.atom_id, atom.source_id, atom.text)) + connection.commit() + connection.close() + + +def test_load_atoms_returns_every_row_in_a_stable_order(tmp_path: Path) -> None: + db = tmp_path / "ely_eye.sqlite" + make_database(db) + loaded = load_atoms(db) + assert [atom.atom_id for atom in loaded] == ["a1", "a2", "a3"] + assert loaded[2].text == "the key hangs on the red door" + + +def test_load_atoms_fails_fast_on_a_missing_or_unusable_database(tmp_path: Path) -> None: + with pytest.raises(SystemExit, match="no merged baseline database"): + load_atoms(tmp_path / "absent.sqlite") + no_table = tmp_path / "no_table.sqlite" + no_table.touch() + with pytest.raises(SystemExit, match="cannot read the atoms table"): + load_atoms(no_table) + empty = tmp_path / "empty.sqlite" + make_database(empty) + connection = sqlite3.connect(empty) + connection.execute("DELETE FROM atoms") + connection.commit() + connection.close() + with pytest.raises(SystemExit, match="atoms table is empty"): + load_atoms(empty) + + +def test_percentile_needs_timings_and_a_valid_fraction() -> None: + runner = ArmRunner("graph_only", ATOMS, Path("/nonexistent")) + with pytest.raises(SystemExit, match="no timings recorded"): + runner.percentile(0.5) + runner.rank("key") + with pytest.raises(SystemExit, match="percentile must be within"): + runner.percentile(1.5) + assert runner.percentile(0.5) >= 0.0 diff --git a/pyproject.toml b/pyproject.toml index 9a4ebc9..7a12311 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [tool.pyright] venvPath = "." venv = ".venv" -extraPaths = ["."] +extraPaths = [".", "backend/src"] pythonVersion = "3.12" include = ["bench", "backend/src"] @@ -10,7 +10,20 @@ python_executable = ".venv/bin/python" mypy_path = "." python_version = "3.12" files = ["bench", "backend/src"] +# bench/corpus is downloaded third-party data (git clones), not our source. +exclude = "bench/corpus/" [[tool.mypy.overrides]] module = "pyarrow.*" ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = [ + "rapidocr_onnxruntime.*", + "cv2.*", + "rank_bm25", + "sentence_transformers.*", + "onnxruntime.*", + "fitz.*", +] +ignore_missing_imports = true diff --git a/pyrightconfig.json b/pyrightconfig.json index 6f38e55..94a9bfb 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -1,7 +1,7 @@ { "venvPath": ".", "venv": ".venv", - "extraPaths": ["."], + "extraPaths": [".", "backend/src"], "include": ["bench", "backend/src"], "pythonVersion": "3.12" }