backends/universal.py

"""
One ASR backend for ~every speech-seq2seq model: built-in (Whisper, CohereAsr, ...) and
remote-code (efficient-speech, custom Whisper variants, etc.).

Flow:
    processor(audio, sr, return_tensors="pt", [language="en"])
    model.generate(**inputs, max_new_tokens=_safe_cap(model))
    processor.decode(outputs, ...)  (falls back to processor.batch_decode)

Loading order:
    1. `AutoModelForSpeechSeq2Seq.from_pretrained(..., trust_remote_code=True)`
       Covers standard HF classes (Whisper, CohereAsrForConditionalGeneration, ...).
    2. `AutoModel.from_pretrained(..., trust_remote_code=True)`
       Covers repos that expose their generator class via `auto_map["AutoModel"]`
       (e.g. `efficient-speech/lite-whisper-*`).

Environment overrides (all optional):
    FFASR_PROCESSOR_ID  — force a fallback processor (e.g. `openai/whisper-large-v3`).
    FFASR_LANGUAGE      — pass language=... to processor/generate when supported. Default: `en`.
    FFASR_MAX_NEW_TOKENS — override generate() length budget. Default: auto (safe cap).
"""

from __future__ import annotations

import os
from collections.abc import Callable
import json
import time

import numpy as np
import torch

from ._audio_utils import safe_pad_audio
from ._model_utils import attach_params

_DEBUG_LOG_PATH = "/home/user/app/.cursor/debug-3654e7.log"
_DEBUG_SESSION_ID = "3654e7"
_DEBUG_RUN_ID = os.environ.get("FFASR_DEBUG_RUN_ID", "initial")


def _debug_log(hypothesis_id: str, location: str, message: str, data: dict) -> None:
    payload = {
        "sessionId": _DEBUG_SESSION_ID,
        "runId": _DEBUG_RUN_ID,
        "hypothesisId": hypothesis_id,
        "location": location,
        "message": message,
        "data": data,
        "timestamp": int(time.time() * 1000),
    }
    try:
        with open(_DEBUG_LOG_PATH, "a", encoding="utf-8") as f:
            f.write(json.dumps(payload, ensure_ascii=True) + "\n")
    except Exception:
        pass


def _safe_max_new_tokens(model, env_override: int | None, default_cap: int = 256) -> int:
    """
    Whisper-style models enforce `len(decoder_prefix) + max_new_tokens ≤ max_target_positions`.
    Derive a cap from the config so we never trip that validation.
    """
    if env_override is not None and env_override > 0:
        return int(env_override)
    cfg = getattr(model, "config", None)
    if cfg is None:
        return default_cap
    mtp = (
        getattr(cfg, "max_target_positions", None)
        or getattr(cfg, "max_length", None)
        or getattr(cfg, "max_position_embeddings", None)
    )
    try:
        mtp = int(mtp) if mtp else 0
    except Exception:
        mtp = 0
    if mtp and mtp <= 2048:
        # Leave headroom for task/language/timestamp prefix tokens.
        return max(32, min(448, mtp - 16))
    return default_cap


def _pick_dtype(device_str: str) -> torch.dtype:
    use_cuda = device_str == "cuda" and torch.cuda.is_available()
    if use_cuda and torch.cuda.is_bf16_supported():
        return torch.bfloat16
    if use_cuda:
        return torch.float16
    return torch.float32


def _load_processor(model_id: str):
    from transformers import AutoProcessor

    override = os.environ.get("FFASR_PROCESSOR_ID", "").strip()
    if override:
        return AutoProcessor.from_pretrained(override, trust_remote_code=True)
    try:
        return AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    except Exception:
        return AutoProcessor.from_pretrained(
            "openai/whisper-large-v3", trust_remote_code=True
        )


def _load_model(model_id: str, dtype: torch.dtype, device_str: str):
    from transformers import AutoModel, AutoModelForSpeechSeq2Seq

    errors: list[str] = []
    for cls in (AutoModelForSpeechSeq2Seq, AutoModel):
        try:
            model = cls.from_pretrained(
                model_id,
                trust_remote_code=True,
                torch_dtype=dtype,
            ).to(device_str)
            return model, cls.__name__
        except Exception as e:
            errors.append(f"{cls.__name__}: {type(e).__name__}: {e}")
    raise RuntimeError("; ".join(errors))


def _call_processor(processor, audio: np.ndarray, sampling_rate: int, language: str):
    """Most ASR processors accept `sampling_rate` and either ignore or consume `language`."""
    tries = (
        dict(sampling_rate=sampling_rate, return_tensors="pt", language=language),
        dict(sampling_rate=sampling_rate, return_tensors="pt"),
    )
    last_exc: Exception | None = None
    for kwargs in tries:
        try:
            return processor(audio, **kwargs)
        except (TypeError, ValueError) as e:
            last_exc = e
            continue
    if last_exc is not None:
        raise last_exc
    return processor(audio, sampling_rate=sampling_rate, return_tensors="pt")


