backends/auto.py

"""
Simplified backend cascade:
    (SpeechBrain shortcut if the repo declares it)
      -> (Granite chat path if the id matches)
      -> pipeline
      -> universal
      -> CTC
      -> SpeechBrain (final fallback).
"""

from __future__ import annotations

from collections.abc import Callable

from . import (
    granite_speech,
    speechbrain_asr,
    transformers_ctc,
    transformers_pipeline,
    universal,
        nemo_asr,
        qwen_asr,
)


def _is_granite_speech_model(model_id: str) -> bool:
    m = model_id.lower()
    return "ibm-granite/" in m and "granite-" in m and "-speech" in m


def _is_nemo_asr_model(model_id: str) -> bool:
    m = model_id.lower()
    return (
        "nvidia/parakeet" in m
        or "parakeet-tdt" in m
        or "nvidia/canary" in m
        or "canary-" in m
    )


def _looks_possibly_speechbrain_model(model_id: str) -> bool:
    m = model_id.lower()
    return m.startswith("speechbrain/") or "speechbrain" in m


def _is_cohere_asr_model(model_id: str) -> bool:
    from .family_resolve import is_cohere_asr_model

    return is_cohere_asr_model(model_id)


def _is_hf_connectivity_cache_error(exc: Exception) -> bool:
    msg = str(exc)
    needles = (
        "we couldn't connect to 'https://huggingface.co'",
        "couldn't find them in the cached files",
        "offline mode",
        "connection error",
        "temporary failure in name resolution",
    )
    lower = msg.lower()
    return any(n in lower for n in needles)


def build_transcriber(
    model_id: str,
    device_str: str,
    device_int: int,
) -> tuple[Callable[..., str], Callable[[], None]]:
    errors: list[str] = []
    saw_hf_connectivity_cache_error = False

    def _record_error(label: str, exc: Exception) -> None:
        nonlocal saw_hf_connectivity_cache_error
        errors.append(f"{label}: {type(exc).__name__}: {exc}")
        if _is_hf_connectivity_cache_error(exc):
            saw_hf_connectivity_cache_error = True

    # SpeechBrain models cannot be loaded by transformers; short-circuit when we can tell.
    if speechbrain_asr.looks_like_speechbrain_repo(model_id):
        try:
            return speechbrain_asr.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("speechbrain", e)

    if _is_granite_speech_model(model_id):
        try:
            return granite_speech.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("granite_speech", e)

    # Nvidia Parakeet/Canary models are distributed as NeMo artifacts; prefer NeMo loader.
    mlow = model_id.lower()
    if _is_nemo_asr_model(model_id):
        try:
            return nemo_asr.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("nemo_asr", e)

    # Qwen ASR models use the qwen_asr runtime (third-party). Prefer qwen backend when detected.
    if mlow.startswith("qwen/") or "qwen3-asr" in mlow or "qwen3_asr" in mlow:
        try:
            return qwen_asr.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("qwen_asr", e)

    # Cohere ASR uses remote code + sentencepiece; pipeline/CTC paths fail first without this.
    if _is_cohere_asr_model(model_id):
        try:
            return universal.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("universal (cohere)", e)

    try:
        return transformers_pipeline.build_transcriber(model_id, device_int)
    except Exception as e:
        _record_error("pipeline", e)

    try:
        return universal.build_transcriber(model_id, device_str)
    except Exception as e:
        _record_error("universal", e)

    # CTC only applies to wav2vec/hubert-style checkpoints; skip when config says seq2seq.
    try:
        from .family_resolve import infer_model_type

        mt = infer_model_type(model_id)
        _ctc_skip_types = frozenset(
            {"cohere_asr", "whisper", "granite_speech", "speech_to_text", "moonshine_streaming"}
        )
        if mt not in _ctc_skip_types:
            return transformers_ctc.build_transcriber(model_id, device_str)
    except Exception as e:
        _record_error("ctc", e)

    # Final fallback in case the SpeechBrain heuristic missed the repo.
    if _looks_possibly_speechbrain_model(model_id):
        try:
            return speechbrain_asr.build_transcriber(model_id, device_str)
        except Exception as e:
            _record_error("speechbrain", e)

    if saw_hf_connectivity_cache_error:
        raise RuntimeError(
            "Could not download/load model files from Hugging Face Hub in this runtime. "
            "The environment appears offline (or the model is not cached locally). "
            f"Model: {model_id}. "
            "Retry when Hub access is available, or pre-cache the model files in this environment."
        )

    raise RuntimeError(
        f"Could not load {model_id}. Tried: " + " | ".join(errors)
    )