pipeline/asr.py

"""Local ASR transcription with faster-whisper, run on the downloaded video file.

Used when the pipeline downloads the source video (via the RapidAPI fast-downloader):
faster-whisper reads the media directly (PyAV decodes the audio track), so no separate
audio extraction is needed. Returns ``[{start, end, text}]`` segments — the same shape the
caption-based path produces — so the rest of the pipeline is unchanged.

Model size / device are configurable via env: ``WHISPER_MODEL`` (default ``base``),
``WHISPER_DEVICE`` (default ``cpu``), ``WHISPER_COMPUTE`` (default ``int8`` — fastest on
CPU). ``base``/``int8`` is a sensible CPU default; use ``small``/``medium`` (or a GPU) for
higher accuracy.
"""
from __future__ import annotations

import os
from functools import lru_cache


class ASRError(RuntimeError):
    """Raised when local transcription fails."""


@lru_cache(maxsize=1)
def _model():
    from faster_whisper import WhisperModel

    size = os.environ.get("WHISPER_MODEL", "base")
    device = os.environ.get("WHISPER_DEVICE", "cpu")
    compute = os.environ.get("WHISPER_COMPUTE", "int8")
    return WhisperModel(size, device=device, compute_type=compute)


def transcribe_file(media_path: str, language: str | None = "en",
                    progress=None) -> list[dict]:
    """Transcribe ``media_path`` (audio or video) into ``[{start, end, text}]``.

    ``language`` pins the language (faster + avoids misdetection); pass None to
    auto-detect.
    """
    if not os.path.exists(media_path):
        raise ASRError(f"media file not found: {media_path}")
    try:
        model = _model()
        # No vad_filter: it depends on onnxruntime (Silero VAD) that isn't always
        # present and, when broken, silently drops every segment. beam_size=1 for speed.
        segments, info = model.transcribe(media_path, language=language, beam_size=1)
    except Exception as exc:
        raise ASRError(f"faster-whisper failed: {type(exc).__name__}: {exc}") from exc

    total = float(getattr(info, "duration", 0) or 0)
    out: list[dict] = []
    for s in segments:  # generator — transcription happens as we iterate
        text = (s.text or "").strip()
        if text:
            out.append({"start": float(s.start), "end": float(s.end), "text": text})
        if progress and total:
            progress(min(1.0, (s.end or 0) / total), desc="Transcribing (faster-whisper)")
    if not out:
        raise ASRError("faster-whisper produced no speech segments.")
    return out