| |
| """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 ( |
| build_freeze_policy_binding, |
| file_evidence, |
| ) |
| from turn_detection.runtime import ControllerConfig |
|
|
|
|
| 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()) |
|
|
