| """Basic grouped scoring and optional scalar-temperature finalization. Policy stays disabled.""" |
|
|
| import contextlib |
| import json |
| import os |
| import tempfile |
| from pathlib import Path |
|
|
| from mdtrain.budget import BudgetLedger |
| from mdtrain.data import DatasetView, EVALUATION_SPLITS |
| from mdtrain.engine import DecisionModel |
| from mdtrain.fixtures import read_requests |
| from mdtrain.metrics import fit_temperature, score_predictions |
| from mdtrain.storage import atomic_json, report_json, verify_seal |
|
|
|
|
| def prediction_file(path: Path, rows: list[dict]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| descriptor, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) |
| temporary = Path(name) |
| try: |
| with os.fdopen(descriptor, "w") as stream: |
| for row in rows: |
| stream.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n") |
| stream.flush() |
| os.fsync(stream.fileno()) |
| os.replace(temporary, path) |
| finally: |
| temporary.unlink(missing_ok=True) |
|
|
|
|
| def gpu_allowance(device: str, path: Path | None): |
| |
| if device == "cuda" and path is not None: |
| return BudgetLedger(path) |
| return contextlib.nullcontext(None) |
|
|
|
|
| def evaluate_model(model: DecisionModel, view: DatasetView, *, ledger: BudgetLedger | None = None) -> tuple[list[dict], dict]: |
| rows = [] |
| for example, canonical in zip(view.examples, view.rows): |
| if ledger is not None: |
| ledger.admit(60) |
| batch = model.compiler.compile((example.input,)) |
| prediction = model.predict((example.input,))[0] |
| expected_keys = tuple(candidate.key for candidate in example.input.candidates) |
| if prediction["id"] != example.input.id or tuple(prediction["probabilities"]) != expected_keys: |
| raise ValueError("Prediction identity/candidate alignment differs from requested view") |
| rows.append({**prediction, "group_id": example.group_id, |
| "target": expected_keys[example.target_index], "source": example.source, |
| "origin": example.origin, "modality": canonical["modality"], |
| "primitive": example.input.kind, "family": example.family, "ood_axis": example.ood_axis, |
| "position_count": int(batch.expanded_lengths[0])}) |
| if ledger is not None: |
| ledger.commit() |
| metrics = score_predictions(rows) |
| real_rows = [row for row in rows if row["origin"] == "real_human_label"] |
| generated_rows = [row for row in rows if row["origin"] != "real_human_label"] |
| return rows, {"all": metrics, "real": score_predictions(real_rows) if real_rows else None, |
| "generated": score_predictions(generated_rows) if generated_rows else None} |
|
|
|
|
| def predict(root: Path, *, artifact: Path, input_path: Path, media_root: Path, output: Path, |
| device: str, budget_ledger: Path | None = None) -> dict: |
| requests = read_requests(input_path) |
| with gpu_allowance(device, budget_ledger) as ledger: |
| if ledger: |
| ledger.admit(300) |
| model = DecisionModel.load(artifact, root=root, device=device, media_root=media_root) |
| rows = [] |
| for item in requests: |
| if ledger: |
| ledger.admit(60) |
| rows.extend(model.predict((item,))) |
| prediction_file(output, rows) |
| return {"status": "completed", "questions": len(rows), "artifact_id": model.artifact_id, |
| "output": str(output), "policy": "disabled"} |
|
|
|
|
| def evaluate(root: Path, *, artifact: Path, splits: tuple[str, ...], output: Path, |
| device: str, budget_ledger: Path | None = None) -> dict: |
| if not splits or len(set(splits)) != len(splits) or any(split not in EVALUATION_SPLITS for split in splits): |
| raise ValueError("Evaluation requires distinct declared evaluation splits") |
| with gpu_allowance(device, budget_ledger) as ledger: |
| if ledger: |
| ledger.admit(300) |
| model = DecisionModel.load(artifact, root=root, device=device) |
| if any(split.startswith("test_") for split in splits) and not model.finalized: |
| raise ValueError("Locked final-test evaluation requires a finalized paired artifact") |
| result = {"status": "completed", "artifact_id": model.artifact_id, |
| "temperature": model.temperature, "policy": "disabled", "splits": {}} |
| for split in splits: |
| view = DatasetView.open(root, split=split, recipe="mixed") |
| rows, metrics = evaluate_model(model, view, ledger=ledger) |
| prediction_file(output.with_name(output.stem + f"-{split}-predictions.jsonl"), rows) |
| result["splits"][split] = {"view_identity": view.identity, "metrics": metrics} |
| result = report_json(result) |
| atomic_json(output, result) |
| return result |
|
|
|
|
| def finalize(root: Path, *, checkpoint: Path, output: Path, device: str, |
| fit_scalar_temperature: bool = False, budget_ledger: Path | None = None) -> dict: |
| checkpoint = checkpoint.resolve(strict=True) |
| checkpoint_id = verify_seal(checkpoint) |
| state = json.loads((checkpoint / "state.json").read_text()) |
| if state["training"]["status"] != "completed" or state["spec"]["stage"] != "main": |
| raise ValueError("Only a completed main checkpoint may be finalized") |
| with gpu_allowance(device, budget_ledger) as ledger: |
| if ledger: |
| ledger.admit(300) |
| if state["ledger_id"] != ledger.record["id"]: |
| raise ValueError("Finalization must use the same durable budget ledger") |
| model = DecisionModel.load(checkpoint / "inference", root=root, device=device) |
| calibration_identity = None |
| if fit_scalar_temperature: |
| view = DatasetView.open(root, split="calibration", recipe="mixed") |
| logits, targets = [], [] |
| for example in view.examples: |
| if ledger: |
| ledger.admit(60) |
| logits.append(model.logits((example.input,))[0]) |
| targets.append(example.target_index) |
| model.temperature = fit_temperature(logits, targets) |
| model.temperature_fitted = True |
| calibration_identity = view.identity |
| model.finalized = True |
| model.save(output, metadata={"checkpoint_id": checkpoint_id, "training_status": "completed", |
| "calibration_view_id": calibration_identity, "policy": "disabled"}) |
| receipt = {"status": "completed", "artifact": str(output), "artifact_id": model.artifact_id, |
| "checkpoint_id": checkpoint_id, "temperature": model.temperature, |
| "temperature_fitted": fit_scalar_temperature, "calibration_view_id": calibration_identity, |
| "policy": "disabled", "weights_selected_using_final_tests": False} |
| atomic_json(output.with_name(output.name + "-freeze-receipt.json"), receipt) |
| return receipt |
|
|
