scripts/recompute_metrics.py

#!/usr/bin/env python3
"""Recompute EM/GM-proxy/MB-proxy on stored eval_results.jsonl (e.g. after a
metrics.py fix) without re-calling any models."""
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from metrics import exact_match, graph_match_proxy, modify_better_proxy  # noqa: E402


def main():
    in_path = Path(sys.argv[1] if len(sys.argv) > 1 else "outputs/eval_results.jsonl")
    tasks = {t["task_id"]: t for t in json.load(open("data/squirrel_mini_tasks.json"))}
    recs = [json.loads(l) for l in open(in_path)]
    for r in recs:
        if r.get("error"):
            continue
        task = tasks[r["task_id"]]
        r["em"] = exact_match(r["pred_sql"], task["reference_sql"])
        r["gm_proxy"] = graph_match_proxy(r["pred_sql"], task["reference_sql"])
        r["mb_proxy"] = modify_better_proxy(r["pred_sql"], task["reference_sql"], task["issue_sql"])
    with open(in_path, "w") as f:
        for r in recs:
            f.write(json.dumps(r) + "\n")
    print(f"Recomputed metrics for {len(recs)} records in {in_path}")


if __name__ == "__main__":
    main()