asr_backends.py

"""
Pluggable ASR backends.

The whole point of this module: scoring.py must not know or care where the
transcript came from. Every backend returns the same Transcript object, so you
can flip ASR_BACKEND and compare error tables across engines with identical
scoring logic.

    ASR_BACKEND=groq      GROQ_API_KEY=gsk_...     # free CPU Space, per-second billing
    ASR_BACKEND=zerogpu                            # free ZeroGPU Space, transformers
    ASR_BACKEND=local                              # your machine / paid GPU Space
    ASR_BACKEND=auto                               # default: groq > zerogpu > local

Why three: CTranslate2 (the engine under faster-whisper) does not allocate
through PyTorch's CUDA allocator, so it will not see a GPU under ZeroGPU's
fork-based allocation. ZeroGPU therefore needs the transformers path.
"""

from __future__ import annotations

import math
import os
import shutil
import subprocess
import tempfile
import threading
import time
from dataclasses import dataclass, field
from typing import Protocol, runtime_checkable

# ---------------------------------------------------------------------------
# Shared data types
# ---------------------------------------------------------------------------

NAN = float("nan")


@dataclass
class Word:
    text: str
    start: float = 0.0
    end: float = 0.0
    prob: float = NAN          # NaN when the backend cannot report confidence


@dataclass
class Transcript:
    text: str
    words: list[Word] = field(default_factory=list)
    speech_seconds: float = 0.0     # voiced time, not wall-clock file length
    backend: str = ""
    model: str = ""
    latency_s: float = 0.0
    has_confidence: bool = False


@runtime_checkable
class ASRBackend(Protocol):
    name: str
    supports_word_confidence: bool

    def describe(self) -> str: ...

    def transcribe(self, audio_path: str, lang: str, hint: str | None = None) -> Transcript: ...


# ---------------------------------------------------------------------------
# Audio preprocessing
# ---------------------------------------------------------------------------

_HAS_FFMPEG = shutil.which("ffmpeg") is not None


def to_16k_mono_flac(path: str) -> str:
    """
    Downsample to 16 kHz mono FLAC. Whisper resamples to 16 kHz internally
    anyway, so this is lossless for accuracy but shrinks the upload ~10x --
    which matters because hosted endpoints cap request size (Groq: 25 MB on
    the free tier) and a phone recording is often 48 kHz stereo.
    """
    if not _HAS_FFMPEG:
        return path
    out = tempfile.NamedTemporaryFile(delete=False, suffix=".flac").name
    try:
        subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", path,
             "-ar", "16000", "-ac", "1", "-c:a", "flac", out],
            check=True, timeout=120,
        )
        return out
    except Exception:
        return path


def _voiced_from_segments(segments) -> float:
    total = 0.0
    for s in segments:
        start = s.get("start") if isinstance(s, dict) else getattr(s, "start", None)
        end = s.get("end") if isinstance(s, dict) else getattr(s, "end", None)
        if start is not None and end is not None:
            total += max(0.0, float(end) - float(start))
    return total


# ---------------------------------------------------------------------------
# Backend 1: Groq  (recommended for a free CPU Space)
# ---------------------------------------------------------------------------

GROQ_URL = "https://api.groq.com/openai/v1/audio/transcriptions"

# large-v3 is meaningfully better than turbo on Hindi; turbo is ~2.8x cheaper
# and fine for English. Override with GROQ_MODEL_HI / GROQ_MODEL_EN.
GROQ_MODELS = {
    "hi": os.getenv("GROQ_MODEL_HI", "whisper-large-v3"),
    "en": os.getenv("GROQ_MODEL_EN", "whisper-large-v3-turbo"),
}


