verify.py

#!/usr/bin/env python3
"""Reproduce MemMesh's LOCOMO scores against Mem0 and Zep — under each system's
OWN judge. This re-scores our published per-question predictions, so you can
confirm we didn't grade ourselves with a friendly judge.

    pip install openai
    export OPENAI_API_KEY=sk-...
    python verify.py            # all three judges
    python verify.py --judge mem0

The three judges are the EXACT rules each system grades LOCOMO with:
  strict  — our internal bar: exact/semantic match (harder than either competitor)
  zep     — Zep's rule (graphiti_core/prompts/eval.py): "same topic = correct"
  mem0    — Mem0's rule (memory-benchmarks): partial-credit, >=1 gold list item = correct
"""
import argparse, json, os, re
from concurrent.futures import ThreadPoolExecutor
from openai import OpenAI

JUDGES = {
    "strict": (
        "You grade whether a predicted answer matches the gold answer. CORRECT only "
        "if it conveys the SAME specific factual answer as the gold (exact/semantic "
        "match). Resolve numbers and dates by value (3.10%==3.1%; a relative date "
        "equals the absolute date it denotes). A vague, partial, or merely on-topic "
        "answer is INCORRECT. Reply JSON {\"correct\": true|false}."
    ),
    "zep": (  # Zep, verbatim: graphiti_core/prompts/eval.py eval_prompt
        "You grade a predicted answer against a gold answer. Although the prediction "
        "may be more verbose, mark it CORRECT as long as it references the SAME TOPIC "
        "as the gold answer. Reply JSON {\"correct\": true|false}."
    ),
    "mem0": (  # Mem0, partial-credit: memory-benchmarks _JUDGE_TEMPLATE
        "You grade a predicted answer against a gold answer with PARTIAL CREDIT. Mark "
        "CORRECT if the prediction is about the same referent and conveys the gold, OR "
        "— for list/'how many' golds — contains AT LEAST ONE correct item from the "
        "gold list. Semantic overlap counts; extra detail is fine; resolve dates and "
        "numbers by value. INCORRECT only if wrong referent, contradictory, or a "
        "refusal. Reply JSON {\"correct\": true|false}."
    ),
}
# Competitors' published LOCOMO figures (answerable categories, LLM-judge).
PUBLISHED = {
    "mem0": {"temporal": 55, "single_hop": 67, "multi_hop": 51, "open_domain": 73, "overall": 66.9},
    "zep":  {"temporal": 49, "single_hop": 62, "multi_hop": 41, "open_domain": 77, "overall": 66.0},
}


def parse(t):
    try:
        s, e = t.find("{"), t.rfind("}")
        return bool(json.loads(t[s:e + 1]).get("correct", False))
    except Exception:
        return t.strip().lower().startswith(("yes", "true", "correct"))


def judge_one(client, model, sysmsg, r):
    msg = (f"Question: {r['question']}\nGold answer: {r['gold']}\n"
           f"Predicted answer: {r['prediction']}\n\nIs the prediction correct? Reply with the JSON.")
    out = client.chat.completions.create(
        model=model, temperature=0, max_tokens=20,
        messages=[{"role": "system", "content": sysmsg}, {"role": "user", "content": msg}],
    )
    return r["category"], parse(out.choices[0].message.content or "")


def score(rows, sysmsg, model, workers):
    client = OpenAI()
    by = {}
    with ThreadPoolExecutor(max_workers=workers) as ex:
        for cat, ok in ex.map(lambda r: judge_one(client, model, sysmsg, r), rows):
            b = by.setdefault(cat, [0, 0]); b[1] += 1
            if ok:
                b[0] += 1
    ans = [sum(by.get(c, [0, 0])[0] for c in by if c != "adversarial"),
           sum(by.get(c, [0, 0])[1] for c in by if c != "adversarial")]
    return by, ans


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--predictions", default="predictions.json")
    ap.add_argument("--judge", choices=list(JUDGES) + ["all"], default="all")
    ap.add_argument("--model", default="gpt-4o")
    ap.add_argument("--workers", type=int, default=10)
    args = ap.parse_args()
    rows = json.load(open(args.predictions))
    names = list(JUDGES) if args.judge == "all" else [args.judge]

    hdr = f"{'judge':<10}{'temporal':>10}{'single':>9}{'multi':>8}{'open':>7}{'overall*':>11}"
    print(f"MemMesh — {len(rows)} predictions ({args.model} judge)\n" + hdr + "\n" + "-" * len(hdr))
    for name in names:
        by, ans = score(rows, JUDGES[name], args.model, args.workers)
        def p(c):
            a, b = by.get(c, [0, 0]); return f"{100*a/b:.0f}" if b else "-"
        print(f"{name:<10}{p('temporal'):>10}{p('single_hop'):>9}{p('multi_hop'):>8}"
              f"{p('open_domain'):>7}{100*ans[0]/max(ans[1],1):>10.1f}%")
    print("\n* overall = answerable categories only (excl. adversarial, which every system drops)")
    for c in ("mem0", "zep"):
        if c in names or args.judge == "all":
            pb = PUBLISHED[c]
            print(f"  {c} published: temporal {pb['temporal']} single {pb['single_hop']} "
                  f"multi {pb['multi_hop']} open {pb['open_domain']} overall {pb['overall']}")
    print("\nRun MemMesh under Mem0's judge and compare to Mem0's published; likewise Zep.")


if __name__ == "__main__":
    main()