analyze.py

#!/usr/bin/env python3
"""Analyze ComposeBench JSONL results.

Reports success rate and mean turns by chain depth or fan-in width, with Wilson
95% intervals. Accepts one or more result files.

Usage:
  python3 analyze.py results.jsonl
  python3 analyze.py results/chain_clean.jsonl results/fanin_clean.jsonl
  python3 analyze.py --self-test
"""
import argparse
import json
import math
from collections import defaultdict

SUBSTRATES = ["repl", "batchfc", "fc", "cli"]


def wilson(k, n, z=1.96):
    if n == 0:
        return 0.0, 0.0
    p = k / n
    d = 1 + z * z / n
    c = p + z * z / (2 * n)
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
    return (c - h) / d, (c + h) / d


def load_rows(paths):
    rows = []
    for path in paths:
        with open(path) as f:
            for line in f:
                if line.strip():
                    rows.append(json.loads(line))
    return rows


def summarize(rows):
    present = [s for s in SUBSTRATES if any(r.get("substrate") == s for r in rows)]
    if not present:
        present = sorted({r.get("substrate") for r in rows})

    print(f"ComposeBench -- {len(rows)} rows")
    print()
    for kind, axis in [("chain", "depth"), ("fanin", "width")]:
        subrows = [r for r in rows if r.get("kind") == kind]
        if not subrows:
            continue
        values = sorted({r.get(axis) for r in subrows if r.get(axis) is not None})
        print(f"== {kind} by {axis} ==")
        print(f"{axis:>6}  " + "  ".join(f"{s:>22}" for s in present))
        for value in values:
            line = f"{value:>6}  "
            cells = []
            for sub in present:
                rs = [r for r in subrows if r.get(axis) == value and r.get("substrate") == sub]
                if not rs:
                    cells.append(f"{'-':>22}")
                    continue
                ok = sum(1 for r in rs if r.get("is_correct"))
                turns = sum(float(r.get("n_turns", 0)) for r in rs) / len(rs)
                lo, hi = wilson(ok, len(rs))
                cells.append(f"{ok/len(rs)*100:5.1f}%/{turns:4.1f}t [{lo*100:2.0f},{hi*100:3.0f}]")
            print(line + "  ".join(cells))
        print()

    print("== overall ==")
    print(f"{'substrate':>10} {'n':>5} {'success':>10} {'turns':>8} {'in_tok':>10} {'out_tok':>10}")
    for sub in present:
        rs = [r for r in rows if r.get("substrate") == sub]
        if not rs:
            continue
        ok = sum(1 for r in rs if r.get("is_correct"))
        turns = sum(float(r.get("n_turns", 0)) for r in rs) / len(rs)
        in_tok = sum(int(r.get("input_tokens", 0)) for r in rs) / len(rs)
        out_tok = sum(int(r.get("output_tokens", 0)) for r in rs) / len(rs)
        print(f"{sub:>10} {len(rs):5d} {ok/len(rs)*100:9.1f}% {turns:8.2f} {in_tok:10.0f} {out_tok:10.0f}")


def self_test_rows():
    rows = []
    for substrate in SUBSTRATES:
        for depth in (1, 2):
            rows.append({
                "task_id": f"chain_d{depth}_00",
                "kind": "chain",
                "depth": depth,
                "width": None,
                "substrate": substrate,
                "rep": 0,
                "is_correct": True,
                "n_turns": 2 if substrate in ("repl", "batchfc") else depth + 1,
                "input_tokens": 100,
                "output_tokens": 20,
            })
        for width in (3, 5):
            rows.append({
                "task_id": f"fanin_w{width}_00",
                "kind": "fanin",
                "depth": None,
                "width": width,
                "substrate": substrate,
                "rep": 0,
                "is_correct": substrate != "cli" or width == 3,
                "n_turns": 2 if substrate in ("repl", "batchfc") else width + 2,
                "input_tokens": 100,
                "output_tokens": 20,
            })
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("paths", nargs="*")
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()
    if args.self_test:
        summarize(self_test_rows())
        return
    if not args.paths:
        raise SystemExit("pass one or more JSONL files, or use --self-test")
    summarize(load_rows(args.paths))


if __name__ == "__main__":
    main()