Linting and Smart Turn Refactor

Author: LArmstrongDev

sdk/voice/speechmatics/voice/_smart_turn.py

@@ -78,6 +78,9 @@ class SmartTurnDetector:
     Further information at https://github.com/pipecat-ai/smart-turn
     """
 
+    WINDOW_SECONDS = 8
+    DEFAULT_SAMPLE_RATE = 16000
+
     def __init__(self, auto_init: bool = True, threshold: float = 0.8):
         """Create the new SmartTurnDetector.
 
@@ -125,7 +128,7 @@ def setup(self) -> None:
             self.session = self.build_session(SMART_TURN_MODEL_LOCAL_PATH)
 
             # Load the feature extractor
-            self.feature_extractor = WhisperFeatureExtractor(chunk_length=8)
+            self.feature_extractor = WhisperFeatureExtractor(chunk_length=self.WINDOW_SECONDS)
 
             # Set initialized
             self._is_initialized = True
@@ -156,83 +159,113 @@ def build_session(self, onnx_path: str) -> ort.InferenceSession:
         # Return the new session
         return ort.InferenceSession(onnx_path, sess_options=so)
 
-    async def predict(
-        self, audio_array: bytes, language: str, sample_rate: int = 16000, sample_width: int = 2
-    ) -> SmartTurnPredictionResult:
-        """Predict whether an audio segment is complete (turn ended) or incomplete.
+    def _prepare_audio(self, audio_array: bytes, sample_rate: int, sample_width: int) -> np.ndarray:
+        """Prepare the audio for inference.
 
         Args:
             audio_array: Numpy array containing audio samples at 16kHz. The function
                 will convert the audio into float32 and truncate to 8 seconds (keeping the end)
                 or pad to 8 seconds.
-            language: Language of the audio.
             sample_rate: Sample rate of the audio.
             sample_width: Sample width of the audio.
 
         Returns:
-            Prediction result containing completion status and probability.
+            Numpy array containing audio samples at 16kHz.
         """
-
-        # Check if initialized
-        if not self._is_initialized:
-            return SmartTurnPredictionResult(error="SmartTurnDetector is not initialized")
-
-        # Check a valid language
-        if not self.valid_language(language):
-            logger.warning(f"Invalid language: {language}. Results may be unreliable.")
-
-        # Record start time
-        start_time = datetime.datetime.now()
-
         # Convert into numpy array
         dtype = np.int16 if sample_width == 2 else np.int8
         int16_array: np.ndarray = np.frombuffer(audio_array, dtype=dtype).astype(np.int16)
 
-        # Truncate to last 8 seconds if needed (keep the tail/end of audio)
-        max_samples = 8 * sample_rate
+        # Truncate to last WINDOW_SECONDS seconds if needed (keep the tail/end of audio)
+        max_samples = self.WINDOW_SECONDS * sample_rate
         if len(int16_array) > max_samples:
             int16_array = int16_array[-max_samples:]
 
         # Convert int16 to float32 in range [-1, 1] (same as reference implementation)
         float32_array: np.ndarray = int16_array.astype(np.float32) / 32768.0
 
-        # Process audio using Whisper's feature extractor
+        return float32_array
+
+    def _get_input_features(self, audio_data: np.ndarray, sample_rate: int) -> np.ndarray:
+        """
+        Get the input features for the audio data using Whisper's feature extractor.
+
+        Args:
+            audio_data: Numpy array containing audio samples.
+            sample_rate: Sample rate of the audio.
+        """
+
         inputs = self.feature_extractor(
-            float32_array,
+            audio_data,
             sampling_rate=sample_rate,
             return_tensors="np",
             padding="max_length",
-            max_length=max_samples,
+            max_length= self.WINDOW_SECONDS * sample_rate,
             truncation=True,
             do_normalize=True,
         )
 
-        # Extract features and ensure correct shape for ONNX
+        # Ensure dimensions are correct shape for ONNX
         input_features = inputs.input_features.squeeze(0).astype(np.float32)
         input_features = np.expand_dims(input_features, axis=0)
 
-        # Run ONNX inference
-        outputs = self.session.run(None, {"input_features": input_features})
+        return input_features
+
+    async def predict(
+        self, audio_array: bytes, language: str, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2
+    ) -> SmartTurnPredictionResult:
+        """Predict whether an audio segment is complete (turn ended) or incomplete.
+
+        Args:
+            audio_array: Numpy array containing audio samples at 16kHz. The function
+                will convert the audio into float32 and truncate to 8 seconds (keeping the end)
+                or pad to 8 seconds.
+            language: Language of the audio.
+            sample_rate: Sample rate of the audio.
+            sample_width: Sample width of the audio.
+
+        Returns:
+            Prediction result containing completion status and probability.
+        """
 
-        # Extract probability (ONNX model returns sigmoid probabilities)
+        # Check if initialized
+        if not self._is_initialized:
+            return SmartTurnPredictionResult(error="SmartTurnDetector is not initialized")
+
+        # Check a valid language
+        if not self.valid_language(language):
+            logger.warning(f"Invalid language: {language}. Results may be unreliable.")
+
+        # Record start time
+        start_time = datetime.datetime.now()
+
+        # Convert the audio into required format
+        prepared_audio = self._prepare_audio(audio_array, sample_rate, sample_width)
+
+        # Feature extraction
+        input_features = self._get_input_features(prepared_audio, sample_rate)
+
+        # Model inference
+        outputs = self.session.run(None, {"input_features": input_features})
         probability = outputs[0][0].item()
 
         # Make prediction (True for Complete, False for Incomplete)
         prediction = probability >= self._threshold
 
-        # Record end time
+        # Result Formatting
         end_time = datetime.datetime.now()
+        duration = float((end_time - start_time).total_seconds())
 
         # Return the result
         return SmartTurnPredictionResult(
             prediction=prediction,
             probability=round(probability, 3),
-            processing_time=round(float((end_time - start_time).total_seconds()), 3),
+            processing_time=round(duration, 3),
         )
 
     @staticmethod
     def truncate_audio_to_last_n_seconds(
-        audio_array: np.ndarray, n_seconds: float = 8.0, sample_rate: int = 16000
+        audio_array: np.ndarray, n_seconds: float = 8.0, sample_rate: int = DEFAULT_SAMPLE_RATE
     ) -> np.ndarray:
         """Truncate audio to last n seconds or pad with zeros to meet n seconds.
 
