asr.py

import numpy as np
import noisereduce as nr
import torch
import io
import time
import wave
import threading
import logging

from groq import Groq
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
from transformers import WhisperFeatureExtractor, WhisperTokenizer, WhisperProcessor

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SAMPLE_RATE          = 16000
GROQ_TIMEOUT         = 3.0
VAD_THRESHOLD        = 0.3
BUFFER_DURATION      = 1.5
BUFFER_SAMPLES       = int(SAMPLE_RATE * BUFFER_DURATION)

LOCAL_MODEL_DIR = "/app/models/whisper"

GROQ_API_KEYS = [
    "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",
]
_groq_key_idx  = 0
_groq_key_lock = threading.Lock()

def _get_groq_client():
    global _groq_key_idx
    with _groq_key_lock:
        key = GROQ_API_KEYS[_groq_key_idx % len(GROQ_API_KEYS)]
        _groq_key_idx += 1
    return Groq(api_key=key)

BLACKLIST = [
    "transkripsi percakapan", "bahasa indonesia", "transkripsi",
    "indonesia.", "indonesia", "terima kasih", "thank you",
    "thanks", "subscribe", "like", "comment",
]

EXACT_BLACKLIST = [
    "indonesia", "terima kasih", "thank you", "thanks",
    "like", "subscribe", "comment", "ok", "oke",
]

ENGLISH_STARTERS = [
    "so ", "we ", "the ", "this ", "that ", "i ", "you ", "it ",
    "and ", "but ", "for ", "are ", "is ", "was ", "in ", "of ",
    "a ", "an ", "to ", "with ", "my ", "your ", "our ", "they ",
    "he ", "she ", "what ", "how ", "when ", "where ", "why ",
]


def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int = SAMPLE_RATE) -> bytes:
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        pcm = (audio * 32767).astype(np.int16)
        wf.writeframes(pcm.tobytes())
    return buf.getvalue()