class GroqBackend:
    name = "groq"
    supports_word_confidence = False     # OpenAI-compatible API returns no logprobs

    def __init__(self, api_key: str | None = None, timeout: float = 90.0):
        self.api_key = api_key or os.environ["GROQ_API_KEY"]
        self.timeout = timeout

    def describe(self) -> str:
        return f"Groq API — {GROQ_MODELS['hi']} (hi) / {GROQ_MODELS['en']} (en)"

    def transcribe(self, audio_path: str, lang: str, hint: str | None = None) -> Transcript:
        import httpx

        model = GROQ_MODELS.get(lang, "whisper-large-v3")
        sent = to_16k_mono_flac(audio_path)
        t0 = time.perf_counter()

        data = {
            "model": model,
            "language": lang,
            "response_format": "verbose_json",
            "timestamp_granularities[]": ["word", "segment"],
            "temperature": "0",
        }
        if hint:
            data["prompt"] = hint[:220]

        payload = None
        last_error: Exception | None = None
        for attempt in range(3):
            try:
                with open(sent, "rb") as fh:
                    resp = httpx.post(
                        GROQ_URL,
                        headers={"Authorization": f"Bearer {self.api_key}"},
                        data=data,
                        files={"file": (os.path.basename(sent), fh, "audio/flac")},
                        timeout=self.timeout,
                    )
                if resp.status_code in (429, 500, 502, 503):
                    raise RuntimeError(f"transient {resp.status_code}: {resp.text[:200]}")
                resp.raise_for_status()
                payload = resp.json()
                break
            except Exception as exc:
                last_error = exc
                if attempt == 2:
                    raise RuntimeError(f"Groq transcription failed: {exc}") from exc
                time.sleep(1.5 * (attempt + 1))     # backoff; free tier is rate-limited
        assert payload is not None, last_error

        if sent != audio_path:
            try:
                os.unlink(sent)
            except OSError:
                pass

        words = [
            Word(w.get("word", "").strip(), float(w.get("start", 0)), float(w.get("end", 0)))
            for w in (payload.get("words") or [])
        ]
        voiced = _voiced_from_segments(payload.get("segments") or [])
        if not voiced and words:
            voiced = words[-1].end - words[0].start

        return Transcript(
            text=(payload.get("text") or "").strip(),
            words=words,
            speech_seconds=voiced,
            backend=self.name,
            model=model,
            latency_s=round(time.perf_counter() - t0, 2),
            has_confidence=False,
        )


# ---------------------------------------------------------------------------
# Backend 2: local faster-whisper
# ---------------------------------------------------------------------------

LOCAL_MODELS = {
    "hi": {"fast": "small", "balanced": "medium", "accurate": "large-v3"},
    "en": {"fast": "distil-small.en", "balanced": "distil-medium.en",
           "accurate": "distil-large-v3"},
}


def _pick_device() -> tuple[str, str]:
    try:
        import torch

        if torch.cuda.is_available():
            major = torch.cuda.get_device_capability()[0]
            return "cuda", "int8_float16" if major >= 7 else "float16"
    except Exception:
        pass
    return "cpu", "int8"


class LocalBackend:
    name = "local"
    supports_word_confidence = True       # the reason to keep this backend around

    def __init__(self, tier: str | None = None):
        self.tier = tier or os.getenv("LOCAL_TIER", "accurate")
        self.device, self.compute_type = _pick_device()
        self._cache: dict[str, object] = {}
        self._lock = threading.Lock()

    def describe(self) -> str:
        return (f"faster-whisper {self.tier} on {self.device} ({self.compute_type})"
                f" — word confidence available")

    def _model(self, lang: str):
        name = LOCAL_MODELS[lang][self.tier]
        with self._lock:
            if name not in self._cache:
                from faster_whisper import WhisperModel

                self._cache[name] = WhisperModel(
                    name, device=self.device, compute_type=self.compute_type,
                    cpu_threads=os.cpu_count() or 4,
                )
            return self._cache[name], name

    def transcribe(self, audio_path: str, lang: str, hint: str | None = None) -> Transcript:
        model, name = self._model(lang)
        t0 = time.perf_counter()

        segments, _info = model.transcribe(
            audio_path,
            language=lang,
            beam_size=1,                        # greedy: ~3x faster, ~1% WER cost
            word_timestamps=True,
            condition_on_previous_text=False,   # stop one bad segment poisoning the rest
            vad_filter=True,                    # skip silence in learner recordings
            vad_parameters={"min_silence_duration_ms": 400},
            initial_prompt=hint,
            temperature=0.0,
        )

        chunks, words, voiced = [], [], 0.0
        for seg in segments:                    # generator -- work happens here
            chunks.append(seg.text)
            voiced += max(0.0, seg.end - seg.start)
            for w in (seg.words or []):
                words.append(Word(w.word.strip(), w.start, w.end,
                                  getattr(w, "probability", NAN)))

        return Transcript(
            text=" ".join(chunks).strip(), words=words, speech_seconds=voiced,
            backend=self.name, model=name,
            latency_s=round(time.perf_counter() - t0, 2), has_confidence=True,
        )


