Integrating with the smart turn model to predict

src/pipecat/audio/turn/base_turn_analyzer.py

@@ -19,6 +19,7 @@ class BaseEndOfTurnAnalyzer(ABC):
     def __init__(self, *, sample_rate: Optional[int] = None):
         self._init_sample_rate = sample_rate
         self._sample_rate = 0
+        self._chunk_size_ms = 0
 
     @property
     def sample_rate(self) -> int:
@@ -27,6 +28,13 @@ class BaseEndOfTurnAnalyzer(ABC):
     def set_sample_rate(self, sample_rate: int):
         self._sample_rate = self._init_sample_rate or sample_rate
 
+    @property
+    def chunk_size_ms(self) -> int:
+        return self._chunk_size_ms
+
+    def set_chunk_size_ms(self, chunk_size_ms: int):
+        self._chunk_size_ms = chunk_size_ms
+
     @abstractmethod
-    def analyze_audio(self, buffer: bytes) -> EndOfTurnState:
+    def analyze_audio(self, buffer: bytes, is_speech: bool) -> EndOfTurnState:
         pass

src/pipecat/audio/turn/local_smart_turn.py

@@ -6,8 +6,10 @@
 
 
 import os
+import time
 
 import numpy as np
+import torch
 from loguru import logger
 
 from pipecat.audio.turn.base_turn_analyzer import BaseEndOfTurnAnalyzer, EndOfTurnState
@@ -23,13 +25,15 @@ except ModuleNotFoundError as e:
     raise Exception(f"Missing module: {e}")
 
 
+# TODO: we should convert all this to params
+STOP_MS = 1000
+PRE_SPEECH_MS = 200
+MAX_DURATION_SECONDS = 16  # Maximum duration for the smart turn model
+
+
 class LocalSmartTurnAnalyzer(BaseEndOfTurnAnalyzer):
     def __init__(self):
         super().__init__()
-        self._audio_buffer = bytearray()
-
-        logger.debug("Loading Local Smart Turn model...")
-
         # To use this locally, set the environment variable LOCAL_SMART_TURN_MODEL_PATH
         # to the path where the smart-turn repo is cloned.
         #
@@ -53,36 +57,145 @@ class LocalSmartTurnAnalyzer(BaseEndOfTurnAnalyzer):
 
         core_ml_model_path = f"{smart_turn_model_path}/coreml/smart_turn_classifier.mlpackage"
 
+        logger.debug("Loading Local Smart Turn model...")
         # Only load the processor, not the torch model
-        processor = AutoFeatureExtractor.from_pretrained(smart_turn_model_path)
+        self._turn_processor = AutoFeatureExtractor.from_pretrained(smart_turn_model_path)
-        model = ct.models.MLModel(core_ml_model_path)
+        self._turn_model = ct.models.MLModel(core_ml_model_path)
-
         logger.debug("Loaded Local Smart Turn")
 
-    def analyze_audio(self, buffer: bytes) -> EndOfTurnState:
+        self._audio_buffer = []
-        self._audio_buffer += buffer
+        self._speech_triggered = False
+        self._silence_frames = 0
+        self._speech_start_time = None
 
-        # TODO: we probably don't need this
+    def analyze_audio(self, buffer: bytes, is_speech: bool) -> EndOfTurnState:
-        # Checking if we have at least 6 seconds of audio
+        state = EndOfTurnState.INCOMPLETE
-        # if len(self._audio_buffer) < 16000 * 2 * 6:
-        #    return EndOfTurnState.INCOMPLETE
 
-        audio_int16 = np.frombuffer(self._audio_buffer, dtype=np.int16)
+        audio_int16 = np.frombuffer(buffer, dtype=np.int16)
 
         # Divide by 32768 because we have signed 16-bit data.
         audio_float32 = np.frombuffer(audio_int16, dtype=np.int16).astype(np.float32) / 32768.0
 
