Score all five arms in one resumable pass
The per-arm runner recomputed BM25 and the dense matmul for every arm, which put the full ablation at roughly 16 hours and left no way to survive a kill. JointRunner computes each signal once per query, attributes latency to the components each arm uses, and checkpoints every query to runs/baseline/rankings.jsonl. - bench/systems.py: JointRunner + ARM_COMPONENTS signal attribution. - bench/run.py: joint pass, rankings cache with a run fingerprint, per-query resume, latency summary, resident-memory progress lines. - bench/tests/test_joint.py: joint rankings equal the per-arm rankings for all five arms; the sparse corpus is scored once even with three arms that need it; cache fingerprint mismatch and malformed rows fail fast. - bench/reports/EXTERNAL_APPENDIX.md: what a real Mem0/Zep run would cost, and which published numbers are comparable at all. Verification: pytest bench/tests -q -> 57 passed; ruff check bench/ -> all checks passed; mypy (fresh cache) -> no issues in 25 source files.
This commit is contained in:
1 parent
a4953335e7
commit
29163632af
6 files changed
+356
-116
No files matched your search
@@ -1,27 +0,0 @@
|
||||
# BASELINE: what the product retrieves today
|
||||
|
||||
- generated 2026-10-06T20:42:42+00:00 at commit fc0698c2426c
|
||||
- corpus 50278 sources, 10 queries (8 answerable)
|
||||
- models: dense `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`, rerank `BAAI/bge-reranker-base`
|
||||
|
||||
| arm | recall@1 | recall@5 | recall@10 | nDCG@10 | MRR | hit@10 (95% lower) | p50 s | p95 s |
|
||||
|---|---|---|---|---|---|---|---|---|
|
||||
| chance (random ranking) | 0.000 | 0.000 | 0.000 | - | - | - | - | - |
|
||||
| bm25_only | 0.625 | 0.875 | 1.000 | 0.788 | 0.725 | 1.000 (0.676) | 0.881 | 1.207 |
|
||||
|
||||
## recall@10 by modality
|
||||
|
||||
| arm | pdf |
|
||||
|---|---|
|
||||
| bm25_only | 1.000 |
|
||||
|
||||
## unanswerable queries (44.8% of the suite)
|
||||
|
||||
- bm25_only: abstains on 0% of unanswerable queries, and puts a wrong-holder source on top for 100% of them
|
||||
|
||||
## 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.
|
||||
@@ -0,0 +1,65 @@
|
||||
# Appendix: what a real Mem0 run would cost, and why its numbers are not ours
|
||||
|
||||
`bench/counterpart.py` rebuilds the store-and-retrieve half of Mem0 and Zep inside this repo so the
|
||||
comparison runs on one corpus, one machine, one scorer. The products themselves are not run here.
|
||||
This appendix records what running them would take, and which of their published numbers can be
|
||||
placed next to ours at all.
|
||||
|
||||
## what the counterpart reproduces, and what it does not
|
||||
|
||||
| behaviour | Mem0 / Zep paper | `bench/counterpart.py` |
|
||||
|---|---|---|
|
||||
| store granularity | LLM-extracted facts, deduplicated against existing facts | deterministic 400-character sentence chunks, deduplicated by `(holder, sha256(text))` |
|
||||
| scope | `user_id` / `agent_id` / `run_id` | `holder` from the frozen corpus |
|
||||
| invalidation | Zep closes an edge when a newer fact contradicts an older one | `supersede` closes the older item when a newer one in the same `(holder, subject)` group overlaps it by less than 0.5 |
|
||||
| retrieval | dense + sparse + graph signals, `top_200` budget in the 2026 hosted stack | dense cosine over the same 384-dimensional embedder the ely-eye dense arm uses, `top_k=10` |
|
||||
| scoring | LLM-as-a-judge on a generated answer | source-level recall@10 against frozen gold |
|
||||
|
||||
The extraction step is the part that cannot be made comparable: Mem0's `add()` calls an LLM per
|
||||
ingest, so its store contents depend on the model, the prompt, and the sampling temperature. Running
|
||||
it against this corpus would add an LLM to the ingest path and make the memory store itself
|
||||
non-reproducible, while leaving the retrieval question untouched. That is why the reference point is
|
||||
in-repo.
|
||||
|
||||
## published numbers worth keeping
|
||||
|
||||
From the Mem0 paper (arXiv 2504.19413v1, LOCOMO, overall LLM-as-a-judge score, mean ± std over 10 runs):
|
||||
|
||||
| system | overall judge score | tokens retrieved per query | search p50 / p95 (s) | total p50 / p95 (s) |
|
||||
|---|---|---|---|---|
|
||||
| full-context | 72.90 ± 0.19 | 26,031 | — | 9.870 / 17.117 |
|
||||
| Mem0 (graph variant) | 68.44 ± 0.17 | 3,616 | 0.476 / 0.657 | 1.091 / 2.590 |
|
||||
| Mem0 | 66.88 ± 0.15 | 1,764 | 0.148 / 0.200 | 0.708 / 1.440 |
|
||||
| Zep | 65.99 ± 0.16 | 3,911 | 0.513 / 0.778 | 1.292 / 2.926 |
|
||||
| best RAG configuration | 60.97 ± 0.20 | k × chunk size | 0.267 / 0.624 | 0.766 / 1.829 |
|
||||
| LangMem | 58.10 ± 0.21 | 127 | 17.99 / 59.82 | 18.53 / 60.40 |
|
||||
| A-Mem | 48.38 ± 0.15 | 2,520 | 0.668 / 1.485 | 1.410 / 4.374 |
|
||||
|
||||
Two things in that table matter more than the scores:
|
||||
|
||||
- **full-context wins the answer-quality column and loses the cost column.** 72.90 at 26,031 tokens
|
||||
per query, p95 17.1 s. Every memory system trades quality for cost. Any claim that a memory layer
|
||||
"beats the context window" has to say which column it means.
|
||||
- **the store-construction cost is not in the table at all.** The paper tabulates search and answer
|
||||
latency, not `add()`. It reports separately that Zep's construction runs multiple asynchronous LLM
|
||||
calls, so retrieval immediately after an add often fails and improves hours later. A benchmark that
|
||||
scores retrieval only, on a store built in advance, hides that entirely.
|
||||
|
||||
Sources: arXiv 2504.19413v1 (table 2 and sections 4.1-4.5), the mem0 repository README's 2026
|
||||
algorithm table (LoCoMo 92.5, LongMemEval 94.4, ~7.0k tokens, p50 0.88-1.09 s, single-pass retrieval
|
||||
at a `top_200` budget, managed platform), and `docs.mem0.ai/core-concepts/memory-evaluation`.
|
||||
|
||||
## why these numbers cannot sit in our ablation table
|
||||
|
||||
1. Different metric. Their column is judge-scored answer correctness on 300 LOCOMO questions. Ours is
|
||||
recall@10 of the gold **source** over 2,928 answerable queries drawn from six third-party sets.
|
||||
A system can rank the right passage last and still answer correctly from a paraphrase.
|
||||
2. Different budget. They retrieve `top_200` items; we score `top_k=10`.
|
||||
3. Different hardware and model. Their run used a hosted production stack; ours is CPU-only with a
|
||||
384-dimensional multilingual MiniLM, no CUDA in this WSL image.
|
||||
4. Different corpus age. LOCOMO has 10 dialogues; our conversation set is LongMemEval-cleaned with
|
||||
7,188 session files and 854 holders, and it also carries five non-conversational modalities that
|
||||
no memory product is evaluated on here.
|
||||
|
||||
The honest use of this appendix is the ordering it implies — full-context > memory systems > naive
|
||||
RAG on answer quality, and the reverse on cost — plus the two cost columns the paper leaves out.
|
||||
@@ -1,77 +0,0 @@
|
||||
{
|
||||
"arms": [
|
||||
{
|
||||
"arm": "bm25_only",
|
||||
"by_modality": {
|
||||
"pdf": {
|
||||
"abstention_rate_answerable": 0.0,
|
||||
"abstention_rate_unanswerable": 0.0,
|
||||
"hit_at_10": 1.0,
|
||||
"hit_at_10_wilson_lower": 0.6755924351161198,
|
||||
"mrr": 0.725,
|
||||
"n_answerable": 8,
|
||||
"n_unanswerable": 2,
|
||||
"ndcg_at_10": 0.7883559233904859,
|
||||
"recall_at_1": 0.625,
|
||||
"recall_at_10": 1.0,
|
||||
"recall_at_5": 0.875,
|
||||
"wrong_holder_confidence_unanswerable": 1.0
|
||||
}
|
||||
},
|
||||
"latency_seconds": {
|
||||
"max": 1.3251221780083142,
|
||||
"mean": 0.8729710359009915,
|
||||
"p50": 0.8812561564991483,
|
||||
"p95": 1.2068847050097244,
|
||||
"queries": 10
|
||||
},
|
||||
"overall": {
|
||||
"abstention_rate_answerable": 0.0,
|
||||
"abstention_rate_unanswerable": 0.0,
|
||||
"hit_at_10": 1.0,
|
||||
"hit_at_10_wilson_lower": 0.6755924351161198,
|
||||
"mrr": 0.725,
|
||||
"n_answerable": 8,
|
||||
"n_unanswerable": 2,
|
||||
"ndcg_at_10": 0.7883559233904859,
|
||||
"recall_at_1": 0.625,
|
||||
"recall_at_10": 1.0,
|
||||
"recall_at_5": 0.875,
|
||||
"wrong_holder_confidence_unanswerable": 1.0
|
||||
}
|
||||
}
|
||||
],
|
||||
"chance": {
|
||||
"corpus_sources": 50278,
|
||||
"recall_at_1": 1.9889414853415013e-05,
|
||||
"recall_at_10": 0.00019889414853415012,
|
||||
"recall_at_5": 9.944707426707506e-05
|
||||
},
|
||||
"meta": {
|
||||
"answerable": 8,
|
||||
"corpus_manifest_sha256": "8fc1ed861c891d778853d41107630869f9eb267a669c5d1038e389de80114347",
|
||||
"corpus_sources": 50278,
|
||||
"generated_at": "2026-10-06T20:42:42+00:00",
|
||||
"git_commit": "fc0698c2426c373e32c361a816523351a1e3fcb7",
|
||||
"limit_per_modality": 10,
|
||||
"machine": {
|
||||
"cpus": 32,
|
||||
"release": "6.18.40.1-microsoft-standard-WSL2",
|
||||
"system": "Linux"
|
||||
},
|
||||
"models": {
|
||||
"dense": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
|
||||
"rerank": "BAAI/bge-reranker-base"
|
||||
},
|
||||
"python": "3.14.4",
|
||||
"queries": 10,
|
||||
"sets": {
|
||||
"code": "ea7e248386350f95132a398f3710f6ce5f09424ca614fd88c618207f667a553e",
|
||||
"conversation": "dc106d90e7749bee8be01b1fab3b645a5459d7527ef40e7fbf983677d5898d11",
|
||||
"image": "1c74488375c7ec9f3182964d51a5e506273bfd5cf3d126725c3cc431a8333128",
|
||||
"pdf": "0cec488e660bf1bfa7ccfc72dbae8b7e3005f0b6a0255028e502296d4ecd4536",
|
||||
"ui_screenshot": "2269dd623e67cd851fb3f6b24a4556a8462a37db98e12ad79577c5830167c7c5",
|
||||
"video": "4973231d9f16d214b14923d79a34b85c3d2b423624dea752b1ee91ba09ff778c"
|
||||
}
|
||||
}
|
||||
}
|
||||
+101
-12
@@ -10,6 +10,7 @@ Run: python -m bench.run [--arms bm25_only] [--limit 200] [--no-live] [--out ben
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
@@ -20,6 +21,8 @@ from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
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 (
|
||||
@@ -29,6 +32,7 @@ from bench.systems import (
|
||||
RERANK_MODEL,
|
||||
ArmRunner,
|
||||
Atom,
|
||||
JointRunner,
|
||||
arm_needs,
|
||||
live_product_rankings,
|
||||
load_atoms,
|
||||
@@ -38,6 +42,7 @@ ROOT = SETS.parent
|
||||
REPORTS = ROOT / "reports"
|
||||
DB_HOME = ROOT / "runs" / "baseline"
|
||||
LIVE_QUERIES = 40 # the product rebuilds BM25 per query, so a live pass costs minutes per query
|
||||
PROGRESS_EVERY = 250 # an arm takes tens of minutes over 5,308 queries; silence looks like a hang
|
||||
SCALE_SIZES = (500, 2_000, 10_000, 50_000)
|
||||
|
||||
|
||||
@@ -291,6 +296,66 @@ def load_config(parser: argparse.ArgumentParser, args: argparse.Namespace, path:
|
||||
raise SystemExit(f"{path}: unknown arm(s) {unknown}, expected one of {list(ARMS)}")
|
||||
|
||||
|
||||
def latency(times: Sequence[float]) -> dict[str, float]:
|
||||
"""Latency summary for one arm's per-query wall times."""
|
||||
values = np.asarray(times, dtype=np.float64)
|
||||
if values.size == 0:
|
||||
raise SystemExit("no latency samples recorded for this arm")
|
||||
return {
|
||||
"queries": int(values.size),
|
||||
"mean": float(values.mean()),
|
||||
"p50": float(np.percentile(values, 50)),
|
||||
"p95": float(np.percentile(values, 95)),
|
||||
"max": float(values.max()),
|
||||
}
|
||||
|
||||
|
||||
def rss_gib() -> float:
|
||||
"""Resident memory of this process; a nine-hour run needs to show its growth."""
|
||||
try:
|
||||
for line in Path("/proc/self/status").read_text(encoding="utf-8").splitlines():
|
||||
if line.startswith("VmRSS:"):
|
||||
return int(line.split()[1]) / 1024 / 1024
|
||||
except OSError:
|
||||
return 0.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def fingerprint(arms: Sequence[str], queries: Sequence[Query], atoms: Sequence[Atom]) -> str:
|
||||
"""Identity of a rankings cache: the arms, the query set, and the corpus behind them."""
|
||||
corpus = hashlib.sha256("\n".join(atom.atom_id for atom in atoms).encode("utf-8")).hexdigest()[:16]
|
||||
payload = "\n".join(
|
||||
[",".join(arms), ",".join(query.query_id for query in queries), str(len(atoms)), corpus]
|
||||
)
|
||||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def load_rankings(path: Path, mark: str) -> dict[str, dict[str, Any]]:
|
||||
"""Read a resumable rankings cache; a header mismatch is an error, not a silent restart."""
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
lines = path.read_text(encoding="utf-8").splitlines()
|
||||
except OSError as error:
|
||||
raise SystemExit(f"{path}: cannot read the rankings cache ({error})") from error
|
||||
rows: dict[str, dict[str, Any]] = {}
|
||||
for index, line in enumerate(lines, start=1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError as error:
|
||||
raise SystemExit(f"{path}:{index}: not JSON ({error})") from error
|
||||
if index == 1:
|
||||
if row.get("fingerprint") != mark:
|
||||
raise SystemExit(
|
||||
f"{path}: cached rankings belong to a different run, delete the file to start fresh"
|
||||
)
|
||||
continue
|
||||
rows[str(row["query_id"])] = dict(row["arms"])
|
||||
return rows
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--arms", nargs="+", choices=ARMS, action="extend", default=None)
|
||||
@@ -325,14 +390,42 @@ def main() -> int:
|
||||
"meta": meta(corpus_sources, queries, args.limit),
|
||||
"chance": chance_row(queries, corpus_sources),
|
||||
}
|
||||
# One pass over the queries: each signal is computed once and every arm is derived
|
||||
# from it, and each query is checkpointed so a kill costs one query, not an arm.
|
||||
mark = fingerprint(arms_requested, queries, atoms)
|
||||
rankings_path = DB_HOME / "rankings.jsonl"
|
||||
done = load_rankings(rankings_path, mark)
|
||||
joint = JointRunner(arms_requested, atoms, MODELS)
|
||||
resumed = 0
|
||||
with rankings_path.open("a", encoding="utf-8") as cache:
|
||||
if not done:
|
||||
cache.write(f"{json.dumps({'fingerprint': mark})}\n")
|
||||
for index, query in enumerate(queries, start=1):
|
||||
if query.query_id in done:
|
||||
resumed += 1
|
||||
continue
|
||||
per_arm = joint.rank_all(query.text)
|
||||
done[query.query_id] = {
|
||||
arm: {"ranked": order, "seconds": seconds} for arm, (order, seconds) in per_arm.items()
|
||||
}
|
||||
cache.write(f"{json.dumps({'query_id': query.query_id, 'arms': done[query.query_id]})}\n")
|
||||
if index % PROGRESS_EVERY == 0:
|
||||
cache.flush()
|
||||
print(f" {index}/{len(queries)} queries, rss {rss_gib():.1f} GiB", flush=True)
|
||||
if resumed:
|
||||
print(f" resumed {resumed} queries from {rankings_path}", flush=True)
|
||||
|
||||
arms: list[dict[str, Any]] = []
|
||||
fast_rankings: dict[str, list[str]] = {}
|
||||
for arm in arms_requested:
|
||||
print(f"== {arm}", flush=True)
|
||||
runner = ArmRunner(arm, atoms, MODELS)
|
||||
for arm in [arm for arm in ARMS if arm in arms_requested]:
|
||||
scores = []
|
||||
times: list[float] = []
|
||||
for query in queries:
|
||||
ranked = runner.rank(query.text)
|
||||
record = done.get(query.query_id)
|
||||
if record is None or arm not in record:
|
||||
raise SystemExit(f"{rankings_path}: no cached ranking for {query.query_id} arm {arm}")
|
||||
times.append(float(record[arm]["seconds"]))
|
||||
ranked = [str(source_id) for source_id in record[arm]["ranked"]]
|
||||
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
|
||||
@@ -341,16 +434,10 @@ def main() -> int:
|
||||
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),
|
||||
}
|
||||
summary["latency_seconds"] = latency(times)
|
||||
arms.append(summary)
|
||||
print(
|
||||
f" recall@10 {summary['overall']['recall_at_10']:.3f} "
|
||||
f"== {arm} 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%} "
|
||||
@@ -358,6 +445,8 @@ def main() -> int:
|
||||
flush=True,
|
||||
)
|
||||
report["arms"] = arms
|
||||
# A crash in the live or scale phase should not cost every finished arm.
|
||||
(DB_HOME / "partial_report.json").write_text(json.dumps(report, indent=1, sort_keys=True), encoding="utf-8")
|
||||
if not args.no_live:
|
||||
report["live_product"] = live_check(queries, fast_rankings, live_queries=args.live_queries)
|
||||
if not args.no_scale:
|
||||
|
||||
@@ -286,6 +286,97 @@ class ArmRunner:
|
||||
return float(np.percentile(np.asarray(self.timings), fraction * 100))
|
||||
|
||||
|
||||
# Which measured signals each arm consumes; JointRunner times each signal once and
|
||||
# attributes it to every arm that uses it.
|
||||
ARM_COMPONENTS: dict[str, tuple[str, ...]] = {
|
||||
"bm25_only": ("sparse",),
|
||||
"dense_only": ("dense",),
|
||||
"graph_only": ("graph",),
|
||||
"no_rerank": ("sparse", "dense", "graph", "fuse"),
|
||||
"fusion_rerank": ("sparse", "dense", "graph", "fuse"),
|
||||
}
|
||||
|
||||
|
||||
class JointRunner:
|
||||
"""Every requested arm over one corpus, computing each signal once per query.
|
||||
|
||||
ArmRunner recomputes the sparse and dense signals per arm; over 166K atoms that is
|
||||
hours of duplicated work. This pass computes each signal once and derives every
|
||||
arm's ranking from it, so a ranking matches ArmRunner's exactly while latency is
|
||||
attributed to the components that arm actually uses.
|
||||
"""
|
||||
|
||||
def __init__(self, arms: Sequence[str], atoms: Sequence[Atom], models: Path) -> None:
|
||||
for arm in arms:
|
||||
if arm not in ARMS:
|
||||
raise SystemExit(f"{arm}: unknown arm, expected one of {ARMS}")
|
||||
if not arms:
|
||||
raise SystemExit("JointRunner needs at least one arm")
|
||||
self.arms = [arm for arm in ARMS if arm in arms]
|
||||
self.atoms = atoms
|
||||
self.models = models
|
||||
self.timings: dict[str, list[float]] = {arm: [] for arm in self.arms}
|
||||
self._sparse: SparseIndex | None = None
|
||||
self._dense: DenseIndex | None = None
|
||||
self._reranker: Reranker | None = None
|
||||
if any(arm in ("bm25_only", "no_rerank", "fusion_rerank") for arm in self.arms):
|
||||
self._sparse = SparseIndex(atoms)
|
||||
if any(arm in ("dense_only", "no_rerank", "fusion_rerank") for arm in self.arms):
|
||||
self._dense = DenseIndex(atoms, local_model_dir(models, DENSE_MODEL))
|
||||
if "fusion_rerank" in self.arms:
|
||||
self._reranker = Reranker(local_model_dir(models, RERANK_MODEL))
|
||||
|
||||
def rank_all(self, query: str, top_k: int = 10) -> dict[str, tuple[list[str], float]]:
|
||||
"""Rank one query with every arm; returns {arm: (source_ids, seconds)}."""
|
||||
signals: dict[str, np.ndarray] = {}
|
||||
seconds: dict[str, float] = {}
|
||||
started = time.perf_counter()
|
||||
signals["graph"] = np.zeros(len(self.atoms), dtype=np.float64)
|
||||
seconds["graph"] = time.perf_counter() - started
|
||||
if self._sparse is not None:
|
||||
started = time.perf_counter()
|
||||
signals["sparse"] = self._sparse.scores(query)
|
||||
seconds["sparse"] = time.perf_counter() - started
|
||||
if self._dense is not None:
|
||||
started = time.perf_counter()
|
||||
signals["dense"] = self._dense.scores(query, local_model_dir(self.models, DENSE_MODEL))
|
||||
seconds["dense"] = time.perf_counter() - started
|
||||
if any(arm in ("no_rerank", "fusion_rerank") for arm in self.arms):
|
||||
started = time.perf_counter()
|
||||
signals["fuse"] = fuse(
|
||||
signals.get("sparse", np.zeros(len(self.atoms))),
|
||||
signals.get("dense", np.zeros(len(self.atoms))),
|
||||
signals["graph"],
|
||||
)
|
||||
seconds["fuse"] = time.perf_counter() - started
|
||||
ranked: dict[str, tuple[list[str], float]] = {}
|
||||
for arm in self.arms:
|
||||
started = time.perf_counter()
|
||||
order = self._order(arm, query, signals, top_k)
|
||||
total = sum(seconds[name] for name in ARM_COMPONENTS[arm]) + (time.perf_counter() - started)
|
||||
self.timings[arm].append(total)
|
||||
ranked[arm] = (order, total)
|
||||
return ranked
|
||||
|
||||
def _order(self, arm: str, query: str, signals: dict[str, np.ndarray], top_k: int) -> list[str]:
|
||||
if arm == "graph_only":
|
||||
return top_sources(signals["graph"], self.atoms, top_k)
|
||||
if arm == "bm25_only":
|
||||
return top_sources(signals["sparse"], self.atoms, top_k)
|
||||
if arm == "dense_only":
|
||||
return top_sources(signals["dense"], self.atoms, top_k)
|
||||
fused = signals["fuse"]
|
||||
if 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"{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 _best_atom_text(atoms: Sequence[Atom], source_ids: Sequence[str]) -> list[tuple[str, str]]:
|
||||
longest: dict[str, str] = {}
|
||||
for atom in atoms:
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
"""JointRunner must rank exactly like ArmRunner, and the rankings cache must resume.
|
||||
|
||||
The ablation switched to a single pass that computes each signal once. If that pass
|
||||
drifted from the per-arm path, the baseline would silently measure a different system.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from bench import systems
|
||||
from bench.run import Query, fingerprint, latency, load_rankings
|
||||
from bench.systems import ARMS, ArmRunner, Atom, JointRunner
|
||||
from bench.tests.test_systems import ATOMS
|
||||
|
||||
MODEL_ARMS = ["bm25_only", "graph_only", "dense_only", "no_rerank", "fusion_rerank"]
|
||||
|
||||
|
||||
def _queries() -> list[str]:
|
||||
return ["alpha report", "beta invoice", "gamma", ""]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arm", MODEL_ARMS)
|
||||
def test_joint_ranking_matches_the_per_arm_runner(arm: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
|
||||
queries = _queries()
|
||||
joint = JointRunner([arm], ATOMS, systems.MODELS)
|
||||
per_arm = ArmRunner(arm, ATOMS, systems.MODELS)
|
||||
for query in queries:
|
||||
expected = per_arm.rank(query)
|
||||
ranked, seconds = joint.rank_all(query)[arm]
|
||||
assert ranked == expected, f"{arm} drifted from the per-arm path on {query!r}"
|
||||
assert seconds >= 0.0
|
||||
|
||||
|
||||
def test_joint_runner_computes_the_sparse_signal_once_per_query(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
|
||||
joint = JointRunner(["bm25_only", "no_rerank", "fusion_rerank"], ATOMS, systems.MODELS)
|
||||
calls = 0
|
||||
assert joint._sparse is not None
|
||||
original = joint._sparse.scores
|
||||
|
||||
def counting_scores(query: str):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return original(query)
|
||||
|
||||
monkeypatch.setattr(joint._sparse, "scores", counting_scores)
|
||||
joint.rank_all("alpha report")
|
||||
assert calls == 1, "three arms that all need BM25 must still score the corpus once"
|
||||
|
||||
|
||||
def test_joint_runner_rejects_unknown_or_empty_arm_lists() -> None:
|
||||
with pytest.raises(SystemExit, match="unknown arm"):
|
||||
JointRunner(["vibes"], ATOMS, systems.MODELS)
|
||||
with pytest.raises(SystemExit, match="at least one arm"):
|
||||
JointRunner([], ATOMS, systems.MODELS)
|
||||
|
||||
|
||||
def test_rankings_cache_resumes_and_refuses_a_mismatched_run(tmp_path: Path) -> None:
|
||||
atoms = [Atom(atom_id=f"a{index}", source_id=f"s{index}", text=f"text {index}") for index in range(3)]
|
||||
queries = [
|
||||
Query(
|
||||
query_id=f"q{index}",
|
||||
text="t",
|
||||
modality="pdf",
|
||||
gold=frozenset({"s0"}),
|
||||
holder="h",
|
||||
answerable=True,
|
||||
as_of=None,
|
||||
)
|
||||
for index in range(2)
|
||||
]
|
||||
mark = fingerprint(["bm25_only"], queries, atoms)
|
||||
path = tmp_path / "rankings.jsonl"
|
||||
assert load_rankings(path, mark) == {}
|
||||
rows = [json.dumps({"fingerprint": mark}), json.dumps({"query_id": "q0", "arms": {"bm25_only": {"ranked": ["s0"]}}})]
|
||||
path.write_text("\n".join(rows) + "\n", encoding="utf-8")
|
||||
assert load_rankings(path, mark) == {"q0": {"bm25_only": {"ranked": ["s0"]}}}
|
||||
with pytest.raises(SystemExit, match="different run"):
|
||||
load_rankings(path, "other")
|
||||
path.write_text("\n".join([rows[0], "{not json"]) + "\n", encoding="utf-8")
|
||||
with pytest.raises(SystemExit, match="not JSON"):
|
||||
load_rankings(path, mark)
|
||||
|
||||
|
||||
def test_latency_needs_samples() -> None:
|
||||
assert latency([0.5, 1.0])["p50"] == pytest.approx(0.75)
|
||||
with pytest.raises(SystemExit, match="no latency samples"):
|
||||
latency([])
|
||||
|
||||
|
||||
def test_every_arm_declares_the_signals_it_consumes() -> None:
|
||||
assert set(systems.ARM_COMPONENTS) == set(ARMS)
|
||||
Reference in new issue
Block a user