app.py

"""Gradio demo for Tiny Hinglish Turn Detection.

The UI intentionally labels the heuristic fallback. A missing model must never
look like a trained result in a hiring submission.
"""

from __future__ import annotations

import os
import sys
from functools import lru_cache
from pathlib import Path
from typing import Any

try:
    import spaces
except ModuleNotFoundError as exc:
    if exc.name != "spaces":
        raise

    class _LocalSpaces:
        """No-op compatibility layer for local installs without ZeroGPU."""

        @staticmethod
        def GPU(*, duration: int = 60) -> Any:  # noqa: N802 - mirrors spaces.GPU
            del duration

            def decorate(function: Any) -> Any:
                return function

            return decorate

    spaces = _LocalSpaces()

PROJECT_ROOT = Path(__file__).resolve().parent
SRC_ROOT = PROJECT_ROOT / "src"
if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))

from turn_detection.runtime import ControllerConfig, TurnController, TurnState  # noqa: E402
from turn_detection.runtime.predictor import load_predictor  # noqa: E402

DEFAULT_MODEL_CANDIDATES = (
    PROJECT_ROOT / "artifacts" / "model.onnx",  # Hugging Face Space layout
    PROJECT_ROOT / "model.onnx",  # Hugging Face model-repository layout
    # Curated GitHub/source-checkout layout.
    PROJECT_ROOT / "artifacts" / "partial-shard-warmstart-lr3e4-5ep" / "model.onnx",
)


def _default_model_path() -> Path:
    """Resolve packaged model layouts without hiding a genuinely missing model."""

    for candidate in DEFAULT_MODEL_CANDIDATES:
        if candidate.is_file():
            return candidate
    return DEFAULT_MODEL_CANDIDATES[0]


@lru_cache(maxsize=1)
def get_predictor() -> Any:
    configured = os.environ.get("TURN_MODEL_PATH")
    return load_predictor(configured or _default_model_path())


def _predictor_metadata(predictor: Any) -> Any | None:
    """Return optional exported metadata without coupling the UI to ONNX."""

    return getattr(predictor, "metadata", None)


def _controller_for_ui(
    metadata: Any | None, threshold: float, max_silence_ms: float
) -> ControllerConfig:
    stored = getattr(metadata, "controller", None)
    if not isinstance(stored, ControllerConfig):
        stored = ControllerConfig(
            endpoint_threshold=threshold,
            long_pause_threshold=max(0.0, threshold - 0.18),
        )
    relaxation_delta = stored.endpoint_threshold - stored.long_pause_threshold
    return ControllerConfig(
        endpoint_threshold=threshold,
        long_pause_threshold=max(0.0, threshold - relaxation_delta),
        min_silence_ms=stored.min_silence_ms,
        relax_after_ms=min(stored.relax_after_ms, max_silence_ms),
        max_silence_ms=max_silence_ms,
        required_confirmations=stored.required_confirmations,
    )


def _timeline_html(probability: float, threshold: float, state: TurnState) -> str:
    probability_width = round(probability * 100, 1)
    threshold_left = round(threshold * 100, 1)
    color = "#16a34a" if state is TurnState.END else "#f59e0b"
    return f"""
    <div aria-label="endpoint probability timeline" style="padding: 10px 2px">
      <div style="position:relative;height:24px;background:#e5e7eb;border-radius:12px;overflow:hidden">
        <div style="height:100%;width:{probability_width}%;background:{color}"></div>
        <div title="decision threshold" style="position:absolute;left:{threshold_left}%;top:0;
             height:100%;border-left:3px solid #111827"></div>
      </div>
      <div style="display:flex;justify-content:space-between;font-size:12px;margin-top:4px">
        <span>HOLD · 0</span><span>threshold {threshold:.2f}</span><span>1 · END</span>
      </div>
    </div>
    """


@spaces.GPU(duration=10)
def analyze_turn(
    audio: tuple[int, Any] | None,
    threshold: float,
    silence_ms: float,
    max_silence_ms: float,
) -> tuple[str, dict[str, float], dict[str, Any], str]:
    if audio is None:
        raise ValueError("Record or upload an utterance first")
    sample_rate, samples = audio
    predictor = get_predictor()
    prediction = predictor.predict(samples, int(sample_rate))
    metadata = _predictor_metadata(predictor)
    controller = TurnController(
        _controller_for_ui(metadata, float(threshold), float(max_silence_ms))
    )
    decision = controller.evaluate_pause(prediction, float(silence_ms))
    is_fallback = prediction.model_name == "heuristic-development-only"
    is_development = bool(getattr(metadata, "development_only", False))
    if is_fallback:
        warning = (
            "\n\n⚠️ **Development fallback active:** exported weights are not present; "
            "this score is not a trained-model result."
        )
    elif is_development:
        scope = getattr(metadata, "data_scope", None) or "limited development data"
        warning = (
            "\n\n⚠️ **Development model:** this score comes from an unqualified preview "
            f"trained on {scope}. It is not evidence of real-world Hinglish accuracy."
        )
    else:
        warning = ""
    status = (
        f"## {decision.state.value}\n\n"
        f"Reason: `{decision.reason}` · p(END): **{prediction.endpoint_probability:.3f}**"
        f"{warning}"
    )
    label = {
        "END": prediction.endpoint_probability,
        "HOLD": 1.0 - prediction.endpoint_probability,
    }
    diagnostics = {
        "state": decision.state.value,
        "emit_response": decision.emit_response,
        "reason": decision.reason,
        "model": prediction.model_name,
        "development_only": is_fallback or is_development,
        "training_status": getattr(metadata, "training_status", "fallback"),
        "data_scope": getattr(metadata, "data_scope", None),
        "data_revision": getattr(metadata, "data_revision", None),
        "parameter_count": getattr(metadata, "parameter_count", None),
        "p_end": round(prediction.endpoint_probability, 6),
        "threshold": round(decision.threshold or threshold, 6),
        "assumed_silence_ms": silence_ms,
        "model_inference_ms": round(prediction.inference_ms, 3),
        "sample_rate_hz": int(sample_rate),
        "samples": int(len(samples)),
    }
    return (
        status,
        label,
        diagnostics,
        _timeline_html(
            prediction.endpoint_probability,
            decision.threshold or threshold,
            decision.state,
        ),
    )


