diff --git a/src/pipecat/audio/turn/base_turn_analyzer.py b/src/pipecat/audio/turn/base_turn_analyzer.py index 1dffa563b..eeb9bfeb8 100644 --- a/src/pipecat/audio/turn/base_turn_analyzer.py +++ b/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 diff --git a/src/pipecat/audio/turn/local_smart_turn.py b/src/pipecat/audio/turn/local_smart_turn.py index 497487e06..d7733adc3 100644 --- a/src/pipecat/audio/turn/local_smart_turn.py +++ b/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) - model = ct.models.MLModel(core_ml_model_path) - + self._turn_processor = AutoFeatureExtractor.from_pretrained(smart_turn_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 += buffer + self._audio_buffer = [] + self._speech_triggered = False + self._silence_frames = 0 + self._speech_start_time = None - # TODO: we probably don't need this - # Checking if we have at least 6 seconds of audio - # if len(self._audio_buffer) < 16000 * 2 * 6: - # return EndOfTurnState.INCOMPLETE + def analyze_audio(self, buffer: bytes, is_speech: bool) -> EndOfTurnState: + state = 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 - # for now it is always returning as complete only for testing it - prediction = 1 - - state = EndOfTurnState.COMPLETE if prediction == 1 else EndOfTurnState.INCOMPLETE - - if state == EndOfTurnState.COMPLETE: - # clears the buffer completely - self._audio_buffer = bytearray() + if is_speech: + if not self._speech_triggered: + self._silence_frames = 0 + self._speech_triggered = True + if self._speech_start_time is None: + self._speech_start_time = time.time() + self._audio_buffer.append((time.time(), audio_float32)) else: - # TODO: implement it - pass + if self._speech_triggered: + 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, + } diff --git a/src/pipecat/transports/base_input.py b/src/pipecat/transports/base_input.py index 946feb063..4221b2f65 100644 --- a/src/pipecat/transports/base_input.py +++ b/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 and VADState.QUIET != vad_state: - end_of_turn_state = await self._handle_end_of_turn(frame, end_of_turn_state) + if self._params.end_of_turn_analyzer: + 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, is_speech + ) # Push audio downstream if passthrough. if audio_passthrough: