| """Summaries of graded generations, and paired comparisons between two runs.""" |
| from __future__ import annotations |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
|
|
| def read_jsonl(path): |
| with open(path) as f: |
| return [json.loads(line) for line in f if line.strip()] |
|
|
|
|
| def _mean(xs): |
| xs = list(xs) |
| return sum(xs) / len(xs) if xs else None |
|
|
|
|
| def _se(xs): |
| xs = list(xs) |
| if len(xs) < 2: |
| return None |
| m = sum(xs) / len(xs) |
| return math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1) / len(xs)) |
|
|
|
|
| def summarize(recs: list[dict]) -> dict: |
| S = [s for r in recs for s in r["samples"]] |
| fin = [s for s in S if not s["cut"]] |
| cut = [s for s in S if s["cut"]] |
| per_problem = [_mean(s["correct"] for s in r["samples"]) for r in recs] |
| out = { |
| "problems": len(recs), |
| "samples": len(S), |
| "acc": _mean(s["correct"] for s in S), |
| "acc_se": _se(per_problem), |
| "cut_share": _mean(s["cut"] for s in S), |
| "mean_tokens": _mean(s["tokens"] for s in S), |
| "closed_share": _mean(s["closed"] for s in S), |
| "p_finish": len(fin) / len(S) if S else None, |
| "p_correct_given_finished": _mean(s["correct"] for s in fin), |
| "p_correct_given_cut": _mean(s["correct"] for s in cut), |
| "no_answer_finished_share": _mean((s.get("pred") is None and not s["cut"]) for s in S), |
| "loop_end_share_of_cut": _mean((s["verbatim_loop_end"] or s["repetitive_end"]) for s in cut), |
| "tail_dup8_cut_mean": _mean(s["tail_dup8"] for s in cut), |
| "median_tokens_correct": _median([s["tokens"] for s in S if s["correct"]]), |
| } |
| srcs = sorted({r.get("source") for r in recs if r.get("source")}) |
| if len(srcs) > 1: |
| out["by_source"] = {src: {k: v for k, v in summarize([r for r in recs if r.get("source") == src]).items() |
| if k in ("problems", "samples", "acc", "acc_se", "cut_share", "mean_tokens")} |
| for src in srcs} |
| return out |
|
|
|
|
| def _median(xs): |
| xs = sorted(xs) |
| if not xs: |
| return None |
| m = len(xs) // 2 |
| return xs[m] if len(xs) % 2 else (xs[m - 1] + xs[m]) / 2 |
|
|
|
|
| def compare(a: list[dict], b: list[dict]) -> dict: |
| """Paired comparison of run b against run a on the problems both answered.""" |
| A = {r["id"]: r for r in a} |
| B = {r["id"]: r for r in b} |
| ids = [i for i in A if i in B] |
| d = [_mean(s["correct"] for s in B[i]["samples"]) - _mean(s["correct"] for s in A[i]["samples"]) for i in ids] |
| a_any = {i for i in ids if any(s["correct"] for s in A[i]["samples"])} |
| b_any = {i for i in ids if any(s["correct"] for s in B[i]["samples"])} |
| return { |
| "problems": len(ids), |
| "gain": _mean(d), |
| "gain_se": _se(d), |
| "problems_better": sum(x > 0 for x in d), |
| "problems_worse": sum(x < 0 for x in d), |
| |
| "lost": len(a_any - b_any), |
| "rescued": len(b_any - a_any), |
| } |
|
|
|
|
| def load_run(run_dir) -> list[dict]: |
| p = Path(run_dir) / "merged.jsonl" |
| return read_jsonl(p) |
|
|
