backends/transformers_pipeline.py

"""
Hugging Face `transformers` ASR pipeline (Whisper, many Wav2Vec2 wrappers, etc.).
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

import numpy as np

from ._audio_utils import safe_pad_audio
from ._model_utils import attach_params


def build_transcriber(model_id: str, device_int: int) -> tuple[Callable[..., str], Callable[[], None]]:
    from transformers import pipeline

    pipe: Any = pipeline(
        "automatic-speech-recognition",
        model=model_id,
        device=device_int,
        trust_remote_code=True,
    )
    pipe({"raw": np.zeros(16000, dtype=np.float32), "sampling_rate": 16000})

    def transcribe(audio_np: np.ndarray, sampling_rate: int = 16000) -> str:
        padded = safe_pad_audio(audio_np)
        out = pipe({"raw": padded, "sampling_rate": sampling_rate})
        return str(out["text"])

    attach_params(transcribe, getattr(pipe, "model", None))

    def cleanup() -> None:
        nonlocal pipe
        del pipe

    return transcribe, cleanup