Measure the counterpart honestly: warm models, vectorised search, scope split

Three defects made the external baseline numbers and labels wrong.

- bench/systems.py: _sentence_model rebuilt the model on every query, so dense arm
  latency was mostly model loading and the latency-vs-size curve was flat because a
  constant dominated it. Models are now cached per process.
- bench/counterpart.py: visible() and search() walked 404k Python objects per query.
  Per-item holder/source/time columns are built once and the query path is numpy;
  12.5 queries/s instead of 0.7.
- bench/metrics.py: the aggregate key was named wrong_holder_confidence_unanswerable
  while it counts top hits that ARE from the asking holder. BASELINE.md read that
  backwards. Renamed to top_hit_from_holder_unanswerable and the sentence rewritten.
- bench/run_external.py: the by-modality table read row['answerable'] but aggregate
  emits n_answerable, so the writer died after scoring finished. Each store is now
  scored twice, [holder scope] and [whole corpus], so the gain from the holder
  boundary is separable from the gain from ranking; progress lines every 500 queries.

Verification: pytest bench/tests -q -> 61 passed; ruff check bench/ -> all checks
passed; mypy (fresh cache) -> no issues in 25 source files.
This commit is contained in:
ZacharyZhang-NY committed 2026-10-07 13:33:13 -04:00
1 parent c4e2982506
commit dbe1d2cc82
8 files changed
+381 -34

No files matched your search

