backend/asr.py

"""Speech-to-text utilities with graceful fallbacks and dynamic stage switching."""

from __future__ import annotations

import io
import os
import wave
from threading import Lock
from typing import Any, Literal, Optional

import numpy as np

try:
    from openai import OpenAI
except ModuleNotFoundError:
    OpenAI = None  # type: ignore[assignment]

from backend.utils import device
import nemo.collections.asr as nemo_asr

try:
    import torch
    from transformers import pipeline
except ModuleNotFoundError:  # PyTorch or transformers not available on Python 3.13 wheels
    torch = None  # type: ignore
    pipeline = None  # type: ignore

try:
    from google.cloud import speech
except ModuleNotFoundError:
    speech = None  # type: ignore


_ASR_PIPELINE = None
_ASR_STAGE: Literal["typhoon", "gpt"] = "typhoon"
_ASR_STAGE_LOCK: Lock = Lock()


def _huggingface_device() -> int | str | None:
    if device == "cuda":
        return 0
    if device == "mps":
        return "mps"
    return "cpu"


def _initialize_typhoon_pipeline():
    if torch is None or pipeline is None:
        return None

    print(f"Using device: {device}")
    print("Initializing Typhoon ASR pipeline...")
    asr_model = nemo_asr.models.ASRModel.from_pretrained(
        model_name="scb10x/typhoon-asr-realtime",
        map_location=device,
    )
    print("Typhoon ASR pipeline initialized.")
    return asr_model


def _initialize_gpt_client() -> Optional[Any]:
    if OpenAI is None:
        print("openai package not available; GPT ASR unavailable.")
        return None

    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("OPENAI_API_KEY not found; GPT ASR unavailable.")
        return None
    try:
        return OpenAI(api_key=api_key)
    except Exception as exc:
        print(f"Failed to initialise GPT ASR client: {exc}")
        return None


_GPT_ASR_MODEL = os.getenv("GPT_ASR_MODEL", "gpt-4o-mini-transcribe")
_GPT_CLIENT = _initialize_gpt_client()


def set_asr_stage(stage: str) -> None:
    """Update the active ASR stage."""
    normalized_stage = stage.lower()
    if normalized_stage not in {"typhoon", "gpt"}:
        raise ValueError(f"Unsupported ASR stage '{stage}'")
    global _ASR_STAGE
    with _ASR_STAGE_LOCK:
        if _ASR_STAGE != normalized_stage:
            print(f"Switching ASR stage to: {normalized_stage}")
        _ASR_STAGE = normalized_stage  # type: ignore[assignment]


def get_asr_stage() -> Literal["typhoon", "gpt"]:
    """Return the current ASR stage."""
    with _ASR_STAGE_LOCK:
        return _ASR_STAGE


def _transcribe_with_pipeline(audio_array: np.ndarray) -> str:
    output = _ASR_PIPELINE(audio_array)  # type: ignore[operator]
    if isinstance(output, dict):
        text = output.get("text", "")
    else:
        text = str(output)
    return text.replace("ทางลัด", "ทางรัฐ")


def _transcribe_with_typhoon(audio_array: np.ndarray) -> str:
    if _ASR_TYPHOON is None:
        raise RuntimeError("Typhoon ASR is unavailable")

    result = _ASR_TYPHOON.transcribe(audio=audio_array)
    if isinstance(result, list):
        transcription = " ".join(result)
    else:
        transcription = str(result)
    return transcription.strip()


def _transcribe_with_google(audio_array: np.ndarray) -> str:
    if speech is None:
        raise RuntimeError("google-cloud-speech is not available")

    int16_audio = (audio_array * 32767.0).astype(np.int16)
    audio_bytes = int16_audio.tobytes()

    client = speech.SpeechClient()
    audio_config = speech.RecognitionConfig(
        encoding=speech.RecognitionConfig.AudioEncoding.LINEAR16,
        sample_rate_hertz=16000,
        language_code="th-TH",
        alternative_language_codes=["en-US"],
        model="telephony",
    )
    audio_data = speech.RecognitionAudio(content=audio_bytes)
    response = client.recognize(config=audio_config, audio=audio_data)
    transcription = " ".join(
        result.alternatives[0].transcript for result in response.results
    )
    return transcription


def _transcribe_with_gpt(audio_array: np.ndarray) -> str:
    if _GPT_CLIENT is None:
        raise RuntimeError("GPT ASR client is unavailable")

    normalized = np.asarray(audio_array, dtype=np.float32)
    if normalized.ndim > 1:
        normalized = normalized.squeeze()
    normalized = np.clip(normalized, -1.0, 1.0)
    int16_audio = (normalized * 32767.0).astype(np.int16)

    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(16000)
        wav_file.writeframes(int16_audio.tobytes())
    buffer.seek(0)

    response = _GPT_CLIENT.audio.transcriptions.create(
        model=_GPT_ASR_MODEL,
        file=("audio.wav", buffer.read(), "audio/wav"),
    )
    text = getattr(response, "text", "")
    if not text and isinstance(response, dict):
        text = response.get("text", "")
    return text.strip()


_ASR_TYPHOON = _initialize_typhoon_pipeline()

def transcribe_typhoon(path: str) -> str:
    text = _ASR_TYPHOON.transcribe(path)
    if text[0].text:
        return text[0].text
    else :
        print(text)
        return ""
    


def transcribe_audio(audio_array: np.ndarray) -> str:
    """Transcribe user audio with the best available backend based on the current stage."""
    if audio_array is None or not np.any(audio_array):
        return ""

    stage = get_asr_stage()

    if stage == "gpt":
        try:
            transcription = _transcribe_with_gpt(audio_array)
            if transcription:
                return transcription.replace("ทางลัด", "ทางรัฐ")
        except Exception as exc:
            print(f"GPT ASR failed: {exc}; falling back to Typhoon.")

    if _ASR_TYPHOON is not None:
        try:
            transcription = _transcribe_with_typhoon(audio_array)
            if transcription:
                return transcription.replace("ทางลัด", "ทางรัฐ")
        except Exception as exc:
            print(f"Typhoon ASR pipeline failed: {exc}")

    try:
        return _transcribe_with_google(audio_array).replace("ทางลัด", "ทางรัฐ")
    except Exception as exc:
        print(f"ASR fallback failed: {exc}")
        return ""