def _move_to_device(inputs, device_str: str, dtype: torch.dtype):
    if hasattr(inputs, "to"):
        inputs = inputs.to(device_str)
        try:
            inputs = inputs.to(dtype=dtype)
        except Exception:
            pass
    batch = dict(inputs)
    meta: dict = {}
    for key in ("audio_chunk_index",):
        if key in batch:
            meta[key] = batch.pop(key)
    # Some processors include optional keys with explicit None values. Passing
    # those through `generate(**batch)` can trigger model-specific failures
    # (e.g. Cohere ASR expecting a tensor decoder_attention_mask once the key exists).
    batch = {k: v for k, v in batch.items() if v is not None}

    # Cohere ASR path: when decoder_input_ids is present but decoder_attention_mask
    # is absent, HF generation internals may try to extend a None mask and crash.
    dec_ids = batch.get("decoder_input_ids")
    dec_attn = batch.get("decoder_attention_mask")
    if dec_ids is not None and dec_attn is None:
        try:
            batch["decoder_attention_mask"] = torch.ones_like(dec_ids, dtype=torch.long)
        except Exception:
            pass
    return batch, meta


def _find_expected_input_feature_dim(model) -> int | None:
    """
    Return the expected final-dimension size for `input_features` when discoverable.
    """
    candidates = (
        ("encoder", "input_linear"),
        ("model", "encoder", "input_linear"),
        ("speech_encoder", "input_linear"),
    )
    for path in candidates:
        node = model
        ok = True
        for name in path:
            if not hasattr(node, name):
                ok = False
                break
            node = getattr(node, name)
        if not ok:
            continue
        in_features = getattr(node, "in_features", None)
        try:
            if in_features is not None:
                return int(in_features)
        except Exception:
            continue
    return None


def _normalize_input_features_layout(model, batch: dict) -> dict:
    """
    Some processors return `input_features` as `[B, F, T]` while models expect `[B, T, F]`.
    If we can infer the expected feature dim and it matches axis 1 (not axis 2),
    transpose to avoid shape mismatches in the encoder projection.
    """
    feats = batch.get("input_features")
    if not hasattr(feats, "shape") or getattr(feats, "dim", lambda: 0)() != 3:
        return batch

    expected = _find_expected_input_feature_dim(model)
    if expected is None:
        return batch

    b, d1, d2 = int(feats.shape[0]), int(feats.shape[1]), int(feats.shape[2])
    if d2 == expected:
        return batch
    if d1 == expected:
        fixed = dict(batch)
        fixed["input_features"] = feats.transpose(1, 2).contiguous()
        _debug_log(
            "H6",
            "backends/universal.py:_normalize_input_features_layout:transpose",
            "Transposed input_features to match model expected feature dimension",
            {
                "batch_size": b,
                "before_shape": [b, d1, d2],
                "after_shape": [b, d2, d1],
                "expected_feature_dim": expected,
            },
        )
        return fixed
    return batch


def _encoder_tensor_from_batch(batch: dict):
    """Return encoder inputs without using ``or`` on tensors (ambiguous bool)."""
    if batch.get("input_features") is not None:
        return batch["input_features"]
    if batch.get("input_values") is not None:
        return batch["input_values"]
    return None


def _generate(model, batch: dict, max_new_tokens: int, language: str):
    """Try language-aware generate first (Whisper path); fall back to plain generate."""
    batch = _normalize_input_features_layout(model, batch)
    base = dict(max_new_tokens=max_new_tokens, num_beams=1)
    model_type = getattr(getattr(model, "config", None), "model_type", "") or ""

    if model_type == "cohere_asr" and "decoder_attention_mask" not in batch:
        # Cohere's remote generate() path can leave decoder_attention_mask as None,
        # then HF generation tries `decoder_attention_mask.new_ones(...)` and crashes.
        # Passing an explicit one-token decoder mask keeps generation state valid.
        src = _encoder_tensor_from_batch(batch)
        if src is not None and hasattr(src, "shape"):
            try:
                batch = dict(batch)
                batch["decoder_attention_mask"] = torch.ones(
                    (int(src.shape[0]), 1),
                    device=src.device,
                    dtype=torch.long,
                )
            except Exception:
                pass

    attempts = (
        {**base, "task": "transcribe", "language": language},
        base,
    )
    # #region agent log
    _debug_log(
        "H4",
        "backends/universal.py:_generate:attempts",
        "Generate attempts and core tensor shapes",
        {
            "model_type": model_type,
            "attempts": [sorted(extra.keys()) for extra in attempts],
            "batch_keys": sorted(batch.keys()),
            "input_features_shape": list(batch.get("input_features").shape) if hasattr(batch.get("input_features"), "shape") else None,
            "input_values_shape": list(batch.get("input_values").shape) if hasattr(batch.get("input_values"), "shape") else None,
        },
    )
    # #endregion
    last: Exception | None = None
    for extra in attempts:
        try:
            return model.generate(**batch, **extra)
        except TypeError as e:
            last = e
        except Exception as e:
            # Whisper sometimes raises when a kwarg (language) is unsupported for this model.
            last = e
            if "language" in extra:
                continue
            raise
    if last is not None:
        raise last
    return model.generate(**batch, **base)