# ---------------------------------------------------------------------------
# Backend 3: ZeroGPU (transformers)
# ---------------------------------------------------------------------------

try:
    import spaces  # type: ignore

    _gpu = spaces.GPU
except Exception:                              # not on a ZeroGPU Space
    def _gpu(func=None, duration=None):        # no-op passthrough
        if func is None:
            return lambda f: f
        return func

ZERO_MODELS = {
    "hi": os.getenv("ZERO_MODEL_HI", "openai/whisper-large-v3"),
    "en": os.getenv("ZERO_MODEL_EN", "distil-whisper/distil-large-v3"),
}


class ZeroGPUBackend:
    """
    Loads on CPU at import, moves to CUDA inside the @spaces.GPU call. ZeroGPU
    forks a GPU-attached process per call, so the .to("cuda") must happen
    inside the decorated function, not at module scope.
    """

    name = "zerogpu"
    supports_word_confidence = False

    def __init__(self):
        self._cache: dict[str, tuple] = {}
        self._lock = threading.Lock()

    def describe(self) -> str:
        return f"ZeroGPU transformers — {ZERO_MODELS['hi']} (hi) / {ZERO_MODELS['en']} (en)"

    def _load(self, lang: str):
        name = ZERO_MODELS[lang]
        with self._lock:
            if name not in self._cache:
                import torch
                from transformers import (AutoProcessor,
                                          WhisperForConditionalGeneration)

                proc = AutoProcessor.from_pretrained(name)
                model = WhisperForConditionalGeneration.from_pretrained(
                    name, torch_dtype=torch.float16, low_cpu_mem_usage=True,
                )
                self._cache[name] = (model, proc, name)
            return self._cache[name]

    def transcribe(self, audio_path: str, lang: str, hint: str | None = None) -> Transcript:
        model, proc, name = self._load(lang)
        t0 = time.perf_counter()
        text, words, voiced = self._run(model, proc, audio_path, lang, hint)
        return Transcript(
            text=text, words=words, speech_seconds=voiced,
            backend=self.name, model=name,
            latency_s=round(time.perf_counter() - t0, 2), has_confidence=False,
        )

    @staticmethod
    @_gpu(duration=90)
    def _run(model, proc, audio_path, lang, hint):
        import torch
        from transformers import pipeline

        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if device == "cuda" else torch.float32
        model = model.to(device=device, dtype=dtype)

        asr = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=proc.tokenizer,
            feature_extractor=proc.feature_extractor,
            torch_dtype=dtype,
            device=device,
            chunk_length_s=30,
            batch_size=8,           # batched long-form: the big transformers speedup
        )
        kwargs = {"language": lang, "task": "transcribe", "num_beams": 1}
        if hint:
            kwargs["prompt_ids"] = proc.get_prompt_ids(hint[:220], return_tensors="pt").to(device)

        out = asr(audio_path, return_timestamps="word", generate_kwargs=kwargs)

        words, voiced = [], 0.0
        for ch in out.get("chunks", []) or []:
            ts = ch.get("timestamp") or (None, None)
            start, end = (ts[0] or 0.0), (ts[1] or 0.0)
            words.append(Word(ch.get("text", "").strip(), float(start), float(end)))
            voiced += max(0.0, float(end) - float(start))
        return out.get("text", "").strip(), words, voiced


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

_backend: ASRBackend | None = None
_factory_lock = threading.Lock()


def get_backend() -> ASRBackend:
    """Resolve once per process. ASR_BACKEND=auto prefers Groq, then ZeroGPU."""
    global _backend
    with _factory_lock:
        if _backend is not None:
            return _backend

        choice = os.getenv("ASR_BACKEND", "auto").lower()
        if choice == "auto":
            if os.getenv("GROQ_API_KEY"):
                choice = "groq"
            elif os.getenv("SPACES_ZERO_GPU") or os.getenv("ZEROGPU"):
                choice = "zerogpu"
            else:
                choice = "local"

        _backend = {"groq": GroqBackend, "zerogpu": ZeroGPUBackend,
                    "local": LocalBackend}[choice]()
        return _backend


def is_nan(x: float) -> bool:
    return isinstance(x, float) and math.isnan(x)