Biomanus/experiments/lab_bench/scripts/summarize_labbench_repeats.py

#!/usr/bin/env python3
"""Summarize repeated LAB-Bench DbQA/SeqQA runs."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from statistics import mean, stdev
from typing import Any


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def numeric_metrics(payloads: list[dict[str, Any]]) -> tuple[dict[str, float], dict[str, float]]:
    keys = sorted({key for payload in payloads for key in payload.get("metrics_all", {})})
    means: dict[str, float] = {}
    stds: dict[str, float] = {}
    for key in keys:
        vals = [payload["metrics_all"][key] for payload in payloads if isinstance(payload["metrics_all"].get(key), (int, float))]
        if not vals:
            continue
        means[key] = mean(vals)
        stds[key] = stdev(vals) if len(vals) > 1 else 0.0
    return means, stds


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-root", type=Path, required=True)
    parser.add_argument("--out", type=Path, default=None)
    args = parser.parse_args()

    groups: dict[str, list[Path]] = {}
    for path in sorted(args.results_root.glob("*_run*.json")):
        if path.name.startswith("combined_"):
            continue
        group = path.stem.rsplit("_run", 1)[0]
        groups.setdefault(group, []).append(path)

    payload = {"results_root": str(args.results_root), "groups": {}}
    for group, paths in sorted(groups.items()):
        loaded = [load_json(path) for path in paths]
        means, stds = numeric_metrics(loaded)
        payload["groups"][group] = {
            "run_count": len(paths),
            "files": [str(path) for path in paths],
            "metrics_all_mean": means,
            "metrics_all_std": stds,
        }

    out = args.out or args.results_root / "summary.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())