File size: 6,196 Bytes
9589849 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python3
"""Compare two paired v1 eval runs, conservatively counting scheduled episodes."""
from __future__ import annotations
import argparse
import json
import math
import tomllib
from collections import Counter
from pathlib import Path
def wilson(successes: int, total: int, z: float = 1.959963984540054) -> list[float] | None:
if total == 0:
return None
p = successes / total
d = 1 + z * z / total
c = (p + z * z / (2 * total)) / d
r = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / d
return [max(0.0, c - r), min(1.0, c + r)]
def read_traces(path: Path) -> list[dict]:
path = path / "traces.jsonl" if path.is_dir() else path
by_id: dict[str, dict] = {}
for line in path.read_text().splitlines():
record = json.loads(line)
for trace in record.get("traces", []):
by_id[trace["id"]] = trace
return list(by_id.values())
def task_key(trace: dict) -> str:
data = (trace.get("task") or {}).get("data") or {}
return str(data.get("name") or data.get("instance_id") or data.get("idx") or trace["id"])
def clean(trace: dict) -> bool:
return bool(trace.get("ok") and trace.get("is_completed") and not trace.get("errors"))
def solved(trace: dict) -> bool:
return float(((trace.get("rewards") or {}).get("solved") or {}).get("score", 0)) > 0
def has_reward(trace: dict) -> bool:
return isinstance((trace.get("rewards") or {}).get("solved"), dict)
def exact_two_sided_sign_p(a_only: int, b_only: int) -> float | None:
"""Exact paired sign-test p-value over discordant outcomes.
This is deliberately dependency-free and conditions only on tasks where the
two harnesses/checkpoints disagree. It is diagnostic rather than a reason
to hide either run's conservative Wilson interval.
"""
discordant = a_only + b_only
if discordant == 0:
return None
tail = sum(math.comb(discordant, k) for k in range(min(a_only, b_only) + 1))
return min(1.0, 2.0 * tail / (2**discordant))
def expected_episodes(path: Path, observed: int) -> int:
run_dir = path if path.is_dir() else path.parent
config_path = run_dir / "config.toml"
if not config_path.is_file():
return observed
config = tomllib.loads(config_path.read_text())
num_tasks = config.get("num_tasks")
num_rollouts = config.get("num_rollouts", 1)
if not isinstance(num_tasks, int) or not isinstance(num_rollouts, int):
return observed
return max(observed, num_tasks * num_rollouts)
def summary(traces: list[dict], path: Path) -> dict:
clean_traces = [trace for trace in traces if clean(trace)]
scored = [trace for trace in traces if has_reward(trace)]
successes = sum(solved(trace) for trace in traces)
calls = [call for trace in traces for call in trace.get("calls", [])]
prompt_tokens = sum(int((call.get("usage") or {}).get("prompt_tokens", 0)) for call in calls)
completion_tokens = sum(
int((call.get("usage") or {}).get("completion_tokens", 0)) for call in calls
)
categories = Counter(
str(((trace.get("task") or {}).get("data") or {}).get("category") or "unknown")
for trace in traces
if solved(trace)
)
scheduled = expected_episodes(path, len(traces))
return {
"scheduled_episodes": scheduled,
"observed_trace_episodes": len(traces),
"missing_trace_episodes": scheduled - len(traces),
"clean_completed": len(clean_traces),
"reward_bearing": len(scored),
"errored_or_unscored": scheduled - len(scored),
"successes": successes,
"score_scheduled": successes / scheduled if scheduled else None,
"ci95_wilson_scheduled": wilson(successes, scheduled),
"diagnostic_score_reward_bearing": successes / len(scored) if scored else None,
"diagnostic_ci95_reward_bearing": wilson(successes, len(scored)),
"stop_conditions": dict(Counter(trace.get("stop_condition", "unknown") for trace in traces)),
"model_calls": len(calls),
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"success_categories": dict(sorted(categories.items())),
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("a", type=Path)
parser.add_argument("b", type=Path)
parser.add_argument("--label-a", default="a")
parser.add_argument("--label-b", default="b")
args = parser.parse_args()
a_traces = read_traces(args.a)
b_traces = read_traces(args.b)
a = {task_key(trace): trace for trace in a_traces}
b = {task_key(trace): trace for trace in b_traces}
common = sorted(a.keys() & b.keys())
paired = Counter()
for key in common:
sa, sb = solved(a[key]), solved(b[key])
paired[
"both" if sa and sb else "a_only" if sa else "b_only" if sb else "neither"
] += 1
paired_delta = (
(paired["b_only"] - paired["a_only"]) / len(common) if common else None
)
print(
json.dumps(
{
args.label_a: summary(a_traces, args.a),
args.label_b: summary(b_traces, args.b),
"paired": {
"common_scheduled_tasks": len(common),
"both_solve": paired["both"],
f"{args.label_a}_only": paired["a_only"],
f"{args.label_b}_only": paired["b_only"],
"neither": paired["neither"],
f"score_delta_{args.label_b}_minus_{args.label_a}": paired_delta,
"discordant_tasks": paired["a_only"] + paired["b_only"],
"exact_two_sided_sign_test_p": exact_two_sided_sign_p(
paired["a_only"], paired["b_only"]
),
f"missing_{args.label_a}": len(b.keys() - a.keys()),
f"missing_{args.label_b}": len(a.keys() - b.keys()),
"both_clean_completed": sum(clean(a[key]) and clean(b[key]) for key in common),
},
},
indent=2,
)
)
if __name__ == "__main__":
main()
|