engine/speech.py

"""Local speech-to-text support for short player utterances."""

from __future__ import annotations

from pathlib import Path
from typing import Any
import wave

MODEL_ID = "nvidia/nemotron-speech-streaming-en-0.6b"
MAX_RECORDING_SECONDS = 10
MAX_RECORDING_DRIFT_SECONDS = 0.75
MIN_RECORDING_SECONDS = 0.15

_ASR_MODEL: Any | None = None


class SpeechTranscriptionError(RuntimeError):
    """Raised when local speech transcription cannot produce usable text."""


def _load_asr_model() -> Any:
    global _ASR_MODEL
    if _ASR_MODEL is None:
        try:
            import nemo.collections.asr as nemo_asr
        except Exception as exc:
            raise SpeechTranscriptionError(
                "Speech recognition is not installed. Install NVIDIA NeMo ASR locally."
            ) from exc
        try:
            _ASR_MODEL = nemo_asr.models.ASRModel.from_pretrained(model_name=MODEL_ID)
        except Exception as exc:
            raise SpeechTranscriptionError(
                f"Could not load local speech model {MODEL_ID}."
            ) from exc
    return _ASR_MODEL


def _duration_seconds(audio_path: Path) -> float | None:
    try:
        with wave.open(str(audio_path), "rb") as fh:
            rate = fh.getframerate()
            if rate <= 0:
                return None
            return fh.getnframes() / float(rate)
    except (wave.Error, EOFError):
        return None


def _validate_audio_path(audio_path: str | Path) -> Path:
    path = Path(audio_path)
    if not path.exists() or not path.is_file():
        raise SpeechTranscriptionError("No audio file was captured.")

    duration = _duration_seconds(path)
    if duration is None:
        raise SpeechTranscriptionError("No speech detected in the recording.")
    if duration < MIN_RECORDING_SECONDS:
        raise SpeechTranscriptionError("No speech detected in the recording.")
    if duration > MAX_RECORDING_SECONDS + MAX_RECORDING_DRIFT_SECONDS:
        raise SpeechTranscriptionError(
            "Recording is too long. Try again with a shorter line."
        )
    return path


def _first_transcript(result: Any) -> str:
    if result is None:
        return ""
    if isinstance(result, str):
        return result
    if hasattr(result, "text"):
        return str(result.text)
    if isinstance(result, dict):
        for key in ("text", "transcript"):
            if key in result:
                return str(result[key])
    if isinstance(result, (list, tuple)):
        if not result:
            return ""
        return _first_transcript(result[0])
    return str(result)


def _normalize_transcript(result: Any) -> str:
    transcript = " ".join(_first_transcript(result).split())
    if not transcript:
        raise SpeechTranscriptionError("No speech detected in the recording.")
    return transcript


def transcribe_audio(audio_path: str | Path, *, model: Any | None = None) -> str:
    path = _validate_audio_path(audio_path)
    asr_model = model if model is not None else _load_asr_model()
    try:
        result = asr_model.transcribe([str(path)], batch_size=1)
    except SpeechTranscriptionError:
        raise
    except Exception as exc:
        raise SpeechTranscriptionError("Speech transcription failed.") from exc
    return _normalize_transcript(result)