-        # TODO: implement to use the smart turn
+        if is_speech:
-        # for now it is always returning as complete only for testing it
+            if not self._speech_triggered:
-        prediction = 1
+                self._silence_frames = 0
-
+            self._speech_triggered = True
-        state = EndOfTurnState.COMPLETE if prediction == 1 else EndOfTurnState.INCOMPLETE
+            if self._speech_start_time is None:
-
+                self._speech_start_time = time.time()
-        if state == EndOfTurnState.COMPLETE:
+            self._audio_buffer.append((time.time(), audio_float32))
-            # clears the buffer completely
-            self._audio_buffer = bytearray()
         else:
-            # TODO: implement it
+            if self._speech_triggered:
-            pass
+                self._audio_buffer.append((time.time(), audio_float32))
+                self._silence_frames += 1
+                if self._silence_frames * self._chunk_size_ms >= STOP_MS:
+                    self._speech_triggered = False
+
+                    # TODO: do we need to stop or do something to prevent ??
+
+                    state = self._process_speech_segment(
+                        self._audio_buffer, self._speech_start_time
+                    )
+                    self._audio_buffer = []
+                    self._speech_start_time = None
+
+                    # TODO: same here for restart
+            else:
+                # Keep buffering some silence before potential speech starts
+                self._audio_buffer.append((time.time(), audio_float32))
+                # Keep the buffer size reasonable, assuming CHUNK is small
+                max_buffer_time = (
+                    PRE_SPEECH_MS + STOP_MS
+                ) / 1000 + MAX_DURATION_SECONDS  # Some extra buffer
+                while (
+                    self._audio_buffer and self._audio_buffer[0][0] < time.time() - max_buffer_time
+                ):
+                    self._audio_buffer.pop(0)
 
         return state
