feat: parallel ONNX+Groq inference

asr.py

@@ -6,6 +6,7 @@ import time
 import wave
 import threading
 import logging
+import concurrent.futures
 
 from groq import Groq
 from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
@@ -15,11 +16,10 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 SAMPLE_RATE          = 16000
-GROQ_TIMEOUT         = 3.0   # dari 5.0 — skip Groq kalau lambat
+GROQ_TIMEOUT         = 3.0
 VAD_THRESHOLD        = 0.3
-BUFFER_DURATION      = 2.0   # dari 2.5 — buffer lebih pendek
+BUFFER_DURATION      = 2.0
 BUFFER_SAMPLES       = int(SAMPLE_RATE * BUFFER_DURATION)
-CONFIDENCE_THRESHOLD = -0.5
 
 LOCAL_MODEL_DIR = "/app/models/whisper"
 
@@ -105,9 +105,14 @@ class ASRPipeline:
         self.is_active      = True
         self.lock           = threading.Lock()
 
-        logger.info("ASRPipeline (Whisper ONNX + Groq hybrid) siap ✓")
+        # Thread pool untuk parallel inference
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
+
+        logger.info("ASRPipeline (Whisper ONNX + Groq parallel) siap ✓")
 
     def stop(self):
+        self.is_active = False
+        time.sleep(0.1)
         self.rolling_buffer = np.array([], dtype=np.float32)
         self.last_text      = ""
         self.is_active      = True
@@ -157,7 +162,8 @@ class ASRPipeline:
                 return False
         return True
 
-    def whisper_transcribe(self, audio: np.ndarray) -> tuple[str, float]:
+    def whisper_transcribe(self, audio: np.ndarray) -> str:
+        """ONNX inference — returns text atau empty string."""
         try:
             audio  = self.preprocess(audio)
             inputs = self.processor(
@@ -168,73 +174,90 @@ class ASRPipeline:
             generated = self.onnx_model.generate(
                 **inputs,
                 forced_decoder_ids=self.forced_decoder_ids,
-                return_dict_in_generate=True,
-                output_scores=True,
             )
             text = self.processor.batch_decode(
-                generated.sequences, skip_special_tokens=True
+                generated, skip_special_tokens=True
             )[0].strip()
-
-            if generated.scores:
-                log_probs = []
-                for score, token_id in zip(
-                    generated.scores,
-                    generated.sequences[0][len(self.forced_decoder_ids) + 1:],
-                ):
-                    lp = torch.nn.functional.log_softmax(score[0], dim=-1)
-                    log_probs.append(lp[token_id].item())
-                avg_logprob = float(np.mean(log_probs)) if log_probs else -1.0
-            else:
-                avg_logprob = -1.0
-
-            logger.info(f"[ONNX] '{text}' | logprob={avg_logprob:.3f}")
-            return text, avg_logprob
-
+            logger.info(f"[ONNX] '{text}'")
+            return text
         except Exception as e:
             logger.error(f"Whisper ONNX error: {e}")
-            return "", -99.0
+            return ""
 
     def groq_transcribe(self, audio: np.ndarray) -> str:
+        """Groq inference — returns text atau empty string."""
         try:
             audio     = self.preprocess(audio)
             wav_bytes = audio_to_wav_bytes(audio)
-            result    = [""]
-            success   = [False]
-            start     = time.time()
-
-            def call_groq():
-                try:
-                    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="",
-                    )
-                    result[0] = (
-                        transcription.strip()
-                        if isinstance(transcription, str)
-                        else str(transcription).strip()
-                    )
-                    success[0] = True
-                except Exception as e:
-                    logger.error(f"Groq error: {e}")
-
-            t = threading.Thread(target=call_groq)
-            t.start()
-            t.join(timeout=GROQ_TIMEOUT)
-
-            elapsed = time.time() - start
-            if success[0] and result[0]:
-                logger.info(f"[Groq {elapsed:.2f}s] >> {result[0]}")
-                return result[0]
-            logger.warning(f"[Groq {elapsed:.2f}s] timeout/empty")
-            return ""
+            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="",
+            )
+            text = (
+                transcription.strip()
+                if isinstance(transcription, str)
+                else str(transcription).strip()
+            )
+            logger.info(f"[Groq] '{text}'")
+            return text
         except Exception as e:
-            logger.error(f"Groq exception: {e}")
+            logger.error(f"Groq error: {e}")
             return ""
 
+    def parallel_transcribe(self, audio: np.ndarray) -> tuple[str, str]:
+        """
+        Jalankan ONNX dan Groq secara paralel.
+        Return (winner_text, winner_source) — siapapun yang selesai duluan.
+        """
+        result_text   = ["", ""]
+        result_done   = [False, False]
+        result_lock   = threading.Lock()
+        winner        = [None, None]  # [text, source]
+        event         = threading.Event()
+
+        def run_onnx():
+            t = self.whisper_transcribe(audio)
+            with result_lock:
+                result_text[0]  = t
+                result_done[0]  = True
+                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:
+                result_text[1]  = t
+                result_done[1]  = True
+                if winner[0] is None and t:
+                    winner[0] = t
+                    winner[1] = "groq"
+                    event.set()
+
+        t1 = threading.Thread(target=run_onnx, daemon=True)
+        t2 = threading.Thread(target=run_groq, daemon=True)
+        t1.start()
+        t2.start()
+
+        # Tunggu siapapun yang selesai duluan, max GROQ_TIMEOUT + 1s
+        event.wait(timeout=GROQ_TIMEOUT + 1.0)
+
+        # Kalau event tidak di-set (keduanya gagal/timeout)
+        if winner[0] is None:
+            # Tunggu sebentar lagi untuk ONNX
+            t1.join(timeout=2.0)
+            with result_lock:
+                if result_text[0]:
+                    winner[0] = result_text[0]
+                    winner[1] = "onnx_late"
+
+        return winner[0] or "", winner[1] or "none"
+
     def transcribe_chunk(self, audio_chunk: np.ndarray) -> str | None:
         # Empty array = signal pause dari extension
         if len(audio_chunk) == 0:
@@ -261,25 +284,18 @@ class ASRPipeline:
                 return None
 
             audio_to_process    = self.rolling_buffer.copy()
-            # Overlap 33% — lebih cepat dari 50%
             self.rolling_buffer = self.rolling_buffer[BUFFER_SAMPLES // 3:]
 
-            # Hybrid logic
-            text, avg_logprob = self.whisper_transcribe(audio_to_process)
-
-            if not text or avg_logprob < CONFIDENCE_THRESHOLD:
-                source = "groq"
-                text   = self.groq_transcribe(audio_to_process)
-            else:
-                source = "onnx"
+        # Parallel inference — di luar lock biar tidak blocking
+        text, source = self.parallel_transcribe(audio_to_process)
 
-            logger.info(f"[{source}] final: '{text}'")
+        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
+        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()
+        self.last_text = text
+        return text.strip()