Merge pull request #3 from concil859856/ops

fix(stream): make /v1/stream actually work end-to-end

Author: concil859856

README.md

@@ -12,19 +12,20 @@ See [`stt-streaming-implementation.md`](../stt-streaming-implementation.md) for
 # 1. Build
 docker build -t vocence/stt-streaming:dev .
 
-# 2. Run on a CUDA-capable GPU
-docker run --rm --gpus all -p 8114:8114 \
+# 2. Run on a CUDA-capable GPU (serves on 8117 by default)
+docker run --rm --gpus all -p 8117:8117 \
   -e ASR_API_KEY=test_key_local_only \
   vocence/stt-streaming:dev
 
 # 3. Smoke test
 python scripts/smoke_test.py \
-  --url ws://localhost:8114/v1/stream \
+  --url ws://localhost:8117/v1/stream \
   --api-key test_key_local_only \
-  --wav tests/fixtures/librispeech_short.wav
+  --wav tests/fixtures/jfk.wav
 ```
 
-Expected: at least one `partial`, one `final` with the known transcript, clean close (1000).
+Expected: streaming `partial`s, one or more `final`s ("And so, my fellow Americans...
+ask what you can do for your country."), clean close (1000).
 
 ## Endpoints
 

src/stt_streaming/model.py

@@ -60,33 +60,50 @@ def __init__(self, model_dir: str) -> None:
     def reset_state(self) -> InferState:
         return InferState()
 
+    def transcribe_window(self, audio: np.ndarray) -> str:
+        """Stateless re-decode of a PCM window; returns the transcript.
+
+        Parakeet TDT via NeMo does not expose a stable per-chunk streaming step
+        for all checkpoints, so we re-decode the (bounded) cumulative utterance
+        window each call. The caller owns the audio buffer; this method keeps no
+        audio state of its own, which is what avoids double-counting.
+        """
+        if audio is None or audio.shape[0] == 0:
+            return ""
+        if audio.dtype != np.float32:
+            audio = audio.astype(np.float32)
+        # Bound the window so latency does not grow without limit on long utterances.
+        if audio.shape[0] > _FALLBACK_CONTEXT_SAMPLES:
+            audio = audio[-_FALLBACK_CONTEXT_SAMPLES:]
+        return self._run_transcribe(audio)
+
     def transcribe_chunk(
-        self, audio_chunk: np.ndarray, state: Optional[InferState]
+        self, audio_window: np.ndarray, state: Optional[InferState]
     ) -> Tuple[str, InferState]:
-        """Run incremental transcription on one PCM chunk; returns (cumulative_text, new_state)."""
+        """Re-decode the current utterance window; returns (cumulative_text, new_state).
+
+        ``audio_window`` is the full cumulative PCM for the utterance so far — the
+        caller (ws session) owns that buffer. We re-decode it rather than appending
+        to per-chunk state, so audio the caller has already buffered is never
+        counted twice. The GPU lock in ``_run_transcribe`` serializes device work.
+        """
         if state is None:
             state = self.reset_state()
-
-        if audio_chunk.dtype != np.float32:
-            audio_chunk = audio_chunk.astype(np.float32)
-
-        # Append to rolling buffer; we always re-decode the buffer because Parakeet TDT
-        # via NeMo 2.0.0 does not expose a stable per-chunk streaming step for all checkpoints.
-        # The lock keeps GPU access serialized; concurrency is achieved across sessions by queueing.
-        state.audio = np.concatenate([state.audio, audio_chunk])
-        # Trim to a bounded sliding window so latency does not grow without bound
-        if state.audio.shape[0] > _FALLBACK_CONTEXT_SAMPLES:
-            state.audio = state.audio[-_FALLBACK_CONTEXT_SAMPLES:]
-
-        text = self._run_transcribe(state.audio)
+        text = self.transcribe_window(audio_window)
         state.last_text = text
         return text, state
 
     def finalize(self, state: InferState) -> str:
-        """Final flush — return best hypothesis on whatever audio is buffered."""
-        if state is None or state.audio.shape[0] == 0:
-            return state.last_text if state else ""
-        return self._run_transcribe(state.audio)
+        """Final flush — best hypothesis on the last decoded window.
+
+        Prefer :meth:`transcribe_window` on the caller's live buffer; this is kept
+        for callers that only hold the opaque state.
+        """
+        if state is None:
+            return ""
+        if state.audio.shape[0] == 0:
+            return state.last_text
+        return self.transcribe_window(state.audio)
 
     def warm_up(self) -> None:
         """One dummy transcribe over 1 s of silence to compile kernels."""
@@ -98,17 +115,51 @@ def warm_up(self) -> None:
         except Exception:
             logger.exception("Warm-up failed (continuing)")
 
+    def _bound(self, audio: np.ndarray) -> Optional[np.ndarray]:
+        if audio is None or audio.shape[0] == 0:
+            return None
+        if audio.dtype != np.float32:
+            audio = audio.astype(np.float32)
+        if audio.shape[0] > _FALLBACK_CONTEXT_SAMPLES:
+            audio = audio[-_FALLBACK_CONTEXT_SAMPLES:]
+        return audio
+
+    def try_transcribe_window(self, audio: np.ndarray) -> Optional[str]:
+        """Best-effort partial decode: returns the transcript, or ``None`` if the
+        GPU is already busy.
+
+        Each NeMo ``transcribe`` call costs a fixed ~55 ms (Python/dataloader
+        overhead, independent of window length), so the lock-serialized model
+        caps out around ~18 calls/s. Partials are disposable — if we blocked
+        here, ticks from many concurrent sessions would queue without bound and
+        starve the event loop (uvicorn's WS keepalive then drops the socket).
+        Skipping a busy tick keeps the loop responsive; the next frame retries.
+        """
+        bounded = self._bound(audio)
+        if bounded is None:
+            return ""
+        if not self._lock.acquire(blocking=False):
+            return None
+        try:
+            return self._infer(bounded)
+        finally:
+            self._lock.release()
+
     def _run_transcribe(self, audio: np.ndarray) -> str:
-        """Underlying NeMo call. Holds the GPU lock for the duration."""
+        """Underlying NeMo call. Holds the GPU lock for the duration (blocking)."""
+        with self._lock:
+            return self._infer(audio)
+
+    def _infer(self, audio: np.ndarray) -> str:
+        """Run NeMo transcribe on one window. Caller must hold ``self._lock``."""
         import torch
 
-        with self._lock:
-            with torch.inference_mode():
-                # NeMo 2.0 transcribe accepts a list of numpy arrays in newer builds
-                try:
-                    out = self.model.transcribe([audio], batch_size=1, verbose=False)
-                except TypeError:
-                    out = self.model.transcribe([audio], batch_size=1)
+        with torch.inference_mode():
+            # NeMo 2.0 transcribe accepts a list of numpy arrays in newer builds
+            try:
+                out = self.model.transcribe([audio], batch_size=1, verbose=False)
+            except TypeError:
+                out = self.model.transcribe([audio], batch_size=1)
         return _extract_text(out)
 
 

src/stt_streaming/vad.py

@@ -1,17 +1,32 @@
 from __future__ import annotations
 
 import logging
+import threading
 from typing import List, Tuple
 
 import numpy as np
 
 log = logging.getLogger(__name__)
 
+# Silero v5 consumes fixed 512-sample windows at 16 kHz (256 @ 8 kHz). Feeding it
+# isolated 20 ms WS frames returns *no* speech ever — it needs streaming context —
+# which is why per-frame get_speech_timestamps() silently gated out every partial.
+_WINDOW_SAMPLES = 512
+# Hysteresis thresholds: enter speech high, leave low, so brief dips mid-word do
+# not bounce the gate and prematurely commit utterances.
+_SPEECH_ON = 0.5
+_SPEECH_OFF = 0.35
+
 
 class SileroVAD:
     def __init__(self) -> None:
         self._model = None
         self._get_speech_timestamps = None
+        # The silero model carries its LSTM state on the object itself, so a single
+        # shared instance is not safe across concurrent sessions. We serialize the
+        # restore->infer->save critical section and stash each session's state in
+        # the per-session dict returned by reset_state().
+        self._lock = threading.Lock()
         try:
             from silero_vad import load_silero_vad, get_speech_timestamps  # type: ignore
 
@@ -20,28 +35,67 @@ def __init__(self) -> None:
         except Exception as e:  # pragma: no cover - exercised only without silero installed
             log.warning("silero-vad unavailable (%s); falling back to RMS energy threshold", e)
 
-    def reset_state(self):
-        return None
+    def reset_state(self) -> dict:
+        # buf: leftover samples not yet forming a full 512 window
+        # triggered: current speech/silence latch (hysteresis)
+        # silero: this session's silero internals — recurrent _state plus the
+        #   64-sample _context carried between windows (None until first window)
+        return {"buf": np.zeros(0, dtype=np.float32), "triggered": False, "silero": None}
 
     def is_speech(self, audio_chunk: np.ndarray, state=None, sample_rate: int = 16000):
         if isinstance(state, int):
             sample_rate = state
             state = None
-        result = self._is_speech_inner(audio_chunk, sample_rate)
-        # ws.py expects (bool, new_state); internal callers ignore the tuple via [0]
-        return result, state
-
-    def _is_speech_inner(self, audio_chunk: np.ndarray, sample_rate: int = 16000) -> bool:
+        if state is None:
+            state = self.reset_state()
         if self._model is None:
-            return _rms_is_speech(audio_chunk)
+            return _rms_is_speech(audio_chunk), state
         try:
-            import torch
-            tensor = torch.from_numpy(_to_float32(audio_chunk))
-            ts = self._get_speech_timestamps(tensor, self._model, sampling_rate=sample_rate)
-            return len(ts) > 0
+            triggered = self._silero_stream(audio_chunk, state, sample_rate)
         except Exception as e:
-            log.warning("silero VAD chunk inference failed (%s); using RMS fallback", e)
-            return _rms_is_speech(audio_chunk)
+            log.warning("silero streaming VAD failed (%s); using RMS fallback", e)
+            return _rms_is_speech(audio_chunk), state
+        return triggered, state
+
+    def _silero_stream(self, audio_chunk: np.ndarray, state: dict, sample_rate: int) -> bool:
+        """Run silero over freshly-completed 512-sample windows, carrying this
+        session's recurrent state across calls. Concurrency-safe: the shared
+        model's state is swapped in/out under a lock per call."""
+        import torch
+
+        buf = np.concatenate([state["buf"], _to_float32(audio_chunk)])
+        max_prob: float | None = None
+        with self._lock:
+            # Load this session's silero internals into the shared model.
+            if state["silero"] is None:
+                self._model.reset_states()
+            else:
+                self._model._state = state["silero"]["state"]
+                self._model._context = state["silero"]["context"]
+                self._model._last_sr = sample_rate
+                self._model._last_batch_size = 1
+            with torch.no_grad():
+                while buf.shape[0] >= _WINDOW_SAMPLES:
+                    window = buf[:_WINDOW_SAMPLES]
+                    buf = buf[_WINDOW_SAMPLES:]
+                    p = float(self._model(torch.from_numpy(window.copy()), sample_rate).item())
+                    max_prob = p if max_prob is None else max(max_prob, p)
+            # Save this session's internals before another session borrows the model.
+            st = getattr(self._model, "_state", None)
+            ctx = getattr(self._model, "_context", None)
+            state["silero"] = {
+                "state": st.clone() if st is not None else None,
+                "context": ctx.clone() if ctx is not None else None,
+            }
+
+        state["buf"] = buf
+        if max_prob is not None:
+            if max_prob >= _SPEECH_ON:
+                state["triggered"] = True
+            elif max_prob < _SPEECH_OFF:
+                state["triggered"] = False
+        # No full window yet (sub-32 ms chunk): hold the previous latch.
+        return state["triggered"]
 
     def process_stream(self, audio_buffer: np.ndarray) -> List[Tuple[int, int, bool]]:
         sample_rate = 16000

src/stt_streaming/ws.py

@@ -163,7 +163,7 @@ async def handle_session(
         await _send_json(websocket, {
             "type": "ready",
             "session_id": sess.session_id,
-            "model": getattr(config, "model_name", model.model_dir),
+            "model": getattr(config, "model_name", None) or getattr(config, "model", None) or model.model_dir,
             "language": sess.language,
             "sample_rate": sess.sample_rate,
         })
@@ -219,14 +219,18 @@ async def _session_loop(
         if "bytes" in msg and msg["bytes"] is not None:
             frame = msg["bytes"]
             sample_count = len(frame) // 2
-            if len(frame) % 2 != 0 or sample_count < MIN_FRAME_SAMPLES:
-                await _send_error(ws, "audio_format_error", "frame too small or unaligned")
+            # Empty or odd-length frames are malformed PCM16 and unrecoverable.
+            if len(frame) == 0 or len(frame) % 2 != 0:
+                await _send_error(ws, "audio_format_error", "frame empty or unaligned")
                 await _close(ws, CLOSE_BAD_REQUEST, "bad frame")
                 return
             if sample_count > MAX_FRAME_SAMPLES:
                 await _send_error(ws, "audio_format_error", "frame too large")
                 await _close(ws, CLOSE_FRAME_TOO_LARGE, "frame too large")
                 return
+            # Frames smaller than MIN_FRAME_SAMPLES (e.g. a real-time client's final
+            # remainder chunk) are still valid audio — buffer them rather than killing
+            # the session over a short tail.
 
             sess.bytes_received += len(frame)
             chunk = _pcm16_to_float32(frame)
@@ -299,15 +303,15 @@ async def _session_loop(
 async def _emit_partial(ws: WebSocket, model: ParakeetModel, sess: WSSession) -> None:
     loop = asyncio.get_running_loop()
     audio = sess.utterance_buffer
-    state = sess.model_state
     try:
-        text, new_state = await loop.run_in_executor(
-            None, model.transcribe_chunk, audio, state
-        )
+        # Best-effort: returns None if the GPU is busy, in which case we drop this
+        # partial tick rather than queueing work behind the other live sessions.
+        text = await loop.run_in_executor(None, model.try_transcribe_window, audio)
     except Exception:
-        logger.exception("transcribe_chunk failed")
+        logger.exception("partial decode failed")
+        return
+    if text is None:
         return
-    sess.model_state = new_state
     if not text or text == sess.last_partial_text:
         return
     sess.last_partial_text = text
@@ -329,11 +333,18 @@ async def _finalize_utterance(
     if sess.utterance_buffer.shape[0] == 0 and not sess.last_partial_text:
         return
     loop = asyncio.get_running_loop()
-    try:
-        text = await loop.run_in_executor(None, model.finalize, sess.model_state)
-    except Exception:
-        logger.exception("finalize failed")
-        text = sess.last_partial_text
+    text = sess.last_partial_text
+    if sess.utterance_buffer.shape[0] > 0:
+        # Re-decode the caller-owned buffer directly. This is correct even when no
+        # partial was ever emitted (e.g. enable_partials=False), where the model
+        # state would otherwise hold no audio.
+        try:
+            text = await loop.run_in_executor(
+                None, model.transcribe_window, sess.utterance_buffer
+            )
+        except Exception:
+            logger.exception("finalize failed")
+            text = sess.last_partial_text
 
     if not text and not force:
         return