code/paired_stats.py

"""Paired per-item statistics for a judged round (ours minus TinyStories-33M), with bootstrap 95% CIs.

    python paired_stats.py eval/<condition>/judge_key.json eval/<condition>/judge_replies.txt
"""
import csv, io, json, sys
import numpy as np

key = json.load(open(sys.argv[1])); rows = {}; hdr = None
for line in open(sys.argv[2]).read().splitlines():
    line = line.strip().strip("`")
    if line.startswith("item,label,"):
        hdr = line.split(","); continue
    if hdr and line:
        v = next(csv.reader(io.StringIO(line)))
        if len(v) == len(hdr) and v[0] in key:
            r = dict(zip(hdr, v)); rows.setdefault(r["item"], {})[key[r["item"]][r["label"].strip().upper()]] = r
O, T = "ts300-10m", "TinyStories-33M"  # "ts300-10m" is the label ts300.py gives our model, whatever its size
rng = np.random.default_rng(0)
for m in ["grammar_errors", "consistency_errors", "grammar", "consistency", "plot", "ending"]:
    d = np.array([float(x[O][m]) - float(x[T][m]) for x in rows.values() if len(x) == 2])
    bs = [rng.choice(d, d.size).mean() for _ in range(5000)]
    print(f"{m:20s} ours-ref {d.mean():+.2f}  95% CI [{np.percentile(bs, 2.5):+.2f}, {np.percentile(bs, 97.5):+.2f}]  n={d.size}")