shared_model.py

"""
Shared ASR model singleton, VAD factory, and common configuration.
Imported by app.py, app_chughtai.py, and main.py so the GPU is
allocated once and both servers share the same vLLM engine.
"""

import os
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

import copy
import logging
import re
import threading
import time
import concurrent.futures
import numpy as np

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("qwen3-asr")

# ── Shared configuration ──────────────────────────────────────────────────────
MODEL_ID                = os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-1.7B")
GPU_MEMORY_UTILIZATION  = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.90"))
STREAMING_MAX_NEW_TOKENS = int(os.getenv("STREAMING_MAX_NEW_TOKENS", "1024"))
MAX_MODEL_LEN           = int(os.getenv("MAX_MODEL_LEN", "0")) or None
ENFORCE_EAGER           = os.getenv("ENFORCE_EAGER", "false").lower() in ("1", "true", "yes")
PORT                    = int(os.getenv("PORT", "7860"))

CHUNK_SIZE_SEC   = float(os.getenv("CHUNK_SIZE_SEC", "4.0"))
UNFIXED_CHUNK_NUM = int(os.getenv("UNFIXED_CHUNK_NUM", "5"))
UNFIXED_TOKEN_NUM = int(os.getenv("UNFIXED_TOKEN_NUM", "15"))
SAMPLE_RATE      = 16000

VAD_THRESHOLD      = float(os.getenv("VAD_THRESHOLD", "0.7"))
VAD_MIN_SILENCE_MS = int(os.getenv("VAD_MIN_SILENCE_MS", "800"))
VAD_SPEECH_PAD_MS  = int(os.getenv("VAD_SPEECH_PAD_MS", "300"))

HALLUCINATION_PHRASES = {
    "transcript", "transcription", "thank you", "thanks for watching",
    "you", "bye", "goodbye", "the end", "subtitle", "subtitles",
}

# Thread pool shared by both sub-apps for running inference off the event loop
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)


def _is_hallucination(text: str) -> bool:
    return text.lower().strip().rstrip(".!?,;:") in HALLUCINATION_PHRASES


def _normalize_ws(s: str) -> str:
    return re.sub(r"\s+", " ", s or "").strip().lower()


def _is_context_echo(text: str, context: str) -> bool:
    """
    True when the model regurgitated the context prompt instead of
    transcribing — a common failure on very short or near-silent
    utterances. Matches when the normalized output is a leading slice of
    the context (partial echo) or begins with the whole context (full echo).
    The 15-char floor avoids false positives on short real speech.
    """
    if not text or not context:
        return False
    nt, nc = _normalize_ws(text), _normalize_ws(context)
    if len(nt) < 15:
        return False
    return nc.startswith(nt) or nt.startswith(nc)


# ── Language resolution ───────────────────────────────────────────────────────
# Mirrors qwen_asr.inference.utils.SUPPORTED_LANGUAGES (v0.0.6). A language the
# model can't handle makes init_streaming_state() raise, which would otherwise
# kill the WebSocket connection mid-stream ("connected but no audio processed").
SUPPORTED_LANGUAGES = {
    "Chinese", "English", "Cantonese", "Arabic", "German", "French", "Spanish",
    "Portuguese", "Indonesian", "Italian", "Korean", "Russian", "Thai",
    "Vietnamese", "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish",
    "Danish", "Finnish", "Polish", "Czech", "Filipino", "Persian", "Greek",
    "Romanian", "Hungarian", "Macedonian",
}

# Client-supplied language (ISO code or full name, any case) -> canonical model
# language. NOTE: Qwen3-ASR has NO "Urdu" — Urdu requests are routed to Hindi so
# the model emits Devanagari, which the Roman-Urdu transliterator then converts.
_LANG_ALIASES = {
    "en": "English",    "english": "English",
    "ms": "Malay",      "malay": "Malay",
    "zh": "Chinese",    "chinese": "Chinese",
    "ja": "Japanese",   "japanese": "Japanese",
    "ko": "Korean",     "korean": "Korean",
    "hi": "Hindi",      "hindi": "Hindi",
    "ur": "Hindi",      "urdu": "Hindi",
    "ar": "Arabic",     "arabic": "Arabic",
    "id": "Indonesian", "indonesian": "Indonesian",
    "th": "Thai",       "thai": "Thai",
    "fa": "Persian",    "persian": "Persian",
}


