pipeline/transcriber.py

"""
pipeline/transcriber.py

Hugging Face Moonshine ASR wrapper for the telecalling agent.

The public API intentionally matches the previous ASR wrapper:

    transcriber = Transcriber()
    text = transcriber.transcribe(sample_rate, audio_np)

Internally this uses UsefulSensors/moonshine-tiny through Transformers:
AutoFeatureExtractor prepares audio features, AutoTokenizer decodes generated
tokens, and MoonshineForConditionalGeneration generates transcripts in memory.
"""

import io
import logging
import threading
import time
from typing import Optional, Sequence

import numpy as np

try:
    import torch
    _TORCH_AVAILABLE = True
except ImportError:  # pragma: no cover
    torch = None
    _TORCH_AVAILABLE = False

import requests
import soundfile as sf
import socket
from requests.exceptions import ConnectionError as RequestsConnectionError

from config import (
    TRANSCRIBE_MODEL_ID,
    TRANSCRIBE_DEVICE,
    HF_TOKEN,
    TRANSCRIBE_LOCAL_ONLY,
)

logger = logging.getLogger(__name__)


class Transcriber:
    """
    Lazy-loading wrapper around UsefulSensors/moonshine-tiny.

    Thread-safe: a single instance can be shared across the whole app.
    """

    # Recommended by the Moonshine model card to reduce hallucination loops.
    _MAX_TOKENS_PER_SECOND = 6.5

    def __init__(self):
        self._model = None
        self._feature_extractor = None
        self._tokenizer = None
        self._device = None
        self._dtype = None
        self._sample_rate = None
        self._lock = threading.Lock()
        self._loaded = False
        self._use_remote = False
        # None = unknown, True = reachable, False = unreachable
        self._remote_ok: Optional[bool] = None

    def transcribe(self, sample_rate: int, audio: np.ndarray) -> str:
        """
        Transcribe a single utterance.

        Parameters
        ----------
        sample_rate : int
            Sample rate of `audio`.
        audio : np.ndarray
            Mono float32 PCM in [-1.0, 1.0].

        Returns
        -------
        str
            Transcribed text, stripped of leading/trailing whitespace.
            Returns "" on error so the pipeline can continue gracefully.
        """
        if audio is None or len(audio) == 0:
            return ""

        self._ensure_loaded()

        audio = self._validate_audio(audio)
        if audio is None:
            return ""

        try:
            t0 = time.perf_counter()
            text = self._generate_text([audio], [sample_rate])[0]
            elapsed = time.perf_counter() - t0

            duration = len(audio) / sample_rate
            rtfx = duration / elapsed if elapsed > 0 else 0
            logger.info(
                f"Transcribed {duration:.2f}s audio in {elapsed:.2f}s "
                f"(RTFx {rtfx:.1f}x): '{text[:80]}{'...' if len(text) > 80 else ''}'"
            )
            return text

        except Exception as exc:
            logger.error(f"Transcription failed: {exc}", exc_info=True)
            return ""

    def transcribe_batch(
        self, utterances: list[tuple[int, np.ndarray]]
    ) -> list[str]:
        """
        Transcribe multiple utterances in one generation call.

        Parameters
        ----------
        utterances : list of (sample_rate, audio_np) tuples

        Returns
        -------
        list of str, one entry per input, "" on individual errors
        """
        if not utterances:
            return []

        self._ensure_loaded()

        sample_rates = [sr for sr, _ in utterances]
        audio_arrays = [self._validate_audio(a) for _, a in utterances]

        audio_arrays = [
            a if a is not None else np.zeros(sample_rates[i], dtype=np.float32)
            for i, a in enumerate(audio_arrays)
        ]

        try:
            return self._generate_text(audio_arrays, sample_rates)
        except Exception as exc:
            logger.error(f"Batch transcription failed: {exc}", exc_info=True)
            return [""] * len(utterances)

    def unload(self):
        """
        Release GPU memory. Call at end of a call session if memory is tight.
        Model will be reloaded lazily on the next call.
        """
        with self._lock:
            if self._loaded:
                self._model = None
                self._feature_extractor = None
                self._tokenizer = None
                self._device = None
                self._dtype = None
                self._sample_rate = None
                self._loaded = False
                if _TORCH_AVAILABLE and torch.cuda.is_available():
                    torch.cuda.empty_cache()
                logger.info("Moonshine transcriber unloaded; VRAM freed.")

    @property
    def is_loaded(self) -> bool:
        return self._loaded

    @property
    def device(self) -> Optional[str]:
        return str(self._device) if self._device else None

    def _ensure_loaded(self):
        """Load model + processor exactly once, thread-safely."""
        if self._loaded:
            return

        with self._lock:
            if self._loaded:
                return

            if not _TORCH_AVAILABLE:
                if TRANSCRIBE_LOCAL_ONLY:
                    logger.warning(
                        "PyTorch is unavailable and offline-only transcription is enabled; "
                        "remote ASR is disabled."
                    )
                    self._use_remote = False
                    self._loaded = True
                    return
                logger.warning(
                    "PyTorch is unavailable in this environment; "
                    "using remote ASR fallback."
                )
                self._use_remote = True
                self._loaded = True
                return

            self._load()

    def prefetch(self) -> None:
        """Download Moonshine into the local HF cache before first inference."""
        if TRANSCRIBE_LOCAL_ONLY:
            logger.info(
                "TRANSCRIBE_LOCAL_ONLY is enabled; skipping startup prefetch."
            )
            return

        import os

        hf_model_cache = os.path.expanduser(
            f"~/.cache/huggingface/hub/models--{TRANSCRIBE_MODEL_ID.replace('/', '--')}"
        )

        if os.path.exists(hf_model_cache):
            logger.info("Moonshine ASR cache already exists; skipping startup prefetch.")
            return

        logger.info(
            "Prefetching Moonshine ASR model into the Hugging Face cache..."
        )

        if not _TORCH_AVAILABLE:
            try:
                import inspect
                from huggingface_hub import snapshot_download

                snapshot_kwargs = {
                    "repo_id": TRANSCRIBE_MODEL_ID,
                    "local_files_only": False,
                    "repo_type": "model",
                }
                sig = inspect.signature(snapshot_download)
                if "token" in sig.parameters:
                    snapshot_kwargs["token"] = HF_TOKEN or None
                elif "use_auth_token" in sig.parameters:
                    snapshot_kwargs["use_auth_token"] = HF_TOKEN or None

                snapshot_download(**snapshot_kwargs)
                logger.info("Moonshine ASR cache download completed.")
            except Exception as exc:
                logger.error(
                    "Failed to prefetch Moonshine model cache without PyTorch: %s",
                    exc,
                    exc_info=True,
                )
            return

        try:
            self._install_torchvision_import_stub_if_needed()
            from transformers import AutoFeatureExtractor, AutoTokenizer
            from transformers import MoonshineForConditionalGeneration

            token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}

            AutoFeatureExtractor.from_pretrained(
                TRANSCRIBE_MODEL_ID,
                local_files_only=False,
                **token_kwargs,
            )
            AutoTokenizer.from_pretrained(
                TRANSCRIBE_MODEL_ID,
                local_files_only=False,
                **token_kwargs,
            )
            MoonshineForConditionalGeneration.from_pretrained(
                TRANSCRIBE_MODEL_ID,
                local_files_only=False,
                **token_kwargs,
            )
            logger.info("Moonshine ASR prefetch completed successfully.")
        except Exception as exc:
            logger.error(
                "Startup prefetch of Moonshine failed: %s",
                exc,
                exc_info=True,
            )

    def _load(self):
        """Pull Moonshine from Hugging Face and move it to the best device."""
        import os

        from transformers import AutoFeatureExtractor, AutoTokenizer

        self._install_torchvision_import_stub_if_needed()
        from transformers import MoonshineForConditionalGeneration

        logger.info(f"Loading Moonshine ASR ({TRANSCRIBE_MODEL_ID})...")
        t0 = time.perf_counter()

        token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
        hf_model_cache = os.path.expanduser(
            f"~/.cache/huggingface/hub/models--{TRANSCRIBE_MODEL_ID.replace('/', '--')}"
        )
        local_only = TRANSCRIBE_LOCAL_ONLY or os.path.exists(hf_model_cache)

        if TRANSCRIBE_LOCAL_ONLY:
            logger.info(
                "Offline-only transcription is enabled; Moonshine must already be cached locally."
            )
        elif not os.path.exists(hf_model_cache):
            logger.info(
                "Moonshine cache not found locally; downloading from Hugging Face Hub."
            )

        if local_only and not os.path.exists(hf_model_cache):
            logger.error(
                "Local Moonshine model cache not found; offline-only transcription "
                "cannot proceed."
            )
            raise RuntimeError("Local Moonshine model is not cached locally.")

        self._feature_extractor = AutoFeatureExtractor.from_pretrained(
            TRANSCRIBE_MODEL_ID,
            local_files_only=local_only,
            **token_kwargs,
        )
        self._tokenizer = AutoTokenizer.from_pretrained(
            TRANSCRIBE_MODEL_ID,
            local_files_only=local_only,
            **token_kwargs,
        )
        self._sample_rate = int(self._feature_extractor.sampling_rate)

        if torch.cuda.is_available():
            try:
                self._device = torch.device(TRANSCRIBE_DEVICE)
                self._dtype = torch.float16
                logger.info(f"CUDA available; loading Moonshine in float16 on {self._device}")
            except Exception:
                self._device = torch.device("cpu")
                self._dtype = torch.float32
                logger.warning("CUDA device init failed; falling back to CPU float32")
        else:
            self._device = torch.device("cpu")
            self._dtype = torch.float32
            logger.info("No CUDA; loading Moonshine on CPU in float32")

        self._model = MoonshineForConditionalGeneration.from_pretrained(
            TRANSCRIBE_MODEL_ID,
            torch_dtype=self._dtype,
            low_cpu_mem_usage=True,
            local_files_only=local_only,
            **token_kwargs,
        ).to(self._device)
        self._model.eval()

        elapsed = time.perf_counter() - t0
        logger.info(
            f"Moonshine ASR ready on {self._device} "
            f"(dtype={self._dtype}, sample_rate={self._sample_rate}, loaded in {elapsed:.1f}s)"
        )

        if torch.cuda.is_available() and self._device.type == "cuda":
            allocated = torch.cuda.memory_allocated(self._device) / 1024**3
            reserved = torch.cuda.memory_reserved(self._device) / 1024**3
            logger.info(
                f"VRAM after load: allocated={allocated:.2f} GB, reserved={reserved:.2f} GB"
            )

        self._loaded = True

    def _generate_text(
        self,
        audio_arrays: Sequence[np.ndarray],
        sample_rates: Sequence[int],
    ) -> list[str]:
        """Run Moonshine generation and decode transcripts."""
        if self._use_remote:
            return [self._remote_transcribe(audio, sr) for audio, sr in zip(audio_arrays, sample_rates)]

        if self._model is None:
            logger.warning("Local ASR model is unavailable; skipping transcription.")
            return ["" for _ in audio_arrays]

        prepared = [
            self._resample(audio, sr, self._sample_rate)
            for audio, sr in zip(audio_arrays, sample_rates)
        ]

        inputs = self._feature_extractor(
            prepared,
            return_tensors="pt",
            sampling_rate=self._sample_rate,
            padding=True,
        )
        inputs = inputs.to(self._device, self._dtype)

        with torch.inference_mode():
            seq_lens = inputs.attention_mask.sum(dim=-1)
            token_limit_factor = self._MAX_TOKENS_PER_SECOND / self._sample_rate
            max_length = max(1, int((seq_lens * token_limit_factor).max().item()))
            generated_ids = self._model.generate(**inputs, max_length=max_length)

        return [
            self._tokenizer.decode(ids, skip_special_tokens=True).strip()
            for ids in generated_ids
        ]

    def _remote_transcribe(self, audio: np.ndarray, sample_rate: int) -> str:
        """Use the Hugging Face Inference API to transcribe audio if local Torch is unavailable."""
        if not HF_TOKEN:
            logger.warning("HF_TOKEN is not set; falling back to mock transcription.")
            return ""

        if requests is None:
            logger.warning(
                "requests is unavailable; cannot perform remote transcription."
            )
            return ""

        endpoint = "https://api-inference.huggingface.co/models/openai/whisper-small"

        if self._remote_ok is False:
            logger.debug("Remote HF has previously failed; skipping remote transcription.")
            return ""

        try:
            headers = {
                "Authorization": f"Bearer {HF_TOKEN}",
                "Accept": "application/json",
            }

            with io.BytesIO() as buffer:
                sf.write(buffer, audio, samplerate=sample_rate, format="WAV")
                buffer.seek(0)
                response = requests.post(
                    endpoint,
                    headers=headers,
                    data=buffer.read(),
                    timeout=90,
                )

            response.raise_for_status()
            payload = response.json()
            if isinstance(payload, dict) and "text" in payload:
                return payload["text"].strip()
            return str(payload)
        except RequestsConnectionError as exc:
            # Connection errors often indicate network/DNS issues — stop retrying
            logger.error(f"Remote transcription connection failed: {exc}", exc_info=True)
            self._remote_ok = False
            return ""
        except Exception as exc:
            logger.error(f"Remote transcription failed: {exc}", exc_info=True)
            return ""

    @staticmethod
    def _install_torchvision_import_stub_if_needed() -> None:
        """
        Keep audio-only Moonshine usable when torchvision is installed but broken.

        Transformers 5.x imports generic image/video utilities while importing
        modeling classes. A mismatched torchvision wheel can raise before any ASR
        code runs, even though Moonshine does not use torchvision at inference.
        """
        try:
            import torchvision  # noqa: F401
            return
        except Exception:
            pass

        import importlib.machinery
        import sys
        import types

        for name in (
            "torchvision",
            "torchvision.transforms",
            "torchvision.transforms.v2",
            "torchvision.transforms.v2.functional",
            "torchvision.io",
        ):
            sys.modules.pop(name, None)

        torchvision = types.ModuleType("torchvision")
        torchvision.__spec__ = importlib.machinery.ModuleSpec("torchvision", None)
        transforms = types.ModuleType("torchvision.transforms")
        transforms.__spec__ = importlib.machinery.ModuleSpec("torchvision.transforms", None)
        transforms_v2 = types.ModuleType("torchvision.transforms.v2")
        transforms_v2.__spec__ = importlib.machinery.ModuleSpec(
            "torchvision.transforms.v2", None
        )
        transforms_v2_functional = types.ModuleType("torchvision.transforms.v2.functional")
        transforms_v2_functional.__spec__ = importlib.machinery.ModuleSpec(
            "torchvision.transforms.v2.functional", None
        )
        torchvision_io = types.ModuleType("torchvision.io")
        torchvision_io.__spec__ = importlib.machinery.ModuleSpec("torchvision.io", None)

        class InterpolationMode:
            NEAREST = 0
            NEAREST_EXACT = 0
            BILINEAR = 2
            BICUBIC = 3
            BOX = 4
            HAMMING = 5
            LANCZOS = 1

        transforms.InterpolationMode = InterpolationMode
        transforms.v2 = transforms_v2
        transforms_v2.functional = transforms_v2_functional
        torchvision.transforms = transforms
        torchvision.io = torchvision_io

        sys.modules["torchvision"] = torchvision
        sys.modules["torchvision.transforms"] = transforms
        sys.modules["torchvision.transforms.v2"] = transforms_v2
        sys.modules["torchvision.transforms.v2.functional"] = transforms_v2_functional
        sys.modules["torchvision.io"] = torchvision_io

    @staticmethod
    def _validate_audio(audio: np.ndarray) -> Optional[np.ndarray]:
        """Ensure audio is mono float32 in [-1.0, 1.0]."""
        if audio is None or len(audio) == 0:
            return None

        audio = np.array(audio, dtype=np.float32)

        if audio.ndim == 2:
            audio = audio.mean(axis=1)
        elif audio.ndim != 1:
            logger.warning(f"Unexpected audio shape {audio.shape}; skipping")
            return None

        if audio.max() > 1.0 or audio.min() < -1.0:
            audio = audio / 32768.0

        return np.clip(audio, -1.0, 1.0)

    @staticmethod
    def _resample(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
        """Simple linear interpolation resample for occasional sample-rate mismatch."""
        if orig_sr == target_sr:
            return audio.astype(np.float32, copy=False)
        if len(audio) == 0:
            return audio.astype(np.float32)

        ratio = target_sr / orig_sr
        new_length = max(1, int(len(audio) * ratio))
        return np.interp(
            np.linspace(0, len(audio) - 1, new_length),
            np.arange(len(audio)),
            audio,
        ).astype(np.float32)


_transcriber: Optional[Transcriber] = None


def get_transcriber() -> Transcriber:
    """Return the module-level singleton, creating it if needed."""
    global _transcriber
    if _transcriber is None:
        _transcriber = Transcriber()
    return _transcriber


def _smoke_test_offline():
    """Validate audio preprocessing and singleton behavior without loading the model."""
    import math

    logging.basicConfig(level=logging.INFO)
    logger.info("Running offline smoke test (pre-processing only)...")

    sr = 16000

    stereo_int16 = (np.random.randn(sr, 2) * 32767).astype(np.int16)
    result = Transcriber._validate_audio(stereo_int16)
    assert result is not None
    assert result.ndim == 1, f"Expected mono, got shape {result.shape}"
    assert result.dtype == np.float32
    assert result.max() <= 1.0
    assert result.min() >= -1.0
    logger.info("Stereo int16 to mono float32 normalization")

    mono_float = np.sin(2 * math.pi * 440 * np.linspace(0, 1, sr)).astype(np.float32)
    result = Transcriber._validate_audio(mono_float)
    assert result is not None and result.shape == (sr,)
    logger.info("Mono float32 passthrough")

    result = Transcriber._validate_audio(np.array([]))
    assert result is None
    logger.info("Empty input returns None")

    t1 = get_transcriber()
    t2 = get_transcriber()
    assert t1 is t2
    logger.info("Module singleton")

    logger.info("Offline smoke test PASSED")
    logger.info(
        "\nTo run the full model test:\n"
        "  python -c \"from pipeline.transcriber import get_transcriber; "
        "import numpy as np; t=get_transcriber(); "
        "audio=np.zeros(16000,dtype='float32'); print(repr(t.transcribe(16000,audio)))\""
    )


if __name__ == "__main__":
    _smoke_test_offline()