| |
| """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() |
|
|
