File size: 5,458 Bytes
35d483e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | #!/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())
|