src/turn_detection/runtime/predictor.py

"""Endpoint predictor interfaces and ONNX implementation."""

from __future__ import annotations

import json
import math
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Protocol, runtime_checkable

from .controller import ControllerConfig
from .features import FrontendConfig, log_mel_spectrogram, normalize_waveform, resample_waveform
from .types import Prediction


@runtime_checkable
class EndpointPredictor(Protocol):
    def predict(self, audio: Any, sample_rate: int) -> Prediction:
        """Estimate the probability that the current user turn is complete."""


@dataclass(frozen=True, slots=True)
class ModelMetadata:
    model_name: str
    architecture: str
    frontend: FrontendConfig = field(default_factory=FrontendConfig)
    threshold: float = 0.60
    controller: ControllerConfig = field(default_factory=ControllerConfig)
    input_features_name: str = "input_features"
    frame_mask_name: str | None = "frame_mask"
    endpoint_output_name: str | None = None
    output_type: str = "logits"
    model_version: str = "unknown"
    development_only: bool = False
    training_status: str = "unknown"
    data_scope: str | None = None
    data_revision: str | None = None
    parameter_count: int | None = None

    def __post_init__(self) -> None:
        if not 0.0 <= self.threshold <= 1.0:
            raise ValueError("threshold must be in [0, 1]")
        if not math.isclose(self.controller.endpoint_threshold, self.threshold, abs_tol=1e-12):
            raise ValueError("controller endpoint_threshold must match threshold")
        if self.output_type not in {"logits", "probability"}:
            raise ValueError("output_type must be logits or probability")
        if self.parameter_count is not None and self.parameter_count <= 0:
            raise ValueError("parameter_count must be positive when provided")

    @classmethod
    def from_path(cls, path: str | Path) -> ModelMetadata:
        payload = json.loads(Path(path).read_text(encoding="utf-8"))
        frontend = FrontendConfig(**payload.pop("frontend", {}))
        controller_payload = payload.pop("controller", None)
        if controller_payload is None:
            threshold = float(payload.get("threshold", 0.60))
            controller = ControllerConfig(
                endpoint_threshold=threshold,
                long_pause_threshold=max(0.0, threshold - 0.18),
            )
        else:
            controller = ControllerConfig(**controller_payload)
        return cls(frontend=frontend, controller=controller, **payload)

    def to_dict(self) -> dict[str, Any]:
        return {
            "model_name": self.model_name,
            "architecture": self.architecture,
            "frontend": self.frontend.to_dict(),
            "threshold": self.threshold,
            "controller": asdict(self.controller),
            "input_features_name": self.input_features_name,
            "frame_mask_name": self.frame_mask_name,
            "endpoint_output_name": self.endpoint_output_name,
            "output_type": self.output_type,
            "model_version": self.model_version,
            "development_only": self.development_only,
            "training_status": self.training_status,
            "data_scope": self.data_scope,
            "data_revision": self.data_revision,
            "parameter_count": self.parameter_count,
        }


class OnnxEndpointPredictor:
    """Batch-one ONNX Runtime predictor with serialized preprocessing contract."""

    def __init__(
        self,
        model_path: str | Path,
        metadata_path: str | Path | None = None,
        *,
        intra_op_threads: int = 1,
    ) -> None:
        try:
            import onnxruntime as ort
        except ImportError as exc:  # pragma: no cover - optional dependency
            raise RuntimeError("Install the 'demo' or 'export' extra for ONNX inference") from exc
        self.model_path = Path(model_path)
        if not self.model_path.is_file():
            raise FileNotFoundError(self.model_path)
        metadata_file = (
            Path(metadata_path)
            if metadata_path
            else self.model_path.with_name("model_metadata.json")
        )
        self.metadata = ModelMetadata.from_path(metadata_file)
        options = ort.SessionOptions()
        options.intra_op_num_threads = intra_op_threads
        options.inter_op_num_threads = 1
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(
            str(self.model_path),
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        self._input_names = {item.name for item in self.session.get_inputs()}

    def predict(self, audio: Any, sample_rate: int) -> Prediction:
        import numpy as np

        started = time.perf_counter_ns()
        features, frame_mask = log_mel_spectrogram(audio, sample_rate, self.metadata.frontend)
        feeds = {self.metadata.input_features_name: features[None, :, :]}
        if self.metadata.frame_mask_name and self.metadata.frame_mask_name in self._input_names:
            feeds[self.metadata.frame_mask_name] = frame_mask[None, :]
        output_names = (
            [self.metadata.endpoint_output_name] if self.metadata.endpoint_output_name else None
        )
        raw = self.session.run(output_names, feeds)[0]
        value = float(np.asarray(raw).reshape(-1)[0])
        probability = (
            1.0 / (1.0 + math.exp(-value)) if self.metadata.output_type == "logits" else value
        )
        elapsed_ms = (time.perf_counter_ns() - started) / 1_000_000
        return Prediction(
            endpoint_probability=min(1.0, max(0.0, probability)),
            inference_ms=elapsed_ms,
            model_name=self.metadata.model_name,
        )


class HeuristicDevelopmentPredictor:
    """Clearly labelled fallback used only when exported weights are absent.

    It exists so the UI and controller can be exercised before training. Scores
    from this class must never be reported as model results.
    """

    def predict(self, audio: Any, sample_rate: int) -> Prediction:
        import numpy as np

        started = time.perf_counter_ns()
        samples = resample_waveform(normalize_waveform(audio), sample_rate, 16_000)
        tail = samples[-4_000:] if len(samples) >= 4_000 else samples
        previous = samples[-12_000:-4_000] if len(samples) >= 12_000 else samples
        tail_rms = float(np.sqrt(np.mean(tail * tail) + 1e-12))
        previous_rms = float(np.sqrt(np.mean(previous * previous) + 1e-12))
        drop = max(0.0, min(1.0, 1.0 - tail_rms / max(previous_rms, 1e-4)))
        duration_signal = min(1.0, len(samples) / 16_000 / 2.0)
        probability = 0.15 + 0.55 * drop + 0.20 * duration_signal
        elapsed_ms = (time.perf_counter_ns() - started) / 1_000_000
        return Prediction(
            endpoint_probability=min(0.95, max(0.05, probability)),
            inference_ms=elapsed_ms,
            model_name="heuristic-development-only",
        )


def load_predictor(model_path: str | Path | None) -> EndpointPredictor:
    if model_path is not None and Path(model_path).is_file():
        return OnnxEndpointPredictor(model_path)
    return HeuristicDevelopmentPredictor()