scripts/analyze_computation_comparison.py

#!/usr/bin/env python3
"""Report canonical HERB run performance by computation/comparison bracket."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parent.parent


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


def bracket_summary(items: list[dict[str, Any]]) -> dict[str, float | int]:
    scope = len(items)
    answered = sum(bool(item.get("answered")) for item in items)
    correct = sum(bool(item.get("correct")) for item in items)
    score_total = sum(
        float(item["score"])
        for item in items
        if isinstance(item.get("score"), (int, float))
    )
    return {
        "scope": scope,
        "answered": answered,
        "correct": correct,
        "score_pct": round(score_total * 100 / scope, 2),
        "perfect_pct": round(correct * 100 / scope, 2),
        "coverage_pct": round(answered * 100 / scope, 2),
    }


def main() -> None:
    eval_rows = load_json(ROOT / "eval.json")
    answerable = [row for row in eval_rows if row.get("kind") == "answerable"]
    labels = {
        row["gid"].replace("#", "_", 1): bool(row["computation_comparison"])
        for row in answerable
    }
    if len(labels) != 815:
        raise ValueError(f"expected 815 answerable labels, found {len(labels)}")

    bracket_counts = {
        "computation_comparison": sum(labels.values()),
        "other": sum(not value for value in labels.values()),
    }
    report: dict[str, Any] = {
        "scope": len(labels),
        "brackets": {
            name: {
                "scope": count,
                "percent": round(count * 100 / len(labels), 2),
            }
            for name, count in bracket_counts.items()
        },
        "runs": [],
    }

    manifest = load_json(ROOT / "runs" / "manifest.json")
    for run in manifest["runs"]:
        index = load_json(ROOT / "runs" / run["slot"] / "index.json")
        by_bracket = {
            "computation_comparison": [
                item for item in index["items"] if labels[item["qid"]]
            ],
            "other": [item for item in index["items"] if not labels[item["qid"]]],
        }
        report["runs"].append(
            {
                "slot": run["slot"],
                "label": run["label"],
                "score_mode": "mean_judge_score",
                "brackets": {
                    name: bracket_summary(items)
                    for name, items in by_bracket.items()
                },
            }
        )

    out_path = ROOT / "computation_comparison_report.json"
    out_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(f"wrote {out_path}")
    print(
        "brackets: "
        f"{bracket_counts['computation_comparison']}/815 computation/comparison, "
        f"{bracket_counts['other']}/815 other"
    )
    for run in report["runs"]:
        comp = run["brackets"]["computation_comparison"]
        other = run["brackets"]["other"]
        print(
            f"{run['label']}: {comp['score_pct']:.2f}% computation/comparison, "
            f"{other['score_pct']:.2f}% other"
        )


if __name__ == "__main__":
    main()