+50 -17
View File
@@ -241,6 +241,18 @@ def supersede(items: Sequence[MemoryItem]) -> int:
return closed
@dataclass(frozen=True)
class _ItemArrays:
"""Per-item columns, built once so a query never walks 400k Python objects."""
holder_ids: dict[str, int]
holders: np.ndarray
sources: list[str]
sources_of_item: np.ndarray
valid_from: np.ndarray
valid_to: np.ndarray
@dataclass
class Counterpart:
"""One counterpart store: the same items, the same embeddings, different memory semantics."""
@@ -252,10 +264,26 @@ class Counterpart:
cache_name: str = "counterpart-items"
timings: list[float] = field(default_factory=list)
_index: DenseIndex | None = None
_arrays: _ItemArrays | None = None
def __post_init__(self) -> None:
"""The store is cheap to build; embeddings are paid for on the first search."""
def _ensure_arrays(self) -> _ItemArrays:
if self._arrays is None:
holder_ids = {holder: code for code, holder in enumerate(sorted({item.holder for item in self.items}))}
ordered_sources = sorted({item.source_id for item in self.items})
source_ids = {source: code for code, source in enumerate(ordered_sources)}
self._arrays = _ItemArrays(
holder_ids=holder_ids,
holders=np.array([holder_ids[item.holder] for item in self.items], dtype=np.int32),
sources=list(ordered_sources),
sources_of_item=np.array([source_ids[item.source_id] for item in self.items], dtype=np.int32),
valid_from=np.array([item.valid_from for item in self.items], dtype="U64"),
valid_to=np.array([item.valid_to or "" for item in self.items], dtype="U64"),
)
return self._arrays
def _ensure_index(self) -> DenseIndex:
if self._index is None:
self._index = DenseIndex(
@@ -272,15 +300,14 @@ class Counterpart:
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
arrays = self._ensure_arrays()
if holder is None:
mask = np.ones(len(self.items), dtype=bool)
else:
mask = arrays.holders == arrays.holder_ids.get(holder, -1)
if self.temporal and as_of is not None:
mask &= arrays.valid_from <= as_of
mask &= (arrays.valid_to == "") | (arrays.valid_to > as_of)
return mask
def search(self, query: str, holder: str | None, as_of: str | None, top_k: int = 10) -> list[str]:
@@ -288,16 +315,22 @@ class Counterpart:
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]:
arrays = self._ensure_arrays()
masked = np.where(mask, scores, -np.inf)
ranked: list[str] = []
seen: set[int] = set()
for position in np.argsort(-masked, kind="stable"):
if not np.isfinite(masked[position]):
break
code = int(arrays.sources_of_item[position])
if code in seen:
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]))
seen.add(code)
ranked.append(arrays.sources[code])
if len(ranked) == top_k:
break
self.timings.append(time.perf_counter() - started)
return [source_id for source_id, _ in ranked[:top_k]]
return ranked
def percentile(self, fraction: float) -> float:
if not self.timings:
+1 -1
View File
@@ -142,7 +142,7 @@ def aggregate(scores: Sequence[QueryScore]) -> dict[str, float | int]:
"abstention_rate_unanswerable": (
sum(1 for s in unanswerable if s.abstained) / len(unanswerable) if unanswerable else 0.0
),
"wrong_holder_confidence_unanswerable": (
"top_hit_from_holder_unanswerable": (
sum(1 for s in holder_known if s.top_hit_from_holder) / len(holder_known)
if holder_known
else 0.0
+256
View File
@@ -0,0 +1,256 @@
{
"meta": {
"answerable": 2928,
"extraction": "lexical (see bench/counterpart.py); LLM extraction deferred with the generator",
"generated_at": "2026-10-07T16:54:23+00:00",
"git_commit": "c4e2982506faaf350a4b06738923899fc44bdedd",
"limit_per_modality": null,
"machine": {
"cpus": 32,
"load_average_1_5_15": [
6.23,
4.72,
4.44
],
"release": "6.18.40.1-microsoft-standard-WSL2",
"system": "Linux"
},
"queries": 5308,
"queries_by_modality": {
"code": 2474,
"conversation": 300,
"image": 480,
"pdf": 984,
"ui_screenshot": 640,
"video": 430
},
"sources": 50789,
"top_k": 10
},
"stores": [
{
"by_modality": {
"code": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.952303961196443,
"hit_at_10_wilson_lower": 0.938963104367135,
"mrr": 0.6949699092787209,
"n_answerable": 1237,
"n_unanswerable": 1237,
"ndcg_at_10": 0.7581529489950144,
"recall_at_1": 0.549717057396928,
"recall_at_10": 0.952303961196443,
"recall_at_5": 0.8884397736459175,
"wrong_holder_confidence_unanswerable": 1.0
},
"conversation": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 1.0,
"hit_at_10_wilson_lower": 0.9750297556319234,
"mrr": 0.9271111111111111,
"n_answerable": 150,
"n_unanswerable": 150,
"ndcg_at_10": 0.9366637917237846,
"recall_at_1": 0.657888888888889,
"recall_at_10": 0.9927777777777778,
"recall_at_5": 0.9814444444444445,
"wrong_holder_confidence_unanswerable": 1.0
},
"image": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.6541666666666667,
"hit_at_10_wilson_lower": 0.5919888820581466,
"mrr": 0.504917328042328,
"n_answerable": 240,
"n_unanswerable": 240,
"ndcg_at_10": 0.5355835725491713,
"recall_at_1": 0.42291666666666666,
"recall_at_10": 0.6458333333333334,
"recall_at_5": 0.6,
"wrong_holder_confidence_unanswerable": 1.0
},
"pdf": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.9164490861618799,
"hit_at_10_wilson_lower": 0.8947140126781276,
"mrr": 0.691579634464752,
"n_answerable": 766,
"n_unanswerable": 218,
"ndcg_at_10": 0.7454292373597072,
"recall_at_1": 0.587467362924282,
"recall_at_10": 0.9164490861618799,
"recall_at_5": 0.825065274151436,
"wrong_holder_confidence_unanswerable": 1.0
},
"ui_screenshot": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.621875,
"hit_at_10_wilson_lower": 0.567595266832662,
"mrr": 0.32982886904761904,
"n_answerable": 320,
"n_unanswerable": 320,
"ndcg_at_10": 0.3983431061821887,
"recall_at_1": 0.21875,
"recall_at_10": 0.621875,
"recall_at_5": 0.465625,
"wrong_holder_confidence_unanswerable": 1.0
},
"video": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.7162790697674418,
"hit_at_10_wilson_lower": 0.6526351336615988,
"mrr": 0.3956736803248431,
"n_answerable": 215,
"n_unanswerable": 215,
"ndcg_at_10": 0.47027044859447076,
"recall_at_1": 0.2837209302325581,
"recall_at_10": 0.7162790697674418,
"recall_at_5": 0.5255813953488372,
"wrong_holder_confidence_unanswerable": 1.0
}
},
"closed_items": 0,
"items": 403859,
"latency_seconds": {
"max": 15.670968882004672,
"mean": 0.09444417075096513,
"p50": 0.08692938649983262,
"p95": 0.10732038790265505,
"queries": 5308
},
"overall": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.8674863387978142,
"hit_at_10_wilson_lower": 0.854722713774439,
"mrr": 0.6285142358400555,
"n_answerable": 2928,
"n_unanswerable": 2380,
"ndcg_at_10": 0.6852635110385413,
"recall_at_1": 0.49903802367941713,
"recall_at_10": 0.8664332877959927,
"recall_at_5": 0.7801286429872496,
"wrong_holder_confidence_unanswerable": 1.0
},
"store": "mem0_flat"
},
{
"by_modality": {
"code": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.952303961196443,
"hit_at_10_wilson_lower": 0.938963104367135,
"mrr": 0.6949699092787209,
"n_answerable": 1237,
"n_unanswerable": 1237,
"ndcg_at_10": 0.7581529489950144,
"recall_at_1": 0.549717057396928,
"recall_at_10": 0.952303961196443,
"recall_at_5": 0.8884397736459175,
"wrong_holder_confidence_unanswerable": 1.0
},
"conversation": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.11333333333333333,
"hit_at_10": 0.9866666666666667,
"hit_at_10_wilson_lower": 0.9526930912996325,
"mrr": 0.9053968253968253,
"n_answerable": 150,
"n_unanswerable": 150,
"ndcg_at_10": 0.9136591750866607,
"recall_at_1": 0.637888888888889,
"recall_at_10": 0.9783333333333334,
"recall_at_5": 0.9458888888888889,
"wrong_holder_confidence_unanswerable": 0.8866666666666667
},
"image": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.5,
"hit_at_10": 0.6541666666666667,
"hit_at_10_wilson_lower": 0.5919888820581466,
"mrr": 0.504917328042328,
"n_answerable": 240,
"n_unanswerable": 240,
"ndcg_at_10": 0.5355835725491713,
"recall_at_1": 0.42291666666666666,
"recall_at_10": 0.6458333333333334,
"recall_at_5": 0.6,
"wrong_holder_confidence_unanswerable": 0.5
},
"pdf": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.9164490861618799,
"hit_at_10_wilson_lower": 0.8947140126781276,
"mrr": 0.691579634464752,
"n_answerable": 766,
"n_unanswerable": 218,
"ndcg_at_10": 0.7454292373597072,
"recall_at_1": 0.587467362924282,
"recall_at_10": 0.9164490861618799,
"recall_at_5": 0.825065274151436,
"wrong_holder_confidence_unanswerable": 1.0
},
"ui_screenshot": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.621875,
"hit_at_10_wilson_lower": 0.567595266832662,
"mrr": 0.32982886904761904,
"n_answerable": 320,
"n_unanswerable": 320,
"ndcg_at_10": 0.3983431061821887,
"recall_at_1": 0.21875,
"recall_at_10": 0.621875,
"recall_at_5": 0.465625,
"wrong_holder_confidence_unanswerable": 1.0
},
"video": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"hit_at_10": 0.7162790697674418,
"hit_at_10_wilson_lower": 0.6526351336615988,
"mrr": 0.3956736803248431,
"n_answerable": 215,
"n_unanswerable": 215,
"ndcg_at_10": 0.47027044859447076,
"recall_at_1": 0.2837209302325581,
"recall_at_10": 0.7162790697674418,
"recall_at_5": 0.5255813953488372,
"wrong_holder_confidence_unanswerable": 1.0
}
},
"closed_items": 27072,
"items": 403859,
"latency_seconds": {
"max": 10.518811015004758,
"mean": 0.09140969146306362,
"p50": 0.0847548475030635,
"p95": 0.10844126994816179,
"queries": 5308
},
"overall": {
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.05756302521008403,
"hit_at_10": 0.8668032786885246,
"hit_at_10_wilson_lower": 0.8540138593279664,
"mrr": 0.6274018236620695,
"n_answerable": 2928,
"n_unanswerable": 2380,
"ndcg_at_10": 0.6840849958419674,
"recall_at_1": 0.4980134335154827,
"recall_at_10": 0.865693306010929,
"recall_at_5": 0.7783071493624771,
"wrong_holder_confidence_unanswerable": 0.942436974789916
},
"store": "zep_temporal"
}
]
}
+2 -1
View File
@@ -292,7 +292,8 @@ def write_markdown(report: dict[str, Any], path: Path) -> None:
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"
f"and its top hit comes from the query's own holder for {overall['top_hit_from_holder_unanswerable']:.0%} of them; "
"the rest are confident answers pulled out of somebody else's memory"
)
if "live_product" in report:
live = report["live_product"]
+28 -8
View File
@@ -31,16 +31,29 @@ from bench.systems import machine_facts
REPORTS = ROOT / "reports"
TOP_K = 10
HOLDER_SCOPE = "holder scope"
WHOLE_CORPUS_SCOPE = "whole corpus"
SCOPES = (HOLDER_SCOPE, WHOLE_CORPUS_SCOPE)
PROGRESS_EVERY = 500
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."""
def score_store(store: Counterpart, queries: list[Query], holder_of: dict[str, str], scope: str) -> dict[str, Any]:
"""Rank every query against one counterpart store and aggregate the same metrics the arms use.
scope "holder" answers only from the asking holder's memory; scope "whole corpus" hides the
holder boundary instead, which is the only difference from what the product's arms are shown.
Without both rows the gain from scoping cannot be told apart from the gain from ranking.
"""
store.timings.clear()
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))
for index, query in enumerate(queries, start=1):
holder = query.holder if scope == HOLDER_SCOPE else None
ranked = store.search(query.text, holder, query.as_of, TOP_K)
scores.append(score_query(query.query_id, ranked, set(query.gold), holder=holder, holder_of=holder_of))
if index % PROGRESS_EVERY == 0:
print(f" {store.name} [{scope}] {index}/{len(queries)}", flush=True)
summary: dict[str, Any] = {
"store": store.name,
"store": f"{store.name} [{scope}]",
"items": len(store.items),
"closed_items": sum(1 for item in store.items if item.valid_to is not None),
"overall": aggregate(scores),
@@ -92,6 +105,13 @@ def write_markdown(report: dict[str, Any], path: Path) -> None:
f"| {latency['p50']:.3f} | {latency['p95']:.3f} |"
)
lines += ["", f"n per modality: {json.dumps(meta['queries_by_modality'], sort_keys=True)}", ""]
lines += [
"",
"`[holder scope]` answers only from the asking holder's memory; `[whole corpus]` searches every",
"holder's memory, which is what the product's arms are shown. The gap between the two rows is the",
"value of the holder boundary, not of a better ranker.",
"",
]
for store in report["stores"]:
lines.append(f"## {store['store']} by modality")
lines.append("")
@@ -99,7 +119,7 @@ def write_markdown(report: dict[str, Any], path: Path) -> None:
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"| {modality} | {row['n_answerable']} | {row['recall_at_10']:.3f} | {row['hit_at_10']:.3f} "
f"| {row['hit_at_10_wilson_lower']:.3f} | {row['mrr']:.3f} |"
)
lines.append("")
@@ -146,7 +166,7 @@ def main() -> int:
"extraction": "lexical (see bench/counterpart.py); LLM extraction deferred with the generator",
"machine": machine_facts(),
},
"stores": [score_store(store, queries, holder_of) for store in stores],
"stores": [score_store(store, queries, holder_of, scope) for store in stores for scope in SCOPES],
}
REPORTS.mkdir(parents=True, exist_ok=True)
args.out.write_text(json.dumps(report, indent=1, sort_keys=True), encoding="utf-8")
+18 -6
View File
@@ -168,13 +168,25 @@ class DenseIndex:
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
_MODELS: dict[str, Any] = {}
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 _sentence_model(model_dir: Path) -> Any:
"""Load a sentence-transformers model once per process.
The query path called this per query, and constructing the model from disk cost
more than the encode itself: dense arm p50 was 1.7s of which most was loading, and
the latency-vs-size curve came out flat because a constant dominated it.
sentence-transformers ships no stubs, so the model is typed Any rather than a fake object.
"""
key = str(model_dir)
if key not in _MODELS:
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")
_MODELS[key] = SentenceTransformer(key, local_files_only=True)
return _MODELS[key]
def top_sources(scores: np.ndarray, atoms: Sequence[Atom], top_k: int) -> list[str]:
+1 -1
View File
@@ -146,7 +146,7 @@ def _overall() -> dict[str, float]:
"hit_at_10_wilson_lower": 0.45,
"abstention_rate_answerable": 0.0,
"abstention_rate_unanswerable": 0.0,
"wrong_holder_confidence_unanswerable": 0.2,
"top_hit_from_holder_unanswerable": 0.2,
}
+25
View File
@@ -162,3 +162,28 @@ def test_percentile_needs_timings_and_a_valid_fraction() -> None:
with pytest.raises(SystemExit, match="percentile must be within"):
runner.percentile(1.5)
assert runner.percentile(0.5) >= 0.0
def test_the_dense_model_is_loaded_once_per_process(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""The query path used to rebuild the model for every query and measured that as retrieval latency."""
import sentence_transformers
from bench import systems
loaded: list[str] = []
class FakeModel:
def __init__(self, path: str, **_kwargs: object) -> None:
loaded.append(path)
monkeypatch.setattr(sentence_transformers, "SentenceTransformer", FakeModel)
systems._MODELS.clear()
model_dir = tmp_path / "dense"
model_dir.mkdir()
try:
assert systems._sentence_model(model_dir) is systems._sentence_model(model_dir)
assert loaded == [str(model_dir)]
with pytest.raises(SystemExit, match="model not downloaded"):
systems._sentence_model(tmp_path / "missing")
finally:
systems._MODELS.clear()