def build_demo() -> Any:
    try:
        import gradio as gr
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise RuntimeError("Install the demo dependencies: uv sync --extra demo") from exc

    predictor = get_predictor()
    metadata = _predictor_metadata(predictor)
    default_threshold = float(getattr(metadata, "threshold", 0.60))
    parameter_count = getattr(metadata, "parameter_count", None)
    frontend = getattr(metadata, "frontend", None)
    window_seconds = getattr(frontend, "max_seconds", None)
    model_summary = (
        f"Loaded `{getattr(metadata, 'model_name', 'unknown')}` · "
        f"{int(parameter_count):,} parameters"
        + (f" · {float(window_seconds):g} s suffix window" if window_seconds else "")
        if parameter_count is not None
        else "No exported model metadata is loaded."
    )
    if metadata is not None and bool(getattr(metadata, "development_only", False)):
        prediction_notice = getattr(metadata, "data_scope", None) or "limited development data"
        evidence_notice = (
            "> ⚠️ **Development preview.** Data scope: "
            f"{prediction_notice}. No official-test or collected-Hinglish claim is made."
        )
    elif predictor.__class__.__name__ == "HeuristicDevelopmentPredictor":
        evidence_notice = (
            "> ⚠️ **Heuristic fallback.** Trained weights are absent; outputs are UI-only."
        )
    else:
        evidence_notice = ""

    with gr.Blocks(title="Tiny Hinglish Turn Detector") as demo:
        gr.Markdown(
            "# Tiny Hinglish Turn Detector\n"
            "Audio-native **HOLD vs END** decisions at VAD pause checkpoints. "
            "Try incomplete phrases, fillers, corrections, and complete Shiprocket-style requests.\n\n"
            f"{evidence_notice}\n\n{model_summary}"
        )
        gr.Markdown(
            "### What to record\n\n"
            "Use natural pacing and leave a short pause at the end of each clip. These are prompts, "
            "not included evaluation examples.\n\n"
            "| Expected | Example prompt | Why |\n"
            "|---|---|---|\n"
            "| HOLD | `mera order number hai... umm...` | filler before missing detail |\n"
            "| END | `mera order cancel kar do` | complete request |\n"
            "| HOLD | `haan matlab... kal wala parcel...` | self-repair / continuation |\n"
            "| END | `haan, kal wala parcel reschedule kar do` | complete after filler |\n"
            "| HOLD | `address change karna hai, flat number...` | slot still missing |\n"
            "| END | `address change karke Flat 12B kar do` | slot supplied |"
        )
        with gr.Row():
            with gr.Column(scale=3):
                audio = gr.Audio(
                    sources=["microphone", "upload"],
                    type="numpy",
                    label="Current user turn",
                )
                analyze = gr.Button("Analyze pause checkpoint", variant="primary")
            with gr.Column(scale=2):
                threshold = gr.Slider(
                    0.0,
                    1.0,
                    value=default_threshold,
                    step=0.01,
                    label="END threshold",
                )
                silence_ms = gr.Slider(
                    200,
                    1800,
                    value=300,
                    step=50,
                    label="Silence at checkpoint (ms)",
                )
                max_silence_ms = gr.Slider(
                    800,
                    3000,
                    value=1800,
                    step=100,
                    label="Maximum response timeout (ms)",
                )
        status = gr.Markdown("## Waiting for audio")
        timeline = gr.HTML()
        with gr.Row():
            scores = gr.Label(num_top_classes=2, label="Decision probabilities")
            diagnostics = gr.JSON(label="Runtime diagnostics")
        gr.Markdown(
            "**Interpretation:** HOLD means the agent should keep listening. END means it may respond. "
            "The production controller also imposes a maximum timeout so uncertain predictions cannot wait forever."
        )
        analyze.click(
            fn=analyze_turn,
            inputs=[audio, threshold, silence_ms, max_silence_ms],
            outputs=[status, scores, diagnostics, timeline],
        )
    return demo


if __name__ == "__main__":
    build_demo().queue(default_concurrency_limit=2).launch()