services: fix STTService and WhisperSTTService

CHANGELOG.md

@@ -13,6 +13,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `VADAnalyzer` arguments have been renamed for more clarity.
 
+### Fixed
+
+- Fixed `STTService`. Add `max_silence_secs` and `max_buffer_secs` to handle
+  better what's being passed to the STT service. Also add exponential smoothing
+  to the RMS.
+
+- Fixed `WhisperSTTService`. Add `no_speech_prob` to avoid garbage output text.
+
 ## [0.0.12] - 2024-05-14
 
 ### Added

src/pipecat/services/ai_services.py

@@ -10,10 +10,11 @@ import math
 import wave
 
 from abc import abstractmethod
-from typing import AsyncGenerator, BinaryIO
+from typing import AsyncGenerator
 
 from pipecat.frames.frames import (
     AudioRawFrame,
+    CancelFrame,
     EndFrame,
     ErrorFrame,
     Frame,
@@ -84,64 +85,80 @@ class STTService(AIService):
     """STTService is a base class for speech-to-text services."""
 
     def __init__(self,
-                 min_rms: int = 400,
+                 min_rms: int = 75,
-                 max_silence_frames: int = 3,
+                 max_silence_secs: float = 0.3,
+                 max_buffer_secs: float = 1.5,
                  sample_rate: int = 16000,
                  num_channels: int = 1):
         super().__init__()
         self._min_rms = min_rms
-        self._max_silence_frames = max_silence_frames
+        self._max_silence_secs = max_silence_secs
+        self._max_buffer_secs = max_buffer_secs
         self._sample_rate = sample_rate
         self._num_channels = num_channels
-        self._current_silence_frames = 0
         (self._content, self._wave) = self._new_wave()
+        self._silence_num_frames = 0
+        # Exponential smoothing
+        self._smoothing_factor = 0.08
+        self._prev_rms = 1 - self._smoothing_factor
 
     @abstractmethod
-    async def run_stt(self, audio: BinaryIO) -> AsyncGenerator[Frame, None]:
+    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
         """Returns transcript as a string"""
         pass
 
     def _new_wave(self):
-        content = io.BufferedRandom(io.BytesIO())
+        content = io.BytesIO()
         ww = wave.open(content, "wb")
         ww.setsampwidth(2)
         ww.setnchannels(self._num_channels)
         ww.setframerate(self._sample_rate)
         return (content, ww)
 
-    def _get_volume(self, audio: bytes) -> float:
+    def _exp_smoothing(self, value: float, prev_value: float, factor: float) -> float:
+        return prev_value + factor * (value - prev_value)
+
+    def _get_smoothed_volume(self, audio: bytes, prev_rms: float, factor: float) -> float:
         # https://docs.python.org/3/library/array.html
         audio_array = array.array('h', audio)
         squares = [sample**2 for sample in audio_array]
         mean = sum(squares) / len(audio_array)
         rms = math.sqrt(mean)
-        return rms
+        return self._exp_smoothing(rms, prev_rms, factor)
+
+    async def _append_audio(self, frame: AudioRawFrame):
+        # Try to filter out empty background noise
+        # (Very rudimentary approach, can be improved)
+        rms = self._get_smoothed_volume(frame.audio, self._prev_rms, self._smoothing_factor)
+        if rms >= self._min_rms:
+            # If volume is high enough, write new data to wave file
+            self._wave.writeframes(frame.audio)
+            self._silence_num_frames = 0
+        else:
+            self._silence_num_frames += frame.num_frames
+        self._prev_rms = rms
+
+        # If buffer is not empty and we have enough data or there's been a long
+        # silence, transcribe the audio gathered so far.
+        silence_secs = self._silence_num_frames / self._sample_rate
+        buffer_secs = self._wave.getnframes() / self._sample_rate
+        if self._content.tell() > 0 and (
+                buffer_secs > self._max_buffer_secs or silence_secs > self._max_silence_secs):
+            self._silence_num_frames = 0
+            self._wave.close()
+            self._content.seek(0)
+            await self.process_generator(self.run_stt(self._content.read()))
+            (self._content, self._wave) = self._new_wave()
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         """Processes a frame of audio data, either buffering or transcribing it."""
-        if not isinstance(frame, AudioRawFrame):
+        if isinstance(frame, CancelFrame) or isinstance(frame, EndFrame):
-            await self.push_frame(frame, direction)
-            return
-
-        audio = frame.audio
-
-        # Try to filter out empty background noise
-        # (Very rudimentary approach, can be improved)
-        rms = self._get_volume(audio)
-        if rms >= self._min_rms:
-            # If volume is high enough, write new data to wave file
-            self._wave.writeframes(audio)
-
-        # If buffer is not empty and we detect a 3-frame pause in speech,
-        # transcribe the audio gathered so far.
-        if self._content.tell() > 0 and self._current_silence_frames > self._max_silence_frames:
-            self._current_silence_frames = 0
             self._wave.close()
-            self._content.seek(0)
+            await self.push_frame(frame, direction)
-            await self.process_generator(self.run_stt(self._content))
+        elif isinstance(frame, AudioRawFrame):
-            (self._content, self._wave) = self._new_wave()
+            await self._append_audio(frame)
-        # If we get this far, this is a frame of silence
+        else:
-        self._current_silence_frames += 1
+            await self.push_frame(frame, direction)
 
 
 class ImageGenService(AIService):
@@ -150,7 +167,7 @@ class ImageGenService(AIService):
         super().__init__()
 
     # Renders the image. Returns an Image object.
-    @abstractmethod
+    @ abstractmethod
     async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
         pass
 
@@ -168,7 +185,7 @@ class VisionService(AIService):
         super().__init__()
         self._describe_text = None
 
-    @abstractmethod
+    @ abstractmethod
     async def run_vision(self, frame: VisionImageRawFrame) -> AsyncGenerator[Frame, None]:
         pass
 

src/pipecat/services/whisper.py

@@ -10,9 +10,11 @@ import asyncio
 import time
 
 from enum import Enum
-from typing import BinaryIO
+from typing_extensions import AsyncGenerator
 
-from pipecat.frames.frames import TranscriptionFrame
+import numpy as np
+
+from pipecat.frames.frames import ErrorFrame, Frame, TranscriptionFrame
 from pipecat.services.ai_services import STTService
 
 from loguru import logger
@@ -39,14 +41,18 @@ class Model(Enum):
 class WhisperSTTService(STTService):
     """Class to transcribe audio with a locally-downloaded Whisper model"""
 
-    def __init__(self, model_name: Model = Model.DISTIL_MEDIUM_EN,
+    def __init__(self,
+                 model: Model = Model.DISTIL_MEDIUM_EN,
                  device: str = "auto",
-                 compute_type: str = "default"):
+                 compute_type: str = "default",
+                 no_speech_prob: float = 0.1,
+                 **kwargs):
 
-        super().__init__()
+        super().__init__(**kwargs)
         self._device: str = device
         self._compute_type = compute_type
-        self._model_name: Model = model_name
+        self._model_name: Model = model
+        self._no_speech_prob = no_speech_prob
         self._model: WhisperModel | None = None
         self._load()
 
@@ -60,15 +66,21 @@ class WhisperSTTService(STTService):
             compute_type=self._compute_type)
         logger.debug("Loaded Whisper model")
 
-    async def run_stt(self, audio: BinaryIO):
+    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
         """Transcribes given audio using Whisper"""
         if not self._model:
+            yield ErrorFrame("Whisper model not available")
             logger.error("Whisper model not available")
             return
 
-        segments, _ = await asyncio.to_thread(self._model.transcribe, audio)
+        # Divide by 32768 because we have signed 16-bit data.
+        audio_float = np.frombuffer(audio, dtype=np.int16).astype(np.float32) / 32768.0
+
+        segments, _ = await asyncio.to_thread(self._model.transcribe, audio_float)
         text: str = ""
         for segment in segments:
-            text += f"{segment.text} "
+            if segment.no_speech_prob < self._no_speech_prob:
+                text += f"{segment.text} "
 
-        await self.push_frame(TranscriptionFrame(text, "", int(time.time_ns() / 1000000)))
+        if text:
+            yield TranscriptionFrame(text, "", int(time.time_ns() / 1000000))