File size: 11,769 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | #!/usr/bin/env python3
"""Command-line entry point.
Method parameters come from configs/methods.yaml (configs/methods.memos.yaml for memos), models from
configs/backbone.yaml and patients from configs/patients.yaml.
Each (method, patient) run is independent and can run as a separate process.
"""
from __future__ import annotations
import argparse
import logging
from pathlib import Path
import yaml
from .dataset import load_patients
from .llm_client import set_endpoints, set_reasoning_models
from .memory_runner import rejudge_patient, run_patient
from .methods import (SUPPORTED_METHODS, DEFAULT_METHODS, INTERNAL_LLM_METHODS, EMBEDDING_DIM_METHODS,
EMBEDDING_MODEL_METHODS, EMBEDDING_BATCH_METHODS, EMBEDDING_TOKENIZER_METHODS, TOP_K_PARAMS,
embedding_params)
from .methods.base import EMBED_ENC, ENC
PKG = Path(__file__).resolve().parent
EVAL_ROOT = PKG.parent
def _setup_logging() -> None:
"""Show the method agents' log records; third-party loggers stay at WARNING."""
logging.basicConfig(level=logging.WARNING,
format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logging.getLogger("memebench").setLevel(logging.INFO)
for noisy in ("httpx", "httpcore", "openai", "urllib3", "filelock"):
logging.getLogger(noisy).setLevel(logging.WARNING)
def load_backbone(path: str) -> dict:
"""Read the backbone config. A <PLACEHOLDER> value is unset (None); the values the selected methods need are
checked in main()."""
bb = yaml.safe_load(open(path, encoding="utf-8")) or {}
for role, sec in bb.items():
for k, v in (sec or {}).items():
if isinstance(v, str) and v.startswith("<") and v.endswith(">"):
sec[k] = None
return bb
def main():
_setup_logging()
ap = argparse.ArgumentParser(description="MEME-BENCH memory-method evaluation")
ap.add_argument("--method", help=f"one of: {', '.join(SUPPORTED_METHODS)}, a comma-list, "
f"or 'all' (= {', '.join(DEFAULT_METHODS)}; memory-system "
f"methods must be named explicitly)")
ap.add_argument("--patient", help="patient id/name from configs/patients.yaml, or 'all'")
ap.add_argument("--list-patients", action="store_true",
help="print the registered patients (id, #sessions, #questions) and exit")
ap.add_argument("--methods-config", default=str(PKG / "configs" / "methods.yaml"))
ap.add_argument("--backbone-config", default=str(PKG / "configs" / "backbone.yaml"))
ap.add_argument("--patients-config", default=str(PKG / "configs" / "patients.yaml"))
ap.add_argument("--output-dir", default=str(EVAL_ROOT / "outputs"))
ap.add_argument("--resume", action="store_true",
help="continue the newest run, skipping qids already graded in answers.jsonl; "
"with --qids, only the selected qids are answered again")
ap.add_argument("--qids", default="", help="with --resume: comma list / a-b ranges of questions to answer again")
ap.add_argument("--judge-workers", type=int, default=8)
ap.add_argument("--top-k", type=int, default=None,
help="number of items every method retrieves for an answer; overrides top_k of the "
"methods config")
ap.add_argument("--rejudge", action="store_true",
help="re-grade every stored answer of the newest completed run with the current judge "
"(no re-answering)")
args = ap.parse_args()
cfg = yaml.safe_load(open(args.methods_config, encoding="utf-8"))
if args.top_k is not None:
if args.top_k < 1:
ap.error("--top-k must be positive")
for m, key in TOP_K_PARAMS.items():
if key in (cfg["methods"].get(m) or {}):
cfg["methods"][m][key] = args.top_k
skipped: list[tuple[str, str]] = []
patients = load_patients(args.patients_config, skipped)
if args.list_patients:
print(f"{'patient':22} {'sessions':>9} {'questions':>10} dir/dialogues")
for name, p in patients.items():
try:
s, q = p.load()
print(f"{name:22} {len(s):>9} {len(q):>10} {p.dir or p.dialogues}")
except SystemExit as e:
print(f"{name:22} {'-':>9} {'-':>10} unusable: {str(e)[:60]}")
for name, reason in skipped:
print(f"SKIPPED {name}: {reason}")
return
if not args.method or not args.patient:
raise SystemExit("--method and --patient are required (or use --list-patients)")
methods = (list(DEFAULT_METHODS) if args.method == "all"
else [m.strip() for m in args.method.split(",") if m.strip()])
for m in methods:
if m not in cfg["methods"]:
hint = (f"; use --methods-config memebench/configs/methods.{m}.yaml"
if (PKG / "configs" / f"methods.{m}.yaml").exists() else "")
raise SystemExit(f"method '{m}' not in {args.methods_config}{hint}")
left = [k for k, v in (cfg["methods"][m] or {}).items()
if isinstance(v, str) and v.startswith("<") and v.endswith(">")]
if left:
raise SystemExit(f"{args.methods_config}: set {', '.join(left)} for method '{m}' before running")
pnames = (list(patients) if args.patient == "all"
else [p.strip() for p in args.patient.split(",") if p.strip()])
for p in pnames:
if p not in patients:
raise SystemExit(f"patient '{p}' not registered (use --list-patients to see the "
f"{len(patients)} entries from {args.patients_config})")
bb = load_backbone(args.backbone_config)
need = [("answer", "model"), ("judge", "model")]
if not args.rejudge and any(m in EMBEDDING_MODEL_METHODS for m in methods):
need.append(("embedding", "model"))
if not args.rejudge and any(m in EMBEDDING_DIM_METHODS for m in methods):
need.append(("embedding", "dim"))
if not args.rejudge and any(m in EMBEDDING_BATCH_METHODS for m in methods):
need.append(("embedding", "batch_size"))
need += [("judge", "reasoning")] + ([] if args.rejudge else [("answer", "tokenizer")])
if not args.rejudge and any(m in INTERNAL_LLM_METHODS for m in methods):
need.append(("answer", "reasoning"))
if not args.rejudge and any(m in EMBEDDING_TOKENIZER_METHODS for m in methods):
need.append(("embedding", "tokenizer"))
missing = [f"{r}.{k}" for r, k in need if (bb.get(r) or {}).get(k) is None]
if missing:
raise SystemExit(f"set these values in {args.backbone_config}: {', '.join(missing)}")
answer_model = bb["answer"]["model"]
judge_model = bb["judge"]["model"]
answer_base_url = bb["answer"].get("base_url")
set_endpoints({role: bb.get(role) or {} for role in ("answer", "judge", "embedding", "memory")})
if answer_base_url or (bb.get("judge") or {}).get("base_url"):
print(f"answering backbone: {answer_model} @ {answer_base_url or 'the default endpoint'}; judge: {judge_model} "
f"@ {(bb.get('judge') or {}).get('base_url') or 'the default endpoint'}")
memory = bb.get("memory") or {}
internal_model = memory.get("model")
internal_base_url = memory.get("base_url")
if internal_model and any(m in INTERNAL_LLM_METHODS for m in methods):
print(f"internal LLM: {internal_model} @ {internal_base_url or 'the default endpoint'} "
f"(methods: {', '.join(m for m in methods if m in INTERNAL_LLM_METHODS)})")
own = [m for m in methods if m in ("hipporag", "remem", "letta", "mirix")]
if internal_model != answer_model and own:
print(f"note: {', '.join(own)} answer inside the method, with the memory model {internal_model}")
else:
internal_model = internal_base_url = None
reasoning = [m for role, m in (("answer", answer_model), ("judge", judge_model))
if (bb.get(role) or {}).get("reasoning") is True]
if internal_model:
mem_reasoning = memory.get("reasoning")
if mem_reasoning is None and internal_model != answer_model:
raise SystemExit(f"set memory.reasoning in {args.backbone_config} for the memory model {internal_model}")
if mem_reasoning is True or (mem_reasoning is None and answer_model in reasoning):
reasoning.append(internal_model)
set_reasoning_models(reasoning)
if not args.rejudge:
ENC.set((bb.get("answer") or {}).get("tokenizer"))
if any(m in EMBEDDING_TOKENIZER_METHODS for m in methods):
EMBED_ENC.set((bb.get("embedding") or {}).get("tokenizer"))
import traceback
failures = []
for pname in pnames:
for m in methods:
print(f"\n===== {m} x {pname} =====", flush=True)
try:
if args.rejudge:
rep = rejudge_patient(
method_name=m, patient=patients[pname],
out_root=Path(args.output_dir),
answer_model=answer_model,
judge_model=judge_model,
judge_workers=args.judge_workers,
)
else:
rep = run_patient(
method_name=m,
patient=patients[pname],
method_params={**embedding_params(m, bb["embedding"]),
**(cfg["methods"][m] or {})},
out_root=Path(args.output_dir),
answer_model=answer_model,
judge_model=judge_model,
answer_max_tokens=int(cfg["answer_max_tokens"]),
answer_temperature=cfg.get("answer_temperature"),
resume=args.resume,
qids=args.qids,
judge_workers=args.judge_workers,
internal_model=internal_model,
internal_base_url=internal_base_url,
answer_base_url=answer_base_url,
)
if rep.get("errors"):
failures.append(f"{m} x {pname}: {rep['errors']} error row(s) "
f"in {rep['total']} (run completed but is not clean)")
if rep.get("judge_failures"):
failures.append(f"{m} x {pname}: judge could not grade {rep['judge_failures']} "
f"answers; re-grade with --rejudge")
if rep.get("memorize_failures") or rep.get("consolidate_failures"):
failures.append(f"{m} x {pname}: memory construction failed on "
f"{rep.get('memorize_failures', 0)} visit(s) and "
f"{rep.get('consolidate_failures', 0)} checkpoint consolidation(s) "
f"(see memory_build.jsonl; the run completed but is not clean)")
except KeyboardInterrupt:
raise
except (SystemExit, Exception) as e:
failures.append(f"{m} x {pname}: {type(e).__name__}: {str(e)[:150]}")
print(f"\nfailed {m} x {pname}: {type(e).__name__}: {e}", flush=True)
traceback.print_exc()
print("continuing with the next pair (re-run with --resume to retry)\n", flush=True)
if failures:
print(f"\n===== {len(failures)} pair(s) failed or unclean =====")
for f in failures:
print(" " + f)
raise SystemExit(1)
if __name__ == "__main__":
main()
|