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:
1 parent
c4e2982506
commit
dbe1d2cc82
8 files changed
+381
-34
No files matched your search
+50
-17
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user