| """Resumable panel scoring, fit-only calibration and one-shot held-out reports.""" |
|
|
| import gzip |
| import hashlib |
| import json |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
| from scipy.optimize import minimize_scalar |
|
|
| from ._vendor.semantics import listed_probs, p_yes |
| from .api import TASKS |
| from .artifacts import ADAPTER_SHA, HEADS_SHA, digest, sha256 |
|
|
|
|
| def load_panel(directory): |
| directory = Path(directory) |
| manifest = json.loads((directory / "manifest.json").read_text()) |
| if ( |
| manifest["adapter_sha256"] != ADAPTER_SHA |
| or manifest["heads_sha256"] != HEADS_SHA |
| or manifest["readout_mode"] != "four_collapsed" |
| ): |
| raise ValueError("Panel belongs to a different checkpoint or answer semantics") |
| raw = gzip.decompress((directory / "jobs.json.gz").read_bytes()) |
| if hashlib.sha256(raw).hexdigest() != manifest["jobs_sha256"]: |
| raise ValueError("Panel jobs checksum mismatch") |
| jobs = json.loads(raw) |
| if len({r["id"] for r in jobs}) != len(jobs): |
| raise ValueError("Duplicate panel branch IDs") |
| return jobs, manifest |
|
|
|
|
| def score_panel(model, panel, output): |
| """Atomically persist each document so interruption never requires rescoring it.""" |
| jobs, manifest = load_panel(panel) |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=True) |
| identity = { |
| "runtime": model.identity, |
| "panel_sha256": manifest["jobs_sha256"], |
| "panel_role": Path(panel).name.split("-")[0], |
| } |
| meta = output / "identity.json" |
| if meta.exists() and json.loads(meta.read_text()) != identity: |
| raise ValueError("Cannot resume with different model code, weights or panel") |
| meta.write_text(json.dumps(identity, indent=2)) |
| documents = defaultdict(list) |
| for row in jobs: |
| documents[row["document_key"]].append(row) |
| for key, group in documents.items(): |
| path = output / (key + ".json") |
| if path.exists(): |
| record = json.loads(path.read_text()) |
| body = {k: v for k, v in record.items() if k != "sha256"} |
| if ( |
| record["sha256"] != digest(body) |
| or record["identity"] != digest(identity) |
| or [r["id"] for r in record["rows"]] != [r["id"] for r in group] |
| ): |
| raise ValueError("Corrupt or mismatched resumed document") |
| continue |
| parts = group[0].get("parts") or [{"text": group[0]["doc"]}] |
| if any((r.get("parts") or [{"text": r["doc"]}]) != parts for r in group): |
| raise ValueError("Document key aliases different sources") |
| with model.prefill(parts) as state: |
| rows = [] |
| for job in group: |
| result = model.engine.ask(state._data, job["block"], job["n"], job["head_key"]) |
| rows.append({**result, **{k: job[k] for k in ("id", "task", "gold", "n", "question_id")}}) |
| body = { |
| "identity": digest(identity), |
| "rows": rows, |
| "prefix_tokens": state.prefix_tokens, |
| "prefill_seconds": state._data["prefill_seconds"], |
| } |
| temp = path.with_suffix(".tmp") |
| temp.write_text(json.dumps({**body, "sha256": digest(body)})) |
| temp.replace(path) |
| print("Scored " + key + " " + str(len(rows)) + " branches", flush=True) |
| completed = { |
| "identity": digest(identity), |
| "documents": len(documents), |
| "branches": len(jobs), |
| "files": {key + ".json": sha256(output / (key + ".json")) for key in documents}, |
| } |
| (output / "complete.json").write_text(json.dumps(completed, indent=2)) |
|
|
|
|
| def read_scores(directory): |
| directory = Path(directory) |
| identity = json.loads((directory / "identity.json").read_text()) |
| completed = json.loads((directory / "complete.json").read_text()) |
| if completed["identity"] != digest(identity): |
| raise ValueError("Score identity mismatch") |
| rows = [] |
| for name, checksum in completed["files"].items(): |
| p = directory / name |
| if not p.resolve().is_relative_to(directory.resolve()) or sha256(p) != checksum: |
| raise ValueError("Score checksum mismatch") |
| record = json.loads(p.read_text()) |
| rows.extend(record["rows"]) |
| if len(rows) != completed["branches"]: |
| raise ValueError("Incomplete score set") |
| return rows, identity |
|
|
|
|
| def unit(row, temperature=1.0): |
| logits = row["letter_logits"] |
| task = row["task"] |
| gold = row["gold"] |
| if task in ("boolean", "entity", "multilabel"): |
| p = p_yes(logits, temperature) |
| return [1 - p, p], int(gold == 0) |
| width = row["n"] - 2 if row["head_key"].endswith("choiceR") else row["n"] |
| if not isinstance(gold, int) or not 0 <= gold < width: |
| return None, None |
| return listed_probs(logits, width, temperature).tolist(), gold |
|
|
|
|
| def fit_calibration(scores, output, *, panel_role): |
| if panel_role != "fit": |
| raise ValueError("Temperature fitting accepts fit panels only") |
| rows, identity = read_scores(scores) |
| if identity["panel_role"] != "fit": |
| raise ValueError("Scores were not generated from a fit panel") |
| output = Path(output) |
| if output.exists(): |
| raise FileExistsError("Calibration artifacts are immutable") |
| temperatures, losses = {}, {} |
| for task in TASKS: |
| selected = [r for r in rows if r["task"] == task and unit(r)[0] is not None] |
| if not selected: |
| raise ValueError("No fit examples for " + task) |
|
|
| def loss(log_t, selected=selected): |
| t = float(np.exp(log_t)) |
| return float(np.mean([-np.log(max(unit(r, t)[0][unit(r, t)[1]], 1e-300)) for r in selected])) |
|
|
| fit = minimize_scalar(loss, bounds=(np.log(0.05), np.log(20)), method="bounded") |
| temperatures[task] = float(np.exp(fit.x)) |
| losses[task] = {"before": loss(0.0), "after": float(fit.fun), "units": len(selected)} |
| payload = { |
| "schema": "solomon-mlx-temperature-v1", |
| "runtime": identity["runtime"]["fingerprint"], |
| "temperatures": temperatures, |
| "fit_panel_sha256": identity["panel_sha256"], |
| "losses": losses, |
| "selection_role": "fit", |
| "heldout_used": False, |
| } |
| output.write_text(json.dumps({**payload, "sha256": digest(payload)}, indent=2)) |
| return payload |
|
|
|
|
| def compare_rows(mlx_rows, cuda_rows, *, temperatures=None, reference_temperatures=None): |
| temperatures = temperatures or dict.fromkeys(TASKS, 1.0) |
| reference_temperatures = reference_temperatures or dict.fromkeys(TASKS, 1.0) |
| reference = {r["id"]: r for r in cuda_rows} |
| if len(reference) != len(cuda_rows) or set(reference) != {r["id"] for r in mlx_rows}: |
| raise ValueError("Comparison panels have different or duplicate branch IDs") |
| units = [] |
| questions = defaultdict(list) |
| for row in mlx_rows: |
| other = {**row, "letter_logits": reference[row["id"]]["letter_logits"]} |
| p, gold = unit(row, temperatures[row["task"]]) |
| q, _ = unit(other, reference_temperatures[row["task"]]) |
| if p is None: |
| continue |
| left, right = int(np.argmax(p)), int(np.argmax(q)) |
| item = { |
| "agreement": left == right, |
| "mlx_correct": left == gold, |
| "cuda_correct": right == gold, |
| "probability_drift": float(np.max(np.abs(np.asarray(p) - q))), |
| } |
| units.append(item) |
| questions[row["question_id"]].append(item) |
| if not units: |
| raise ValueError("No defined comparison targets") |
| agreement = float(np.mean([r["agreement"] for r in units])) |
| question_agreement = float(np.mean([all(x["agreement"] for x in r) for r in questions.values()])) |
| mlx_accuracy = float(np.mean([all(x["mlx_correct"] for x in r) for r in questions.values()])) |
| cuda_accuracy = float(np.mean([all(x["cuda_correct"] for x in r) for r in questions.values()])) |
| return { |
| "units": len(units), |
| "questions": len(questions), |
| "unit_decision_agreement": agreement, |
| "question_decision_agreement": float( |
| np.mean([all(x["agreement"] for x in r) for r in questions.values()]) |
| ), |
| "mlx_whole_question_accuracy": mlx_accuracy, |
| "cuda_whole_question_accuracy": cuda_accuracy, |
| "accuracy_degradation_percentage_points": 100 * (cuda_accuracy - mlx_accuracy), |
| "max_probability_drift": max(r["probability_drift"] for r in units), |
| "probability_comparison": { |
| "mlx_temperatures": temperatures, |
| "cuda_temperatures": reference_temperatures, |
| }, |
| "mean_probability_drift": float(np.mean([r["probability_drift"] for r in units])), |
| "quality_gate_passed": agreement >= 0.999 |
| and question_agreement >= 0.999 |
| and cuda_accuracy - mlx_accuracy <= 0.0025, |
| } |
|
|
|
|
| def read_cuda_scores(directory, panel, reference_identity): |
| """Reuse only scores bound to the exact pinned CUDA runtime and panel.""" |
| jobs, manifest = load_panel(panel) |
| directory = Path(directory) |
| result = {} |
| for file in sorted(directory.glob("scores*.json.gz")): |
| payload = json.loads(gzip.decompress(file.read_bytes())) |
| identity = payload["identity"] |
| if not payload["complete"] or identity["runtime"] != reference_identity: |
| raise ValueError("Existing CUDA scores do not match the fresh reference runtime") |
| if identity["manifest"]["jobs_sha256"] != manifest["jobs_sha256"]: |
| raise ValueError("CUDA scores use another panel") |
| for key, row in payload["scores"].items(): |
| if key in result: |
| raise ValueError("Duplicate CUDA score ID") |
| result[key] = row |
| if set(result) != {r["id"] for r in jobs}: |
| raise ValueError("CUDA score set is incomplete") |
| return [{**r, **result[r["id"]]} for r in jobs] |
|
|
|
|
| def select_calibration(fitted, dev_scores, output): |
| fitted, output = Path(fitted), Path(output) |
| if output.exists(): |
| raise FileExistsError("Selected calibration is immutable") |
| fit = json.loads(fitted.read_text()) |
| fit_payload = {k: v for k, v in fit.items() if k != "sha256"} |
| rows, identity = read_scores(dev_scores) |
| if ( |
| fit["sha256"] != digest(fit_payload) |
| or identity["runtime"]["fingerprint"] != fit["runtime"] |
| or identity["panel_role"] != "dev" |
| ): |
| raise ValueError("Calibration or development identity mismatch") |
| temperatures, selection = {}, {} |
| for task in TASKS: |
| selected = [r for r in rows if r["task"] == task and unit(r)[0] is not None] |
| if not selected: |
| raise ValueError("Missing development task " + task) |
|
|
| def loss(t, selected=selected): |
| values = [unit(r, t) for r in selected] |
| return float(np.mean([-np.log(max(p[g], 1e-300)) for p, g in values])) |
|
|
| original, candidate = loss(1.0), loss(fit["temperatures"][task]) |
| temperatures[task] = fit["temperatures"][task] if candidate < original else 1.0 |
| selection[task] = {"untempered_nll": original, "fit_temperature_nll": candidate} |
| payload = { |
| **fit_payload, |
| "temperatures": temperatures, |
| "selection_role": "dev_selected", |
| "fit_artifact_sha256": sha256(fitted), |
| "dev_panel_sha256": identity["panel_sha256"], |
| "development_selection": selection, |
| } |
| output.write_text(json.dumps({**payload, "sha256": digest(payload)}, indent=2)) |
| return payload |
|
|
|
|
| def heldout_report( |
| scores, |
| cuda_directory, |
| panel, |
| calibration, |
| reference, |
| output, |
| *, |
| reference_binding="evaluations/cuda-acceptance/input/serving-binding.json", |
| ): |
| """Evaluate a frozen configuration once; an existing output cannot be replaced.""" |
| output = Path(output) |
| if output.exists(): |
| raise FileExistsError("Held-out report already exists; do not reuse it for selection") |
| rows, identity = read_scores(scores) |
| cal = json.loads(Path(calibration).read_text()) |
| payload = {k: v for k, v in cal.items() if k != "sha256"} |
| if ( |
| cal["sha256"] != digest(payload) |
| or cal["runtime"] != identity["runtime"]["fingerprint"] |
| or cal["selection_role"] != "dev_selected" |
| or identity["panel_role"] != "cert" |
| ): |
| raise ValueError( |
| "Held-out evaluation requires frozen development-selected calibration and cert scores" |
| ) |
| ref = json.loads(Path(reference).read_text()) |
| cuda = read_cuda_scores(cuda_directory, panel, ref["identity"]) |
| binding_path = Path(reference_binding) |
| source_manifest = json.loads((binding_path.parent / "manifest.json").read_text()) |
| if sha256(binding_path) != source_manifest["files"][binding_path.name]: |
| raise ValueError("CUDA acceptance binding checksum mismatch") |
| binding = json.loads(binding_path.read_text()) |
| for key in ( |
| "adapter_sha256", |
| "trained_heads_sha256", |
| "model_sha256", |
| "numerics", |
| "placement", |
| "arithmetic", |
| ): |
| if binding["runtime"][key] != ref["identity"][key]: |
| raise ValueError("CUDA calibration belongs to another reference runtime") |
| reference_temperatures = {task: binding["temperatures"]["models"][task]["temperature"] for task in TASKS} |
| report = { |
| **compare_rows( |
| rows, cuda, temperatures=cal["temperatures"], reference_temperatures=reference_temperatures |
| ), |
| "cuda_calibration_binding_sha256": sha256(binding_path), |
| "runtime": identity["runtime"], |
| "panel_sha256": identity["panel_sha256"], |
| "calibration_sha256": sha256(calibration), |
| "reference_sha256": sha256(reference), |
| "scope": "text-only held-out panel", |
| "image_qualification": False, |
| } |
| output.write_text(json.dumps(report, indent=2)) |
| return report |
|
|
