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.
This commit is contained in:
ZacharyZhang-NY committed 2026-10-06 14:36:02 -04:00
1 parent 77219c930b
commit fc0698c242
27 files changed
+2589 -24

No files matched your search

+2
View File
@@ -31,3 +31,5 @@ designreference.html
bench/corpus/*
!bench/corpus/manifest.jsonl
.pi/
bench/runs/
bench/models/
+13 -12
View File
@@ -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
---
+1 -1
View File
@@ -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",
View File
Whitespace-only changes.
+62
View File
@@ -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`.
+278
View File
@@ -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())
+1 -1
View File
@@ -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(
+1 -1
View File
@@ -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(
+3 -2
View File
@@ -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))
+1 -1
View File
@@ -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(
+1 -1
View File
@@ -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(
+1 -1
View File
@@ -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))
+5
View File
@@ -0,0 +1,5 @@
{
"arms": ["bm25_only", "dense_only", "graph_only", "no_rerank", "fusion_rerank"],
"no_live": false,
"no_scale": false
}
+321
View File
@@ -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]
+2 -2
View File
@@ -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:
+45
View File
@@ -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())
+404
View File
@@ -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())
+150
View File
@@ -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
),
}
+359
View File
@@ -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())
+161
View File
@@ -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())
+294
View File
@@ -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
+102
View File
@@ -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]}"
+125
View File
@@ -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)
+103
View File
@@ -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())])
+139
View File
@@ -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
+14 -1
View File
@@ -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
+1 -1
View File
@@ -1,7 +1,7 @@
{
"venvPath": ".",
"venv": ".venv",
"extraPaths": ["."],
"extraPaths": [".", "backend/src"],
"include": ["bench", "backend/src"],
"pythonVersion": "3.12"
}