+
+    def _process_speech_segment(self, audio_buffer, speech_start_time) -> EndOfTurnState:
+        state = EndOfTurnState.INCOMPLETE
+
+        if not audio_buffer:
+            return state
+
+        # Find start and end indices for the segment
+        start_time = speech_start_time - (PRE_SPEECH_MS / 1000)
+        start_index = 0
+        for i, (t, _) in enumerate(audio_buffer):
+            if t >= start_time:
+                start_index = i
+                break
+
+        end_index = len(audio_buffer) - 1
+
+        # Extract the audio segment
+        segment_audio_chunks = [chunk for _, chunk in audio_buffer[start_index : end_index + 1]]
+        segment_audio = np.concatenate(segment_audio_chunks)
+
+        # Remove (STOP_MS - 200)ms from the end of the segment
+        samples_to_remove = int((STOP_MS - 200) / 1000 * self.sample_rate)
+        segment_audio = segment_audio[:-samples_to_remove]
+
+        # Limit maximum duration
+        if len(segment_audio) / self.sample_rate > MAX_DURATION_SECONDS:
+            segment_audio = segment_audio[: int(MAX_DURATION_SECONDS * self.sample_rate)]
+
+        # No resampling needed as both recording and prediction use 16000 Hz
+        segment_audio_resampled = segment_audio
+
+        if len(segment_audio_resampled) > 0:
+            # Call the new predict_endpoint function with the audio data
+            start_time = time.perf_counter()
+
+            result = self._predict_endpoint(segment_audio_resampled)
+
+            state = (
+                EndOfTurnState.COMPLETE if result["prediction"] == 1 else EndOfTurnState.INCOMPLETE
+            )
+
+            end_time = time.perf_counter()
+
+            logger.debug("--------")
+            logger.debug(f"Prediction: {'Complete' if result['prediction'] == 1 else 'Incomplete'}")
+            logger.debug(f"Probability of complete: {result['probability']:.4f}")
+            logger.debug(f"Prediction took {(end_time - start_time) * 1000:.2f}ms seconds")
+        else:
+            logger.debug("Captured empty audio segment, skipping prediction.")
+
+        return state
+
+    def _predict_endpoint(self, audio_array):
+        """
+        Predict whether an audio segment is complete (turn ended) or incomplete.
+
+        Args:
+            audio_array: Numpy array containing audio samples at 16kHz
+
+        Returns:
+            Dictionary containing prediction results:
+            - prediction: 1 for complete, 0 for incomplete
+            - probability: Probability of completion class
+        """
+
+        inputs = self._turn_processor(
+            audio_array,
+            sampling_rate=16000,
+            padding="max_length",
+            truncation=True,
+            max_length=800,  # Maximum length as specified in training
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+
+        output = self._turn_model.predict(dict(inputs))
+        logits = output["logits"]  # Core ML returns numpy array
+        logits_tensor = torch.tensor(logits)
+        probabilities = torch.nn.functional.softmax(logits_tensor, dim=1)
+        completion_prob = probabilities[0, 1].item()  # Probability of class 1 (Complete)
+        prediction = 1 if completion_prob > 0.5 else 0
+
+        return {
+            "prediction": prediction,
+            "probability": completion_prob,
+        }

src/pipecat/transports/base_input.py

@@ -79,6 +79,9 @@ class BaseInputTransport(FrameProcessor):
         # Configure End of turn analyzer.
         if self._params.end_of_turn_analyzer:
             self._params.end_of_turn_analyzer.set_sample_rate(self._sample_rate)
+            self._params.end_of_turn_analyzer.set_chunk_size_ms(
+                self._params.audio_out_10ms_chunks * 10
+            )
         # Start audio filter.
         if self._params.audio_in_filter:
             await self._params.audio_in_filter.start(self._sample_rate)
@@ -214,18 +217,23 @@ class BaseInputTransport(FrameProcessor):
             vad_state = new_vad_state
         return vad_state
 
-    async def _end_of_turn_analyze(self, audio_frame: InputAudioRawFrame) -> EndOfTurnState:
+    async def _end_of_turn_analyze(
+        self, audio_frame: InputAudioRawFrame, is_speech: bool
+    ) -> EndOfTurnState:
         state = EndOfTurnState.INCOMPLETE
         if self.end_of_turn_analyzer:
             state = await self.get_event_loop().run_in_executor(
-                self._executor, self.end_of_turn_analyzer.analyze_audio, audio_frame.audio
+                self._executor,
+                self.end_of_turn_analyzer.analyze_audio,
+                audio_frame.audio,
+                is_speech,
             )
         return state
 
     async def _handle_end_of_turn(
-        self, audio_frame: InputAudioRawFrame, end_of_turn_state: EndOfTurnState
+        self, audio_frame: InputAudioRawFrame, end_of_turn_state: EndOfTurnState, is_speech: bool
     ):
-        new_eot_state = await self._end_of_turn_analyze(audio_frame)
+        new_eot_state = await self._end_of_turn_analyze(audio_frame, is_speech)
         if new_eot_state != end_of_turn_state:
             await self._handle_user_interruption(UserEndOfTurnFrame())
         return new_eot_state
@@ -246,14 +254,13 @@ class BaseInputTransport(FrameProcessor):
             # changes from QUIET to SPEAKING and vice versa.
             if self._params.vad_enabled:
                 vad_state = await self._handle_vad(frame, vad_state)
-                # TODO: need to check if we need to keep it later
-                if vad_state == VADState.QUIET:
-                    end_of_turn_state = EndOfTurnState.INCOMPLETE
                 audio_passthrough = self._params.vad_audio_passthrough
 
-            # We only need to check for completion if the user is speaking
+            if self._params.end_of_turn_analyzer:
-            if self._params.end_of_turn_analyzer and VADState.QUIET != vad_state:
+                is_speech = vad_state == VADState.SPEAKING or vad_state == VADState.STARTING
-                end_of_turn_state = await self._handle_end_of_turn(frame, end_of_turn_state)
+                end_of_turn_state = await self._handle_end_of_turn(
+                    frame, end_of_turn_state, is_speech
+                )
 
             # Push audio downstream if passthrough.
             if audio_passthrough: