#!/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 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()