class ASRPipeline:
    def __init__(self):
        logger.info("Loading processor dari openai/whisper-small...")
        fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
        tk = WhisperTokenizer.from_pretrained(
            "openai/whisper-small",
            language="indonesian",
            task="transcribe"
        )
        self.processor = WhisperProcessor(feature_extractor=fe, tokenizer=tk)
        logger.info("Processor loaded ✓")

        logger.info(f"Loading ONNX dari {LOCAL_MODEL_DIR} ...")
        self.onnx_model = ORTModelForSpeechSeq2Seq.from_pretrained(
            LOCAL_MODEL_DIR,
            export=False,
        )
        self.forced_decoder_ids = self.processor.get_decoder_prompt_ids(
            language="indonesian", task="transcribe"
        )
        logger.info("Whisper ONNX loaded ✓")

        self.vad_model, self.vad_utils = torch.hub.load(
            repo_or_dir="snakers4/silero-vad",
            model="silero_vad",
            force_reload=False,
            trust_repo=True,
        )
        (self.get_speech_timestamps, _, self.read_audio, *_) = self.vad_utils
        logger.info("Silero VAD loaded ✓")

        self.rolling_buffer    = np.array([], dtype=np.float32)
        self.last_text         = ""
        self.is_active         = False
        self.lock              = threading.Lock()
        # Semaphore — hanya 1 inference paralel sekaligus, cegah backlog
        self.inference_sem     = threading.Semaphore(1)

        logger.info("ASRPipeline (Whisper ONNX + Groq parallel) siap ✓")

    def stop(self):
        self.is_active = False
        self.rolling_buffer = np.array([], dtype=np.float32)
        self.last_text = ""
        logger.info("Pipeline stopped ✓")

    def clear_buffer(self):
        with self.lock:
            self.rolling_buffer = np.array([], dtype=np.float32)

    def preprocess(self, audio: np.ndarray) -> np.ndarray:
        if len(audio) > SAMPLE_RATE * 0.1:
            try:
                audio = nr.reduce_noise(y=audio, sr=SAMPLE_RATE, stationary=True)
            except Exception:
                pass
        return audio.astype(np.float32)

    def get_vad_confidence(self, audio: np.ndarray) -> float:
        VAD_CHUNK   = 512
        confidences = []
        for i in range(0, len(audio) - VAD_CHUNK + 1, VAD_CHUNK):
            chunk  = audio[i : i + VAD_CHUNK]
            tensor = torch.from_numpy(chunk)
            conf   = self.vad_model(tensor, SAMPLE_RATE).item()
            confidences.append(conf)
        return max(confidences) if confidences else 0.0

    def has_enough_energy(self, audio: np.ndarray) -> bool:
        return np.sqrt(np.mean(audio ** 2)) > 0.01

    def is_valid_text(self, text: str) -> bool:
        t = text.lower().strip()
        if len(t) < 3:
            return False
        if t in EXACT_BLACKLIST:
            return False
        for bl in BLACKLIST:
            if t.startswith(bl):
                return False
        for s in ENGLISH_STARTERS:
            if t.startswith(s):
                return False
        words = t.replace(".", "").replace(",", "").replace("!", "").split()
        if len(words) >= 3:
            mc = max(set(words), key=words.count)
            if words.count(mc) / len(words) > 0.4:
                return False
        return True

    def whisper_transcribe(self, audio: np.ndarray) -> str:
        if not self.is_active:
            return ""
        try:
            audio  = self.preprocess(audio)
            inputs = self.processor(
                audio,
                sampling_rate=SAMPLE_RATE,
                return_tensors="pt",
            )
            if not self.is_active:
                return ""
            generated = self.onnx_model.generate(
                **inputs,
                forced_decoder_ids=self.forced_decoder_ids,
            )
            text = self.processor.batch_decode(
                generated, skip_special_tokens=True
            )[0].strip()
            logger.info(f"[ONNX] '{text}'")
            return text
        except Exception as e:
            logger.error(f"Whisper ONNX error: {e}")
            return ""

    def groq_transcribe(self, audio: np.ndarray) -> str:
        if not self.is_active:
            return ""
        try:
            audio     = self.preprocess(audio)
            wav_bytes = audio_to_wav_bytes(audio)
            if not self.is_active:
                return ""
            start  = time.time()
            client = _get_groq_client()
            transcription = client.audio.transcriptions.create(
                file=("audio.wav", wav_bytes),
                model="whisper-large-v3-turbo",
                language="id",
                response_format="text",
                prompt="",
                timeout=GROQ_TIMEOUT,
            )
            text = (
                transcription.strip()
                if isinstance(transcription, str)
                else str(transcription).strip()
            )
            elapsed = time.time() - start
            logger.info(f"[Groq {elapsed:.2f}s] '{text}'")
            return text
        except Exception as e:
            logger.error(f"Groq error: {e}")
            return ""

    def parallel_transcribe(self, audio: np.ndarray) -> tuple[str, str]:
        winner      = [None, None]
        result_lock = threading.Lock()
        event       = threading.Event()

        def run_onnx():
            t = self.whisper_transcribe(audio)
            with result_lock:
                if winner[0] is None and t:
                    winner[0] = t
                    winner[1] = "onnx"
                    event.set()

        def run_groq():
            t = self.groq_transcribe(audio)
            with result_lock:
                if winner[0] is None and t:
                    winner[0] = t
                    winner[1] = "groq"
                    event.set()

        threading.Thread(target=run_onnx, daemon=True).start()
        threading.Thread(target=run_groq, daemon=True).start()

        # Tunggu sampai 5s — cukup untuk Groq (3s) + buffer
        event.wait(timeout=5.0)
        return winner[0] or "", winner[1] or "none"

    def transcribe_chunk(self, audio_chunk: np.ndarray) -> str | None:
        # Empty = pause signal
        if len(audio_chunk) == 0:
            self.clear_buffer()
            return None

        if not self.is_active:
            self.is_active = True

        with self.lock:
            if not self.is_active:
                return None

            vad_conf = self.get_vad_confidence(audio_chunk)
            if vad_conf <= VAD_THRESHOLD:
                self.rolling_buffer = np.array([], dtype=np.float32)
                return None

            if not self.has_enough_energy(audio_chunk):
                return None

            self.rolling_buffer = np.concatenate([self.rolling_buffer, audio_chunk])
            if len(self.rolling_buffer) > BUFFER_SAMPLES:
                self.rolling_buffer = self.rolling_buffer[-BUFFER_SAMPLES:]

            if len(self.rolling_buffer) < BUFFER_SAMPLES:
                return None

            audio_to_process    = self.rolling_buffer.copy()
            self.rolling_buffer = self.rolling_buffer[BUFFER_SAMPLES // 2:]

        if not self.is_active:
            return None

        # Semaphore — skip kalau ada inference yang sedang jalan
        if not self.inference_sem.acquire(blocking=False):
            logger.debug("Inference busy, skip chunk")
            return None

        try:
            text, source = self.parallel_transcribe(audio_to_process)
        finally:
            self.inference_sem.release()

        logger.info(f"[{source}] final: '{text}'")

        if not text:
            return None
        if not self.is_valid_text(text):
            logger.info(f"Filtered: {text}")
            return None

        self.last_text = text
        return text.strip()