shared_model.py

"""
Shared helpers for the async servers: Silero VAD factory (one template,
deep-copied per WebSocket connection), text post-filters (hallucination /
context-echo), and client-language resolution.

The vLLM engine itself lives in async_streaming.py (a single shared
vllm.AsyncLLM). This module deliberately holds NO model/engine state.
"""

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

import copy
import logging
import re
import threading
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")

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",
}


# ── Text post-filters ─────────────────────────────────────────────────────────
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
    # Full echo, or output is the leading slice of the context (or vice-versa).
    if nc.startswith(nt) or nt.startswith(nc):
        return True
    # Mid-context echo: a long verbatim run of the context anywhere in the
    # output. High floor (40 chars) so short real overlaps with the term lists
    # (e.g. "fever, cough, headache") don't false-trigger.
    if len(nt) >= 40 and (nt in nc or nc in nt):
        return True
    return False


# ── Language resolution ───────────────────────────────────────────────────────
# Mirrors qwen_asr.inference.utils.SUPPORTED_LANGUAGES (v0.0.6). A language the
# model can't handle makes the engine 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


# ── 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