"""Assemble initialization_provenance, evaluation, decision observations, activation and metrics from saved evidence.""" import json, os, glob, hashlib from research_exchange.research_decisions import paired_interval ARMS = [f"rank{r}_lr{lr}" for r in (1, 32) for lr in ("1e-4", "3e-4", "1e-3")] + [a for a in ("rank1_lr3e-5", "rank1_lr3e-3", "rank32_lr3e-5", "rank32_lr3e-3") if os.path.exists(f"ev/{a}")] BASE_CKPT = "Qwen/Qwen2.5-3B@3aab1f1954e9cc14eb9509a215f9e5ca08227a9b" REQ = {"learning_rate": None, "beta1": 0.9, "beta2": 0.95, "eps": 1e-08} def lr_of(a): return float(a.split("_lr")[1]) def jl(p): return [json.loads(l) for l in open(p)] out = {} # initialization provenance prov = {"version": 1, "arms": []} for a in ARMS: p = f"ev/{a}/init_provenance.json" if not os.path.exists(p): continue r = json.load(open(p)); info = r["get_info_before_first_update"] calls = json.load(open(f"ev/{a}/optim_calls.json")) if os.path.exists(f"ev/{a}/optim_calls.json") else [] prov["arms"].append({ "name": a, "initialization": {"base_checkpoint": BASE_CKPT, "source": "server_response", "adapter_sha256": None, "adapter_unavailable_reason": r.get("adapter_error"), "server_model_id": info["model_id"], "server_lora_rank": info["lora_rank"], "evidence_ref": f"{a}-init_provenance.json", "before_first_update": True}, "restore_count": 0, "optimizer_checks": [{"after_restore": 0, "actual": info["optimizer_config"]["type"], "source": "server_response", "actual_config": {k: v for k, v in info["optimizer_config"].items() if k != "type"}, "client_sent_adam_params_first_step": calls[0] if calls else None, "evidence_ref": f"{a}-init_provenance.json", "before_next_update": True}], "initial_output": {"available": True, "batch_sha256": r["batch_sha256"], "output_sha256": r["output_sha256"], "evidence_ref": f"{a}-init_provenance.json", "before_first_update": True}}) rp = f"ev/{a}/init_provenance_restore1.json" if os.path.exists(rp): rr = json.load(open(rp)); ri = rr["get_info_before_first_update"]; req = json.load(open(f"ev/{a}/restore1_request.json")) rc = json.load(open(f"ev/{a}/optim_calls_restore1.json")) arm = prov["arms"][-1]; arm["restore_count"] = 1 arm["restores"] = [{"after_restore": 1, "reason": "410 model_lost (shared Spindle engine terminated) during first attempt", "restored_from": req["restored_from"], "method": "create_training_client_from_state_with_optimizer (stock recipe resume)", "server_lora_rank": ri["lora_rank"], "server_model_id": ri["model_id"], "fixed_batch_output_sha256_after_restore": rr["output_sha256"], "evidence_ref": f"{a}-init_provenance_restore1.json"}] arm["optimizer_checks"].append({"after_restore": 1, "actual": ri["optimizer_config"]["type"], "source": "server_response", "actual_config": {k: v for k, v in ri["optimizer_config"].items() if k != "type"}, "client_sent_adam_params_first_step": rc[0] if rc else None, "evidence_ref": f"{a}-init_provenance_restore1.json", "before_next_update": True}) out["initialization_provenance"] = prov # training metrics per arm arms = {} for a in ARMS: m = f"runs/{a}/metrics.jsonl" if not os.path.exists(m): continue rows = [x for x in jl(m) if "env/all/correct" in x] if not rows: continue arms[a] = {"train_steps_completed": len(rows), "train_correct_step0": rows[0]["env/all/correct"], "train_correct_last": rows[-1]["env/all/correct"], "frac_mixed_step0": rows[0]["env/all/by_group/frac_mixed"], "learning_rate": lr_of(a), "lora_rank": int(a[4:a.index("_")])} for f in glob.glob("ev/score_*.json"): s = json.load(open(f)); arms.setdefault(s["arm"], {})[f"{s['split']}_accuracy"] = s["accuracy"] out["arms"] = arms # evaluation + decision evaluators, scores = {}, [] for f in sorted(glob.glob("ev/score_*.json")): s = json.load(open(f)); evaluators[s["evaluator_id"]] = s["provenance"] sel = json.load(open("ev/selection.json")) if os.path.exists("ev/selection.json") else None for f in sorted(glob.glob("ev/score_selection_*.json")): s = json.load(open(f)) scores.append({"arm": s["arm"], "metric": "selection_accuracy", "value": s["accuracy"], "evaluator_id": s["evaluator_id"], "seed": 0, "predictions_ref": f"pred_selection_{s['arm']}.jsonl"}) comparisons = [{"metric": "selection_accuracy", "arms": sorted({s["arm"] for s in scores})}] if scores else [] obs = [] if sel: roles = {"treatment": sel["ranks"]["1"]["selected_arm"], "baseline": sel["ranks"]["32"]["selected_arm"], "base": "base"} have = all(os.path.exists(f"ev/pred_confirmation_{v}.jsonl") for v in roles.values()) if have: per = {} for role, arm in roles.items(): rows = jl(f"ev/pred_confirmation_{arm}.jsonl"); per[role] = {r["item_id"]: r["value"] for r in rows} eid = rows[0]["evaluator_id"] for r in rows: obs.append({"item_id": r["item_id"], "arm": role, "seed": 0, "metric": "accuracy", "value": r["value"], "evaluator_id": eid}) mean = sum(per[role].values()) / len(per[role]) scores.append({"arm": role, "metric": "accuracy", "value": mean, "evaluator_id": eid, "seed": 0, "predictions_ref": f"pred_confirmation_{arm}.jsonl", "checkpoint_arm": arm}) comparisons.append({"metric": "accuracy", "arms": ["treatment", "baseline"]}) comparisons.append({"metric": "accuracy", "arms": ["treatment", "baseline", "base"]}) items = sorted(per["treatment"], key=lambda x: int(x.split("/")[-1])) diffs = [per["treatment"][i] - per["baseline"][i] for i in items] iv = paired_interval(diffs, -1.0, 1.0, 0.05) out["primary"] = {"interval": iv, "treatment_acc": sum(per["treatment"].values()) / len(items), "baseline_acc": sum(per["baseline"].values()) / len(items), "base_acc": sum(per["base"].values()) / len(items), "n": len(items), "roles": roles} with open("ev/scored-observations.jsonl", "w") as f: for o in obs: f.write(json.dumps(o, sort_keys=True) + "\n") out["observations_sha256"] = hashlib.sha256(open("ev/scored-observations.jsonl", "rb").read()).hexdigest() out["evaluation"] = {"version": 1, "evaluators": evaluators, "scores": scores, "comparisons": comparisons} json.dump(out, open("ev/report_parts.json", "w"), indent=1) print(json.dumps({k: v for k, v in out.items() if k not in ("evaluation", "initialization_provenance")}, indent=1))