src/turn_detector/inference.py

"""
Final inference module for the chosen architecture:

    Audio -> 16kHz mono -> Whisper Tiny frozen encoder -> mean pooling
    -> Logistic Regression -> P(END) -> threshold/hysteresis -> END/CONTINUE

Chosen based on measured evidence (EXP-001 through EXP-004 — see
experiments/EXPERIMENTS.md and docs/RESULTS.md): Whisper Tiny frozen
encoder + mean pooling + Logistic Regression scored F1=0.693 on a 75-clip
real-audio validation split, a directional improvement over the acoustic
baseline's F1=0.575 measured on an independent sample from the same
dataset (not a paired comparison — see docs/RESULTS.md for the exact
wording this project commits to about that comparison).

SANDBOX NOTE (this specific development environment, not a statement about
the architecture itself): `torch` and `transformers` are not installed
here and cannot be installed (no network route to PyPI — the same
`host_not_allowed` block documented throughout this project). This module
is written to run correctly in a normal ML environment where those are
available (imports are lazy/guarded so the rest of the codebase — and this
file's own non-Whisper logic — can still be imported and tested here).
The trained classifier head (`models/whisper_classifier.joblib`) and a
real 250-clip embedding matrix from the actual EXP-004 Colab run
(`artifacts/exp004/whisper_embeddings.npz` if present) ARE available here,
so the classifier stage of this pipeline is genuinely testable in this
sandbox even though the Whisper encoder stage is not — see
tests/test_inference.py for exactly what is and isn't verified where.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

import numpy as np

DEFAULT_WHISPER_MODEL = "openai/whisper-tiny"
DEFAULT_CLASSIFIER_PATH = Path(__file__).resolve().parents[2] / "models" / "whisper_classifier.joblib"
TARGET_SAMPLE_RATE = 16_000
WHISPER_HIDDEN_SIZE = 384  # whisper-tiny d_model — used only for a defensive shape check


class InferenceError(RuntimeError):
    pass


class AudioValidationError(InferenceError):
    """Raised for malformed/invalid audio input — never silently coerced."""
    pass


# ---------------------------------------------------------------------------
# Turn decision logic (threshold + optional min-confidence + optional
# debounce/hysteresis across repeated calls in a streaming-style usage)
# ---------------------------------------------------------------------------

@dataclass
class TurnDecisionConfig:
    """Why this layer exists (per project requirement): a turn detector
    that flips to END the instant P(END) crosses 0.5 on a single noisy
    call will end turns prematurely on one uncertain pause. This layer
    lets a caller require sustained evidence (min_confidence, and/or
    consecutive-call debounce) before committing to END, without training
    a second model — it's a deterministic rule on top of the classifier's
    output probability, same category of design as EXP-001's threshold
    and EXP-006 in the experiment ladder (temporal decision logic).
    """
    end_threshold: float = 0.5
    min_confidence: Optional[float] = None  # if set, |P(END)-0.5|*2 must exceed this to commit either way
    debounce_consecutive_calls: int = 1     # >1 requires N consecutive END calls before committing to END


class TurnDecisionState:
    """Stateful wrapper for streaming-style repeated calls. Not a model —
    pure bookkeeping over a sequence of P(END) values from TurnDetector.
    """

    def __init__(self, config: TurnDecisionConfig = TurnDecisionConfig()):
        self.config = config
        self._consecutive_end_calls = 0

    def reset(self) -> None:
        self._consecutive_end_calls = 0

    def update(self, end_probability: float) -> dict:
        """Feed one new P(END) observation, return the decision at this
        point in the (simulated) stream, accounting for debounce.
        """
        cfg = self.config
        raw_decision_end = end_probability >= cfg.end_threshold

        if cfg.min_confidence is not None:
            confidence = abs(end_probability - 0.5) * 2  # 0 at p=0.5, 1 at p=0 or p=1
            if confidence < cfg.min_confidence:
                self._consecutive_end_calls = 0
                return {"decision": "CONTINUE", "reason": "below_min_confidence", "end_probability": end_probability}

        if raw_decision_end:
            self._consecutive_end_calls += 1
        else:
            self._consecutive_end_calls = 0

        if self._consecutive_end_calls >= cfg.debounce_consecutive_calls and raw_decision_end:
            return {"decision": "END", "reason": "threshold_and_debounce_met", "end_probability": end_probability}
        return {"decision": "CONTINUE", "reason": "threshold_not_met_or_debouncing", "end_probability": end_probability}


def decide(end_probability: float, config: TurnDecisionConfig = TurnDecisionConfig()) -> str:
    """Stateless single-call decision (no debounce applied — debounce is
    only meaningful across repeated calls, use TurnDecisionState for that).
    """
    if config.min_confidence is not None:
        confidence = abs(end_probability - 0.5) * 2
        if confidence < config.min_confidence:
            return "CONTINUE"
    return "END" if end_probability >= config.end_threshold else "CONTINUE"


# ---------------------------------------------------------------------------
# Audio validation (shared by TurnDetector.predict and the Gradio app)
# ---------------------------------------------------------------------------

def validate_audio(audio: np.ndarray, sr: int) -> None:
    if audio is None:
        raise AudioValidationError("Audio is None.")
    audio = np.asarray(audio)
    if audio.ndim not in (1, 2):
        raise AudioValidationError(f"Audio must be 1D (mono) or 2D (multi-channel), got shape {audio.shape}.")
    if audio.size == 0:
        raise AudioValidationError("Audio is empty (zero samples).")
    if not np.isfinite(audio).all():
        raise AudioValidationError("Audio contains NaN or infinite values.")
    if sr is None or sr <= 0:
        raise AudioValidationError(f"Invalid sample rate: {sr}.")
    duration_sec = audio.shape[0] / sr
    if duration_sec < 0.05:
        raise AudioValidationError(f"Audio too short ({duration_sec*1000:.1f}ms) to be a meaningful clip.")


# ---------------------------------------------------------------------------
# TurnDetector
# ---------------------------------------------------------------------------

@dataclass
class TurnDetectorConfig:
    whisper_model_name: str = DEFAULT_WHISPER_MODEL
    classifier_path: Path = field(default_factory=lambda: DEFAULT_CLASSIFIER_PATH)
    device: Optional[str] = None  # None = auto-detect (cuda if available else cpu)
    decision: TurnDecisionConfig = field(default_factory=TurnDecisionConfig)


class TurnDetector:
    """
    Audio -> 16kHz mono -> Whisper Tiny frozen encoder -> mean pooling
    -> Logistic Regression -> P(END) -> threshold/hysteresis -> END/CONTINUE

    Usage:
        detector = TurnDetector()
        result = detector.predict(audio_array, sr=16000)
        print(result)
        # {"decision": "END", "end_probability": 0.87, "continue_probability": 0.13, "latency_ms": 15.4}

    The Whisper encoder is loaded lazily on first use (not in __init__),
    so constructing a TurnDetector never fails just because torch/
    transformers aren't installed — only calling `.predict()` (or
    `.load_whisper()` explicitly) does, with a clear `InferenceError`.
    The trained classifier head IS loaded eagerly in __init__ (it's a
    small local file, no torch/network needed), so a missing/corrupt
    classifier artifact fails fast at construction time rather than
    silently inside the first `.predict()` call. This also lets `embed_fn`
    be injected for testing (see tests/test_inference.py), so the
    classifier stage — which IS runnable in this project's dev sandbox
    using the real trained classifier and real embeddings from the actual
    EXP-004 Colab run — is genuinely tested end-to-end without needing
    torch/transformers present.
    """

    def __init__(
        self,
        config: TurnDetectorConfig = TurnDetectorConfig(),
        embed_fn: Optional[Callable[[np.ndarray, int], np.ndarray]] = None,
    ):
        self.config = config
        self._whisper_model = None
        self._feature_extractor = None
        self._classifier = None
        self._embed_fn_override = embed_fn  # for testing / alternate backends
        self._decision_state = TurnDecisionState(config.decision)
        # Classifier is loaded eagerly (cheap, local disk file) so its load
        # time is never mistakenly counted as part of a predict() call's
        # reported inference latency. Whisper stays lazy (network/GPU-
        # dependent, and constructing a TurnDetector shouldn't require it
        # just to e.g. inspect config or run classifier-only tests).
        self._load_classifier()

    # -- lazy loading -----------------------------------------------------

    def _load_classifier(self):
        if self._classifier is not None:
            return self._classifier
        import joblib
        path = Path(self.config.classifier_path)
        if not path.exists():
            raise InferenceError(
                f"Classifier not found at {path}. Expected the trained Logistic Regression "
                f"pipeline saved from EXP-004 (models/whisper_classifier.joblib)."
            )
        self._classifier = joblib.load(path)
        return self._classifier

    def load_whisper(self):
        """Explicitly load Whisper Tiny (also called automatically by
        predict() on first use). Raises InferenceError with a clear
        message if torch/transformers aren't available, rather than an
        opaque ImportError.
        """
        if self._whisper_model is not None:
            return
        try:
            import torch
            from transformers import WhisperFeatureExtractor, WhisperModel
        except ImportError as e:
            raise InferenceError(
                "torch/transformers are required to load Whisper Tiny but are not installed "
                "in this environment. Install with `pip install torch transformers` in an "
                "environment with network access. (This is a known limitation of the Claude "
                "sandbox this project was developed in — see docs/INITIAL_ANALYSIS.md.)"
            ) from e

        device = self.config.device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._feature_extractor = WhisperFeatureExtractor.from_pretrained(self.config.whisper_model_name)
        self._whisper_model = WhisperModel.from_pretrained(self.config.whisper_model_name).to(device)
        self._whisper_model.eval()
        for p in self._whisper_model.parameters():
            p.requires_grad = False
        self._device = device

    # -- embedding extraction ---------------------------------------------

    def _extract_embedding(self, audio: np.ndarray, sr: int) -> np.ndarray:
        if self._embed_fn_override is not None:
            return self._embed_fn_override(audio, sr)

        self.load_whisper()
        import torch
        with torch.no_grad():
            inputs = self._feature_extractor(audio, sampling_rate=sr, return_tensors="pt")
            input_features = inputs["input_features"].to(self._device)
            encoder_out = self._whisper_model.encoder(input_features)
            hidden = encoder_out.last_hidden_state
            pooled = hidden.mean(dim=1).squeeze(0).cpu().numpy()
        return pooled

    # -- prediction ---------------------------------------------------------

    def predict(self, audio: np.ndarray, sr: int = TARGET_SAMPLE_RATE) -> dict:
        """Run the full pipeline on one clip. Returns a structured dict —
        never hardcoded, every field computed from the actual input.
        """
        t0 = time.perf_counter()

        validate_audio(audio, sr)
        audio = np.asarray(audio, dtype=np.float32)
        if audio.ndim == 2:
            audio = audio.mean(axis=-1)  # downmix to mono

        if sr != TARGET_SAMPLE_RATE:
            from .audio_io import resample_linear
            audio = resample_linear(audio, sr_in=sr, sr_out=TARGET_SAMPLE_RATE)
            sr = TARGET_SAMPLE_RATE

        embedding = self._extract_embedding(audio, sr)
        embedding = np.asarray(embedding).reshape(1, -1)

        clf = self._load_classifier()
        proba = clf.predict_proba(embedding)[0]
        # sklearn's classes_ ordering determines which column is P(True)=P(END)
        classes = list(clf.named_steps["clf"].classes_) if hasattr(clf, "named_steps") else list(clf.classes_)
        end_idx = classes.index(1) if 1 in classes else classes.index(True)
        end_probability = float(proba[end_idx])
        continue_probability = float(1.0 - end_probability)

        decision = decide(end_probability, self.config.decision)

        latency_ms = (time.perf_counter() - t0) * 1000.0

        return {
            "decision": decision,
            "end_probability": end_probability,
            "continue_probability": continue_probability,
            "latency_ms": latency_ms,
        }

    def predict_streaming_step(self, audio: np.ndarray, sr: int = TARGET_SAMPLE_RATE) -> dict:
        """Like predict(), but routes the decision through the stateful
        TurnDecisionState (debounce/hysteresis across repeated calls) —
        for simulating a sequence of growing-buffer calls. Call
        `.reset_stream()` between independent turns.
        """
        result = self.predict(audio, sr)
        stateful = self._decision_state.update(result["end_probability"])
        result["decision"] = stateful["decision"]
        result["decision_reason"] = stateful["reason"]
        return result

    def reset_stream(self) -> None:
        self._decision_state.reset()