# Held-out loss record via the spec's loss_eval.LossEval on confirmation (test) arms.
# run.py eval_nll stores per-example (NLL sum over weight>0 tokens, count) == loss_eval.example_nll.
import json, loss_eval
spec = json.load(open("spec.json"))
C = [json.loads(l) for l in open("confirm_log.jsonl")]
arms = {}
for e in C:
    if e["stage"] != "eval": continue
    le = loss_eval.LossEval("token_mean")
    le.examples = [(float(s), int(n)) for s, n in e["per_example"]]
    arms[e["arm"]] = le
rep = loss_eval.report("token_mean", arms)
for k, a in rep["arms"].items():
    ev = next(e for e in C if e["stage"] == "eval" and e["arm"] == k)
    assert abs(a["token_mean"] - ev["token_mean"]) <= 1e-9 * ev["token_mean"], k
print("problems:", loss_eval.check(spec, rep))
json.dump(rep, open("loss_eval_record.json", "w"), indent=1)
print(json.dumps({k: (v["token_mean"], v["n_tokens"], v["n_examples"]) for k, v in rep["arms"].items()}))