@@ -300,7 +333,8 @@ def model_exists() -> bool:
 
     @staticmethod
     def valid_language(language: str) -> bool:
-        """Check if the language is valid.
+        """Check if the language is valid against list of supported languages
+        for the Pipecat model.
 
         Args:
             language: Language code to validate.

sdk/voice/speechmatics/voice/_vad.py

@@ -225,7 +225,7 @@ def process_chunk(self, chunk_f32: np.ndarray) -> float:
 
         # Return probability (out shape is (1, 1))
         return float(out[0][0])
-    
+
     def _validate_input(self, sample_rate: int) -> bool:
         """
         Ensures the VAD is ready and the incoming audio format
@@ -244,7 +244,7 @@ def _validate_input(self, sample_rate: int) -> bool:
         if sample_rate != SILERO_SAMPLE_RATE:
             logger.error(f"Sample rate must be {SILERO_SAMPLE_RATE}Hz, got {sample_rate}Hz")
             return False
-        
+
         return True
 
     def _get_audio_chunks(self, sample_width: int):
@@ -282,17 +282,17 @@ def _prepare_chunk(self, chunk_bytes: bytes, sample_width: int) -> np.ndarray:
         """
         if sample_width == 2:
             dtype = np.int16
-            divisor = 32768.0 
+            divisor = 32768.0
         elif sample_width == 1:
             dtype = np.int8
             divisor = 128.0
         else:
             raise ValueError(f"Unsupported sample_width {sample_width}")
 
-        # Decode and normalize the chunk data 
+        # Decode and normalize the chunk data
         int_array = np.frombuffer(chunk_bytes, dtype=dtype)
         float32_array: np.ndarray = int_array.astype(np.float32) / divisor
-        
+
         return float32_array
 
     def _evaluate_activity_change(self) -> None:
@@ -353,7 +353,7 @@ def _trigger_on_state_change(self, is_speech: bool, probability: float) -> None:
 
         # Trigger callback with result
         self._on_state_change(result)
-        
+
     async def process_audio(self, audio_bytes: bytes, sample_rate: int = 16000, sample_width: int = 2) -> None:
         """Process incoming audio bytes and invoke callback on state changes.
 
@@ -367,21 +367,21 @@ async def process_audio(self, audio_bytes: bytes, sample_rate: int = 16000, samp
         """
         if not self._validate_input(sample_rate):
             return
-        
+
         # Add new bytes to the buffer
         self._audio_buffer += audio_bytes
-        
+
         # Process all complete chunks in the buffer
         for chunk in self._get_audio_chunks(sample_width):
             audio_f32 = self._prepare_chunk(chunk, sample_width)
-            
+
             try:
                 probability = self.process_chunk(audio_f32)
                 self._prediction_window.append(probability)
             except Exception as e:
                 logger.error(f"Error processing VAD chunk: {e}")
                 continue
-        
+
         # Check if VAD state has changed
         self._evaluate_activity_change()