scripts/freeze_candidate.py

#!/usr/bin/env python3
"""Freeze model, code, preprocessing, and controller before opening official test."""

from __future__ import annotations

import argparse
import json
import math
import sys
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = ROOT / "src"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))

from turn_detection.provenance import (  # noqa: E402
    build_freeze_policy_binding,
    file_evidence,
)
from turn_detection.runtime import ControllerConfig  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--config", required=True)
    parser.add_argument("--output", default="artifacts/final/frozen_manifest.json")
    parser.add_argument(
        "--long-pause-threshold",
        type=float,
        help="default: endpoint threshold minus 0.18, floored at zero",
    )
    parser.add_argument("--min-silence-ms", type=float, default=200.0)
    parser.add_argument("--relax-after-ms", type=float, default=800.0)
    parser.add_argument("--max-silence-ms", type=float, default=1800.0)
    parser.add_argument("--required-confirmations", type=int, default=1)
    return parser.parse_args()


def _resolve(value: str) -> Path:
    path = Path(value)
    return (path if path.is_absolute() else ROOT / path).resolve()


def main() -> int:
    args = parse_args()
    try:
        import torch
    except ImportError as exc:
        raise SystemExit("freezing a candidate requires PyTorch") from exc

    checkpoint_path = _resolve(args.checkpoint)
    config_path = _resolve(args.config)
    output_path = _resolve(args.output)
    for path in (checkpoint_path, config_path):
        try:
            path.relative_to(ROOT)
        except ValueError as exc:
            raise SystemExit(f"frozen input must stay inside the project: {path}") from exc
        if not path.is_file() or path.is_symlink():
            raise SystemExit(f"frozen input is not a regular file: {path}")
    try:
        output_path.relative_to(ROOT / "artifacts")
    except ValueError as exc:
        raise SystemExit("--output must stay under artifacts/") from exc

    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    if not isinstance(checkpoint, dict):
        raise SystemExit("checkpoint must be a mapping")
    metadata = checkpoint.get("metadata", {})
    run_metadata = metadata.get("run_metadata", {}) if isinstance(metadata, dict) else {}
    training_status = str(run_metadata.get("status", "")) if isinstance(run_metadata, dict) else ""
    if training_status.lower() != "final":
        raise SystemExit(
            "only a checkpoint trained with run.status='final' may be frozen for official test"
        )
    threshold = float(checkpoint.get("threshold", math.nan))
    if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0:
        raise SystemExit("checkpoint has no valid calibrated threshold")
    controller = ControllerConfig(
        endpoint_threshold=threshold,
        long_pause_threshold=(
            args.long_pause_threshold
            if args.long_pause_threshold is not None
            else max(0.0, threshold - 0.18)
        ),
        min_silence_ms=args.min_silence_ms,
        relax_after_ms=args.relax_after_ms,
        max_silence_ms=args.max_silence_ms,
        required_confirmations=args.required_confirmations,
    )
    catalog = json.loads((ROOT / "configs/datasets.json").read_text(encoding="utf-8"))
    source_files = [
        *sorted((ROOT / "src/turn_detection").rglob("*.py")),
        *sorted((ROOT / "scripts").glob("*.py")),
        *sorted((ROOT / "scripts").glob("*.sh")),
        ROOT / "pyproject.toml",
        ROOT / "configs/datasets.json",
        config_path,
        checkpoint_path,
    ]
    unique_files = sorted(set(source_files), key=lambda path: path.as_posix())
    manifest = {
        "format_version": 1,
        "status": "frozen_for_official_test",
        "frozen_at": datetime.now(timezone.utc).isoformat(),
        "checkpoint": file_evidence(checkpoint_path, ROOT),
        "training_config": file_evidence(config_path, ROOT),
        "threshold": threshold,
        "model_config": checkpoint.get("model_config"),
        "frontend": metadata.get("feature_config") if isinstance(metadata, dict) else None,
        "controller": {
            "endpoint_threshold": controller.endpoint_threshold,
            "long_pause_threshold": controller.long_pause_threshold,
            "min_silence_ms": controller.min_silence_ms,
            "relax_after_ms": controller.relax_after_ms,
            "max_silence_ms": controller.max_silence_ms,
            "required_confirmations": controller.required_confirmations,
        },
        "train_data": catalog["train"],
        "official_test": catalog["official_test"],
        "files": [file_evidence(path, ROOT) for path in unique_files],
    }
    manifest["policy_binding"] = build_freeze_policy_binding(manifest)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    print(f"frozen manifest: {output_path.relative_to(ROOT)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())