def resolve_language(raw):
    """
    Map a client-supplied language to a canonical Qwen3-ASR language.
    Accepts ISO codes ("hi"), full names ("hindi"), any case. Returns the
    canonical name if supported, else None (caller keeps its own default
    rather than forcing an unsupported value that would crash the stream).
    """
    if not raw:
        return None
    key = str(raw).strip().lower()
    name = _LANG_ALIASES.get(key)
    if name is None:                       # try the raw value as a canonical name
        cand = key[:1].upper() + key[1:]
        name = cand if cand in SUPPORTED_LANGUAGES else None
    return name if name in SUPPORTED_LANGUAGES else None


# ── ASR model singleton ───────────────────────────────────────────────────────

_asr_model  = None
_asr_lock   = threading.Lock()
_model_ready = threading.Event()

# Serializes ALL inference on the shared vLLM engine. The offline vllm.LLM is
# not safe for concurrent generate() calls, and the batch endpoint temporarily
# patches generate / mutates the shared sampling_params — so streaming and batch
# must be mutually exclusive. Hold it only around a single inference call.
model_lock = threading.Lock()


def _get_dtype():
    try:
        import torch
        if torch.cuda.is_available():
            cap = torch.cuda.get_device_capability()
            if cap[0] * 10 + cap[1] >= 80:
                return "bfloat16"
    except Exception:
        pass
    return "half"


def get_asr_model():
    global _asr_model
    if _asr_model is None:
        with _asr_lock:
            if _asr_model is None:
                log.info(f"Loading {MODEL_ID}...")
                start = time.time()
                dtype = _get_dtype()
                from qwen_asr import Qwen3ASRModel
                llm_kwargs = {
                    "model": MODEL_ID,
                    "gpu_memory_utilization": GPU_MEMORY_UTILIZATION,
                    "dtype": dtype,
                    "max_new_tokens": STREAMING_MAX_NEW_TOKENS,
                    "enforce_eager": ENFORCE_EAGER,
                }
                if MAX_MODEL_LEN:
                    llm_kwargs["max_model_len"] = MAX_MODEL_LEN
                _asr_model = Qwen3ASRModel.LLM(**llm_kwargs)
                log.info(f"Model loaded in {time.time() - start:.1f}s (dtype={dtype})")
                _model_ready.set()
    return _asr_model


def is_model_ready() -> bool:
    return _model_ready.is_set()


# ── Silero VAD (one template; deep-copied per WebSocket connection) ────────────

_vad_model_template = None
_vad_utils          = None
_vad_lock           = threading.Lock()


def _ensure_vad_loaded():
    global _vad_model_template, _vad_utils
    if _vad_model_template is None:
        with _vad_lock:
            if _vad_model_template is None:
                import torch
                _vad_model_template, _vad_utils = torch.hub.load(
                    "snakers4/silero-vad", "silero_vad",
                    trust_repo=True, verbose=False,
                )
                log.info("Silero VAD model loaded")


def create_vad(threshold=VAD_THRESHOLD, min_silence_ms=VAD_MIN_SILENCE_MS,
               speech_pad_ms=VAD_SPEECH_PAD_MS):
    try:
        _ensure_vad_loaded()
        model_copy   = copy.deepcopy(_vad_model_template)
        VADIterator  = _vad_utils[3]
        return VADIterator(
            model_copy,
            threshold=threshold,
            sampling_rate=SAMPLE_RATE,
            min_silence_duration_ms=min_silence_ms,
            speech_pad_ms=speech_pad_ms,
        )
    except Exception as e:
        log.warning(f"Silero VAD failed ({e}), using RMS fallback")
        return RMSVad(threshold=0.01, silence_frames=int(min_silence_ms / 20))


def is_vad_ready() -> bool:
    return _vad_model_template is not None


class RMSVad:
    """Energy-based VAD fallback used when Silero fails to load."""

    def __init__(self, threshold=0.01, silence_frames=30, speech_frames=3):
        self.threshold     = threshold
        self.silence_frames = silence_frames
        self.speech_frames  = speech_frames
        self.is_speaking   = False
        self._silent_count = 0
        self._speech_count = 0

    def __call__(self, audio_chunk):
        import torch
        if isinstance(audio_chunk, np.ndarray):
            audio_chunk = torch.from_numpy(audio_chunk)
        rms = float(torch.sqrt(torch.mean(audio_chunk.float() ** 2)))

        if rms > self.threshold:
            self._speech_count += 1
            self._silent_count  = 0
            if not self.is_speaking and self._speech_count >= self.speech_frames:
                self.is_speaking = True
                return {"start": 0}
        else:
            self._silent_count += 1
            self._speech_count  = 0
            if self.is_speaking and self._silent_count >= self.silence_frames:
                self.is_speaking = False
                return {"end": 0}
        return None

    def reset_states(self):
        self.is_speaking   = False
        self._silent_count = 0
        self._speech_count = 0