wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Integrating with the smart turn model to predict

Browse Source
This commit is contained in:
Filipi Fuchter
2025-04-15 16:01:09 -03:00
parent 3588b06718
commit e6325a8229
3 changed files with 164 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/pipecat/audio/turn/base_turn_analyzer.py
View File

@@ -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

163
src/pipecat/audio/turn/local_smart_turn.py
View File

@@ -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,
}

27
src/pipecat/transports/base_input.py
View File

@@ -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: