backends/transformers_ctc.py

"""
Hugging Face CTC ASR (Wav2Vec2, HuBERT-CTC, etc.): logits -> argmax -> decode.
"""

from __future__ import annotations

from collections.abc import Callable

import numpy as np
import torch

from ._model_utils import attach_params


def build_transcriber(model_id: str, device_str: str) -> tuple[Callable[..., str], Callable[[], None]]:
    from transformers import AutoModelForCTC, AutoProcessor

    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    model = AutoModelForCTC.from_pretrained(model_id, trust_remote_code=True).to(device_str)
    model.eval()

    def transcribe(audio_np: np.ndarray, sampling_rate: int = 16000) -> str:
        inputs = processor(
            audio_np,
            sampling_rate=sampling_rate,
            return_tensors="pt",
            padding=True,
        )
        inputs = {k: v.to(device_str) if hasattr(v, "to") else v for k, v in inputs.items()}
        with torch.no_grad():
            logits = model(**inputs).logits
        pred_ids = torch.argmax(logits, dim=-1)
        text = processor.batch_decode(pred_ids)[0]
        return text

    attach_params(transcribe, model)

    def cleanup() -> None:
        nonlocal model, processor
        del model, processor

    return transcribe, cleanup