def _decode(processor, outputs, meta: dict, language: str) -> str:
    """Prefer `processor.decode` (Cohere / long-form chunking); else `batch_decode`."""
    if hasattr(processor, "decode"):
        try:
            if meta.get("audio_chunk_index") is not None:
                return str(
                    processor.decode(
                        outputs,
                        skip_special_tokens=True,
                        audio_chunk_index=meta["audio_chunk_index"],
                        language=language,
                    )
                ).strip()
            result = processor.decode(outputs, skip_special_tokens=True)
            if isinstance(result, str):
                return result.strip()
        except TypeError:
            pass
        except Exception:
            pass
    try:
        decoded = processor.batch_decode(outputs, skip_special_tokens=True)
        return str(decoded[0]).strip() if decoded else ""
    except Exception:
        if hasattr(processor, "tokenizer"):
            decoded = processor.tokenizer.batch_decode(outputs, skip_special_tokens=True)
            return str(decoded[0]).strip() if decoded else ""
        raise


def build_transcriber(model_id: str, device_str: str) -> tuple[Callable[..., str], Callable[[], None]]:
    processor = _load_processor(model_id)
    dtype = _pick_dtype(device_str)
    model, _cls = _load_model(model_id, dtype, device_str)
    model.eval()
    # #region agent log
    _debug_log(
        "H1",
        "backends/universal.py:build_transcriber:model_loaded",
        "Universal backend selected and model loaded",
        {
            "model_id": model_id,
            "model_type": str(getattr(getattr(model, "config", None), "model_type", "")),
            "processor_class": processor.__class__.__name__,
            "override_processor_id": bool(os.environ.get("FFASR_PROCESSOR_ID", "").strip()),
        },
    )
    # #endregion

    language = os.environ.get("FFASR_LANGUAGE", "en").strip() or "en"
    env_cap_raw = os.environ.get("FFASR_MAX_NEW_TOKENS", "").strip()
    env_cap = int(env_cap_raw) if env_cap_raw.isdigit() else None
    max_new = _safe_max_new_tokens(model, env_cap)

    def transcribe(audio_np: np.ndarray, sampling_rate: int = 16000) -> str:
        arr = safe_pad_audio(audio_np)
        # #region agent log
        _debug_log(
            "H3",
            "backends/universal.py:transcribe:audio_in",
            "Audio input after safe_pad_audio",
            {
                "sampling_rate": int(sampling_rate),
                "arr_shape": list(arr.shape) if hasattr(arr, "shape") else None,
                "arr_dtype": str(arr.dtype) if hasattr(arr, "dtype") else type(arr).__name__,
            },
        )
        # #endregion
        inputs = _call_processor(processor, arr, int(sampling_rate), language)
        # #region agent log
        _debug_log(
            "H2",
            "backends/universal.py:transcribe:processor_out",
            "Processor output keys and shapes",
            {
                "keys": sorted(list(dict(inputs).keys())) if hasattr(inputs, "keys") else [],
                "input_features_shape": list(dict(inputs).get("input_features").shape) if hasattr(dict(inputs).get("input_features"), "shape") else None,
                "input_values_shape": list(dict(inputs).get("input_values").shape) if hasattr(dict(inputs).get("input_values"), "shape") else None,
            },
        )
        # #endregion
        batch, meta = _move_to_device(inputs, device_str, dtype)
        # #region agent log
        _debug_log(
            "H5",
            "backends/universal.py:transcribe:batch_out",
            "Batch sent to generate after device move",
            {
                "batch_keys": sorted(batch.keys()),
                "meta_keys": sorted(meta.keys()),
                "input_features_shape": list(batch.get("input_features").shape) if hasattr(batch.get("input_features"), "shape") else None,
                "input_values_shape": list(batch.get("input_values").shape) if hasattr(batch.get("input_values"), "shape") else None,
            },
        )
        # #endregion
        with torch.no_grad():
            outputs = _generate(model, batch, max_new, language)
        return _decode(processor, outputs, meta, language)

    attach_params(transcribe, model)

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

    return transcribe, cleanup