File size: 7,346 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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """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()
|