cortexm/bench/ood.py
Rama main · réplica de lectura del repositorio original.
"""Out-of-distribution (OOD) benchmark — breaking the circularity.
The in-distribution (ID) benchmark generates conversations from the SAME
template families the μ=0 extractor's patterns were authored against, so ID
scores are an upper bound for template-shaped text, not a capability claim.
This module measures the honest generalization gap:
1. EXTRACTION RECALL per style — ground-truth facts from persona
registries are re-rendered by an independent LLM in styles the pattern
author never saw (paraphrase / negation / indirect / informal /
non_english / code_switch). The μ=0 extractor runs on the renderings;
recall is matched against the ground-truth match keys.
2. END-TO-END RETRIEVAL — the OOD corpus (renderings + the same
distractor machinery) flows through the full memory fabric and is
probed by the SAME probe builder + judges as the ID benchmark, so
ID-vs-OOD deltas are apples-to-apples.
3. LLM-JUDGE CROSS-CHECK — probe/context pairs are exported in the
canonical BEAM judge format for independent LLM grading.
Renderer omissions (facts the LLM failed to convey) are tracked separately
from extraction failures so neither layer can silently blame the other.
"""
from __future__ import annotations
import random
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from cortexm.bench.abilities import ABILITIES, build_probes, judge
from cortexm.bench.generator import (Corpus, distractor_paragraph,
smalltalk_message)
from cortexm.util import month_name, normalize, token_estimate
T0 = datetime(2026, 3, 1, tzinfo=timezone.utc)
# Match-tier vocabulary: how hard each fact is to extract from re-phrased text.
CATEGORY_TIERS = {
"core": "explicit entity facts (name, employer, city, family, prefs)",
"relational": "multi-hop chains (manager->team->tech)",
"temporal": "dated events and interval changes",
}
def _d(y, m, day=1):
return f"{y:04d}-{m:02d}-{day:02d}"
def _mn(iso_date: str) -> str:
return month_name(int(iso_date[5:7]))
# ---------------------------------------------------------------- manifest
def export_manifest(personas: list, t0: datetime = T0) -> dict:
"""Persona ground truth as LLM-renderable fact manifests.
Mirrors the generator's session structure (parts 0-7) so the SAME probes
remain answerable; facts are stated as semantic text, never as the
generator's template phrasings.
"""
out = []
for p in personas:
e0, e1 = p.employers[0], p.employers[1]
e_last = p.employers[-1]
c0, c1 = p.cities
sister_full = p.family[0][0]
sessions: list[dict] = []
def S(part: int, facts: list[dict]) -> None:
date = t0 + timedelta(days=part * 21)
sessions.append({"session": part, "date": date.date().isoformat(),
"facts": [
{"id": f"s{part}f{i}", "text": t,
"type": ty, "match": m}
for i, (ty, t, m) in enumerate(facts)]})
# ---- part 0: introduction
f0 = [
("name", f"The user's full name is {p.full_name}.",
[p.full_name, p.first]),
("employment", f"The user works at {e0[0]} as a {p.roles[0][0]}.",
[e0[0]]),
("city", f"The user lives in {c0[0]}.", [c0[0]]),
("birthday",
f"The user's birthday is {month_name(p.birthday[0])} "
f"{p.birthday[1]}.",
[f"{month_name(p.birthday[0]).lower()} {p.birthday[1]}"]),
("family",
f"The user has a sister named {sister_full}.",
[sister_full.split()[0]]),
]
if p.nickname:
f0.append(("alias", f"The user goes by the nickname "
f"\"{p.nickname}\".", [p.nickname]))
_instr_keys = (["french"] if "french" in p.instruction[1].lower()
else ["short", "concise", "brief"])
f0.append(("instruction",
f"The user gave the assistant a standing instruction: "
f"\"{p.instruction[0]}\"", _instr_keys))
f0.append(("employment_since",
f"The user has been at {e0[0]} since {_mn(e0[1])} {e0[1][:4]}.",
[e0[0]]))
S(0, f0)
# ---- part 1: preferences + skills
f1 = []
for i in range(0, len(p.prefs), 2):
cat, v_old = p.prefs[i][0], p.prefs[i][1]
v_new = p.prefs[i + 1][1] if i + 1 < len(p.prefs) else v_old
f1.append(("preference",
f"For {cat}, the user used to like {v_old} but now "
f"prefers {v_new}.", [v_old, v_new]))
for s in p.skills:
f1.append(("skill", f"The user knows {s}.", [s]))
f1.append(("hobby", f"In their free time the user enjoys "
f"{p.hobbies[0]}.", [p.hobbies[0]]))
if len(p.hobbies) > 1:
f1.append(("hobby", f"The user also enjoys {p.hobbies[1]}.",
[p.hobbies[1]]))
S(1, f1)
# ---- part 2: job change
f2 = []
m_end = int(e0[2][5:7]) if e0[2] else 6
f2.append(("left_job",
f"The user left {e0[0]} in {_mn(_d(2024, m_end))} "
f"{e0[2][:4] if e0[2] else ''}.".strip(), [e0[0]]))
m_new = int(e1[1][5:7])
f2.append(("joined_job",
f"The user joined {e1[0]} in {_mn(e1[1])} {e1[1][:4]} "
f"as a {p.roles[0][0]}.", [e1[0]]))
if len(p.employers) > 2:
mid = p.employers[1]
m_mid = int(mid[2][5:7]) if mid[2] else 6
f2.append(("left_job",
f"The user later left {mid[0]} in {_mn(_d(2024, m_mid))} "
f"{mid[2][:4] if mid[2] else ''}.".strip(), [mid[0]]))
f2.append(("employment",
f"These days the user works at {e_last[0]}.",
[e_last[0]]))
S(2, f2)
# ---- part 3: relocation + family
m_move = int(c1[1][5:7])
f3 = [
("moved",
f"The user moved to {c1[0]} in {_mn(c1[1])} {c1[1][:4]}, "
f"previously living in {c0[0]}.", [c1[0], c0[0]]),
]
S(3, f3)
# ---- part 4: work structure (multi-hop)
mgr, team = p.manager
tname, tech = p.team_tech
S(4, [
("manager", f"The user's manager is {mgr}.", [mgr]),
("manager_team", f"{mgr} manages the {tname} team.", [tname]),
("team_tech", f"The {tname} team uses {tech}.", [tech]),
("on_team", f"The user is on the {tname} team.", [tname]),
])
# ---- part 5: projects
f5 = []
for name, start, end in p.projects:
if end:
f5.append(("project",
f"The user worked on {name}, finished in "
f"{_mn(end)} {end[:4]}.", [name]))
else:
f5.append(("project",
f"The user is currently working on {name}.", [name]))
S(5, f5)
# ---- parts 6/7: dated events + preference flip
f6 = []
for date, desc in p.events[:2]:
f6.append(("event",
f"On {month_name(int(date[5:7]))} {int(date[8:10])}, "
f"{date[:4]}, the user {desc}.", [desc]))
S(6, f6)
f7 = []
for date, desc in p.events[2:]:
f7.append(("event",
f"On {month_name(int(date[5:7]))} {int(date[8:10])}, "
f"{date[:4]}, the user {desc}.", [desc]))
cat = p.prefs[2][0] if len(p.prefs) > 2 else "coffee"
vals = [v for (c, v, s, e) in p.prefs if c == cat]
if len(vals) >= 2:
f7.append(("preference",
f"The user has since switched to {vals[-1]} for {cat}.",
[vals[-1]]))
S(7, f7)
out.append({"user_id": p.user_id, "full_name": p.full_name,
"sessions": sessions})
return {"personas": out}
# ------------------------------------------------------------ corpus build
def build_ood_corpus(rendered: dict, persona, t0: datetime = T0,
target_tokens: int = 120_000, seed: int = 7) -> Corpus:
"""Assemble an evaluation Corpus from one rendered persona-style row."""
rng = random.Random(seed)
sessions = []
total = 0
for s in sorted(rendered.get("sessions", []), key=lambda x: x.get("session", 0)):
date = t0 + timedelta(days=int(s.get("session", 0)) * 21
+ rng.randrange(0, 5))
msgs = [("user", str(t)) for t in s.get("messages", []) if str(t).strip()]
for _, txt in msgs:
total += token_estimate(txt)
sessions.append((persona.user_id, date, msgs))
# distractor volume — same machinery as the ID generator
guard = 0
while total < target_tokens and guard < 200_000:
guard += 1
uid, date, msgs = sessions[rng.randrange(len(sessions))]
if rng.random() < 0.45:
txt = distractor_paragraph(rng)
else:
txt = smalltalk_message(rng) + " " + smalltalk_message(rng)
k = rng.randrange(0, max(1, len(msgs)))
msgs.insert(k, ("user", txt))
total += token_estimate(txt)
return Corpus(bucket="ood", target_tokens=target_tokens, sessions=sessions,
personas=[persona], total_tokens=total,
generation_seconds=0.0)
# ------------------------------------------------- extraction-layer recall
def extraction_recall(rendered: dict, persona, config) -> dict:
"""Run the μ=0 extractor over one rendered persona-style row.
Returns per-fact match results against the manifest match keys — the
direct, honest measure of pattern generalization.
"""
from cortexm.bridge.extractor import Extractor
from cortexm.bridge.patterns import ExtractionContext
manifest = export_manifest([persona], T0)["personas"][0]
manifest_facts = {f["id"]: f for s in manifest["sessions"]
for f in s["facts"]}
extractor = Extractor(config)
candidates = []
name = None
for s in sorted(rendered.get("sessions", []),
key=lambda x: x.get("session", 0)):
ts = T0 + timedelta(days=int(s.get("session", 0)) * 21)
for text in s.get("messages", []):
ctx = ExtractionContext(user_id=persona.user_id, ts=ts,
speaker="user", subject_name=name,
lexicon=set())
try:
cands = extractor.extract(str(text), ctx)
except Exception:
cands = []
candidates.extend(cands)
for c in cands:
if c.relation == "name" and name is None:
name = c.value
blob = " \n ".join(
f"{normalize(c.subject)} | {normalize(c.relation)} | {normalize(c.value)}"
for c in candidates if c.pattern != "mention_fallback")
matches = {}
for fid, f in manifest_facts.items():
hit = any(normalize(k) and normalize(k) in blob for k in f["match"])
matches[fid] = {"hit": hit, "type": f["type"]}
conveyed = set(rendered.get("conveyed", []))
total = hit = 0
per_type: dict[str, list[int]] = {}
for fid, r in matches.items():
if fid not in conveyed:
continue # renderer omitted it — not an extraction failure
total += 1
hit += int(r["hit"])
per_type.setdefault(r["type"], [0, 0])
per_type[r["type"]][1] += 1
per_type[r["type"]][0] += int(r["hit"])
return {
"user_id": persona.user_id,
"recall": round(hit / total, 4) if total else None,
"n_ground_truth": total,
"n_renderer_omitted": len(manifest_facts) - len(conveyed & set(manifest_facts)),
"n_candidates": len([c for c in candidates
if c.pattern != "mention_fallback"]),
"per_type": {k: round(v[0] / v[1], 4) for k, v in per_type.items() if v[1]},
"missed": [fid for fid, r in matches.items()
if not r["hit"] and fid in conveyed],
"definition": "entity-level recall: share of ground-truth facts "
"whose key appears in any non-fallback candidate; "
"temporal precision is measured end-to-end by the "
"TR/EO probes, not here",
}
# ------------------------------------------------------- end-to-end eval
@dataclass
class OODResult:
style: str
overall: float = 0.0
per_ability: dict = field(default_factory=dict)
n_questions: int = 0
ingest: dict = field(default_factory=dict)
extraction: dict = field(default_factory=dict)
details: list = field(default_factory=list)
def to_dict(self) -> dict:
return {k: v for k, v in self.__dict__.items()}
def run_ood_eval(corpus: Corpus, personas: list, style: str,
db_path: str = ":memory:", max_probes: int | None = None,
judge_fn=None, enrich_fn=None) -> OODResult:
"""Ingest an OOD corpus and evaluate with the standard probe/judge pair.
With ``enrich_fn(memory) -> report`` the probe set is evaluated twice —
before and after async LLM enrichment — quantifying the graceful-
degradation fallback's recovery on OOD text.
"""
import time
from cortexm import metrics
from cortexm.api.memory import Memory
from cortexm.config import Config
t0 = time.time()
res = OODResult(style=style)
cfg = Config(db_path=db_path) if db_path != ":memory:" else Config()
cfg.apply_rules_each_add = False
memory = Memory(cfg)
metrics.reset_counters()
t_ing = time.time()
n_msgs = 0
for user_id, date, msgs in corpus.sessions:
payload = [{"role": role, "content": text, "timestamp": date}
for role, text in msgs]
n_msgs += len(msgs)
memory.add(payload, user_id=user_id, timestamp=date)
memory.apply_rules()
ingest_s = time.time() - t_ing
stats = memory.stats()
res.ingest = {
"wall_seconds": round(ingest_s, 2),
"messages": n_msgs,
"tokens": corpus.total_tokens,
"tokens_per_second": int(corpus.total_tokens / max(ingest_s, 1e-9)),
"llm_calls": metrics.counters()["llm_calls"],
"u0_protocol": stats["u0_protocol"],
"facts": stats["facts"],
}
def _eval() -> tuple[float, dict, list]:
rng = random.Random(hash((style, "probes")) & 0xFFFFFFFF)
probes = build_probes(personas, rng)
by_ability: dict[str, list] = {a: [] for a in ABILITIES}
for p in probes:
by_ability[p.ability].append(p)
if max_probes:
for a in ABILITIES:
by_ability[a] = by_ability[a][:max_probes]
n_q = sum(len(v) for v in by_ability.values())
score_fn = judge_fn or judge
per_ability = {a: 0.0 for a in ABILITIES}
counts = {a: len(by_ability[a]) for a in ABILITIES}
details = []
for ability in ABILITIES:
for probe in by_ability[ability]:
out = memory.search(probe.question, user_id=probe.user_id, k=12)
score, detail = score_fn(probe, out["context_block"])
per_ability[ability] += score
details.append({"ability": ability, "question": probe.question,
"score": score, "detail": detail,
"context": out["context_block"][:600]})
overall = round(sum(per_ability.values()) / max(n_q, 1), 4)
per_ability = {a: round(per_ability[a] / max(counts[a], 1), 4)
for a in ABILITIES if counts[a]}
return overall, per_ability, details
res.overall, res.per_ability, res.details = _eval()
res.n_questions = len(res.details)
res.extraction = {"facts": stats["facts"]}
if enrich_fn is not None:
rep = enrich_fn(memory)
res.ingest["enrichment"] = rep
post_overall, post_per_ability, _post = _eval()
res.enriched = {"overall": post_overall,
"per_ability": post_per_ability,
"report": rep} # type: ignore[attr-defined]
memory.close()
res.ingest["wall_seconds_total"] = round(time.time() - t0, 2)
return res
# ---------------------------------------------------- LLM-judge item export
def describe_expected(probe) -> list[str]:
"""Human-readable nugget description for the canonical LLM judge."""
exp = probe.expected
if probe.ability == "AB":
return [f"NOTHING — '{probe.question}' concerns an attribute never "
f"mentioned; correct behaviour is abstention."]
if "events" in exp:
return [f"Event '{d}' occurred on {iso}." for d, iso in exp["events"]]
if "set" in exp:
return ["The complete set of projects: " + ", ".join(exp["set"]) + "."]
if "current" in exp and "old" in exp:
return [f"Current employer: {exp['current']}.",
f"Previous employer (superseded): {exp['old']}."]
out = []
for k in exp.get("contains_any", []):
out.append(f"The answer must mention: {k}.")
for k in exp.get("also_any", []):
out.append(f"Supporting evidence should also mention: {k}.")
return out or ["(no nugget description)"]
def export_judge_items(result: OODResult, personas: list,
per_ability_limit: int = 4) -> list[dict]:
"""Probe/context pairs in the canonical judge JSONL format.
The deterministic judge's score is attached (det_score) so agreement
between the two graders can be computed after the LLM pass.
"""
items = []
per_ability: dict[str, int] = {}
probe_by_q = {pr.question: pr for pr in build_probes(personas, random.Random(1234))}
for i, d in enumerate(result.details):
if per_ability.get(d["ability"], 0) >= per_ability_limit:
continue
per_ability[d["ability"]] = per_ability.get(d["ability"], 0) + 1
probe = probe_by_q.get(d["question"])
items.append({
"id": f"{result.style}-{i}",
"ability": d["ability"],
"question": d["question"],
"expected": describe_expected(probe) if probe
else ["(probe metadata unavailable)"],
"context": d["context"],
"det_score": d["score"],
})
return items