From ed8f30ec71d2e7e753214e9e9594fbc061560a09 Mon Sep 17 00:00:00 2001 From: marcus-daily <111281783+marcus-daily@users.noreply.github.com> Date: Wed, 16 Jul 2025 12:13:25 +0100 Subject: [PATCH] Add support for running smart-turn-v2 locally --- examples/foundational/38b-smart-turn-local.py | 9 +- .../turn/smart_turn/local_smart_turn_v2.py | 190 ++++++++++++++++++ 2 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 src/pipecat/audio/turn/smart_turn/local_smart_turn_v2.py diff --git a/examples/foundational/38b-smart-turn-local.py b/examples/foundational/38b-smart-turn-local.py index 08f21f875..66e223098 100644 --- a/examples/foundational/38b-smart-turn-local.py +++ b/examples/foundational/38b-smart-turn-local.py @@ -12,6 +12,7 @@ from loguru import logger from pipecat.audio.turn.smart_turn.base_smart_turn import SmartTurnParams from pipecat.audio.turn.smart_turn.local_smart_turn import LocalSmartTurnAnalyzer +from pipecat.audio.turn.smart_turn.local_smart_turn_v2 import LocalSmartTurnAnalyzerV2 from pipecat.audio.vad.silero import SileroVADAnalyzer from pipecat.audio.vad.vad_analyzer import VADParams from pipecat.pipeline.pipeline import Pipeline @@ -37,7 +38,7 @@ load_dotenv(override=True) # # Hugging Face uses LFS to store large model files, including .mlpackage # git lfs install # # Clone the repo with the smart_turn_classifier.mlpackage -# git clone https://huggingface.co/pipecat-ai/smart-turn +# git clone https://huggingface.co/pipecat-ai/smart-turn-v2 # # Then set the env variable: # export LOCAL_SMART_TURN_MODEL_PATH=./smart-turn @@ -52,7 +53,7 @@ transport_params = { audio_in_enabled=True, audio_out_enabled=True, vad_analyzer=SileroVADAnalyzer(params=VADParams(stop_secs=0.2)), - turn_analyzer=LocalSmartTurnAnalyzer( + turn_analyzer=LocalSmartTurnAnalyzerV2( smart_turn_model_path=smart_turn_model_path, params=SmartTurnParams() ), ), @@ -60,7 +61,7 @@ transport_params = { audio_in_enabled=True, audio_out_enabled=True, vad_analyzer=SileroVADAnalyzer(params=VADParams(stop_secs=0.2)), - turn_analyzer=LocalSmartTurnAnalyzer( + turn_analyzer=LocalSmartTurnAnalyzerV2( smart_turn_model_path=smart_turn_model_path, params=SmartTurnParams() ), ), @@ -68,7 +69,7 @@ transport_params = { audio_in_enabled=True, audio_out_enabled=True, vad_analyzer=SileroVADAnalyzer(params=VADParams(stop_secs=0.2)), - turn_analyzer=LocalSmartTurnAnalyzer( + turn_analyzer=LocalSmartTurnAnalyzerV2( smart_turn_model_path=smart_turn_model_path, params=SmartTurnParams() ), ), diff --git a/src/pipecat/audio/turn/smart_turn/local_smart_turn_v2.py b/src/pipecat/audio/turn/smart_turn/local_smart_turn_v2.py new file mode 100644 index 000000000..07f28c901 --- /dev/null +++ b/src/pipecat/audio/turn/smart_turn/local_smart_turn_v2.py @@ -0,0 +1,190 @@ +# +# Copyright (c) 2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Local PyTorch turn analyzer for on-device ML inference using the smart-turn-v2 model. + +This module provides a smart turn analyzer that uses PyTorch models for +local end-of-turn detection without requiring network connectivity. +""" + +from typing import Any, Dict + +import numpy as np +from loguru import logger + +from pipecat.audio.turn.smart_turn.base_smart_turn import BaseSmartTurn + +try: + import torch + from torch import nn + import torch.nn.functional as F + from transformers import Wav2Vec2PreTrainedModel, Wav2Vec2Model, Wav2Vec2Processor, Wav2Vec2Config +except ModuleNotFoundError as e: + logger.error(f"Exception: {e}") + logger.error( + "In order to use LocalSmartTurnAnalyzerV2, you need to `pip install pipecat-ai[local-smart-turn]`." + ) + raise Exception(f"Missing module: {e}") + + +class LocalSmartTurnAnalyzerV2(BaseSmartTurn): + """Local turn analyzer using the smart-turn-v2 PyTorch model. + + Provides end-of-turn detection using locally-stored PyTorch models, + enabling offline operation without network dependencies. Uses + Wav2Vec2 architecture for audio sequence classification. + """ + + def __init__(self, *, smart_turn_model_path: str, **kwargs): + """Initialize the local PyTorch smart-turn-v2 analyzer. + + Args: + smart_turn_model_path: Path to directory containing the PyTorch model + and feature extractor files. If empty, uses default HuggingFace model. + **kwargs: Additional arguments passed to BaseSmartTurn. + """ + super().__init__(**kwargs) + + if not smart_turn_model_path: + # Define the path to the pretrained model on Hugging Face + smart_turn_model_path = "pipecat-ai/smart-turn-v2" + + logger.debug("Loading Local Smart Turn v2 model...") + # Load the pretrained model for sequence classification + self._turn_model = Wav2Vec2ForEndpointing.from_pretrained( + smart_turn_model_path + ) + # Load the corresponding feature extractor for preprocessing audio + self._turn_processor = Wav2Vec2Processor.from_pretrained(smart_turn_model_path) + # Set device to GPU if available, else CPU + self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # Move model to selected device and set it to evaluation mode + self._turn_model = self._turn_model.to(self._device) + self._turn_model.eval() + logger.debug("Loaded Local Smart Turn v2") + + async def _predict_endpoint(self, audio_array: np.ndarray) -> Dict[str, Any]: + """Predict end-of-turn using local PyTorch model.""" + inputs = self._turn_processor( + audio_array, + sampling_rate=16000, + padding="max_length", + truncation=True, + max_length=16000 * 16, # 16 seconds at 16kHz + return_attention_mask=True, + return_tensors="pt" + ) + + # Move inputs to device + inputs = {k: v.to(self._device) for k, v in inputs.items()} + + # Run inference + with torch.no_grad(): + outputs = self._turn_model(**inputs) + + # The model returns sigmoid probabilities directly in the logits field + probability = outputs["logits"][0].item() + + # Make prediction (1 for Complete, 0 for Incomplete) + prediction = 1 if probability > 0.5 else 0 + + return { + "prediction": prediction, + "probability": probability, + } + +class Wav2Vec2ForEndpointing(Wav2Vec2PreTrainedModel): + def __init__(self, config: Wav2Vec2Config): + super().__init__(config) + self.wav2vec2 = Wav2Vec2Model(config) + + self.pool_attention = nn.Sequential( + nn.Linear(config.hidden_size, 256), + nn.Tanh(), + nn.Linear(256, 1) + ) + + self.classifier = nn.Sequential( + nn.Linear(config.hidden_size, 256), + nn.LayerNorm(256), + nn.GELU(), + nn.Dropout(0.1), + nn.Linear(256, 64), + nn.GELU(), + nn.Linear(64, 1) + ) + + for module in self.classifier: + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=0.1) + if module.bias is not None: + module.bias.data.zero_() + + for module in self.pool_attention: + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=0.1) + if module.bias is not None: + module.bias.data.zero_() + + def attention_pool(self, hidden_states, attention_mask): + # Calculate attention weights + attention_weights = self.pool_attention(hidden_states) + + if attention_mask is None: + raise ValueError("attention_mask must be provided for attention pooling") + + attention_weights = attention_weights + ( + (1.0 - attention_mask.unsqueeze(-1).to(attention_weights.dtype)) * -1e9 + ) + + attention_weights = F.softmax(attention_weights, dim=1) + + # Apply attention to hidden states + weighted_sum = torch.sum(hidden_states * attention_weights, dim=1) + + return weighted_sum + + def forward(self, input_values, attention_mask=None, labels=None): + outputs = self.wav2vec2(input_values, attention_mask=attention_mask) + hidden_states = outputs[0] + + # Create transformer padding mask + if attention_mask is not None: + input_length = attention_mask.size(1) + hidden_length = hidden_states.size(1) + ratio = input_length / hidden_length + indices = (torch.arange(hidden_length, device=attention_mask.device) * ratio).long() + attention_mask = attention_mask[:, indices] + attention_mask = attention_mask.bool() + else: + attention_mask = None + + pooled = self.attention_pool(hidden_states, attention_mask) + + logits = self.classifier(pooled) + + if torch.isnan(logits).any(): + raise ValueError("NaN values detected in logits") + + if labels is not None: + # Calculate positive sample weight based on batch statistics + pos_weight = ((labels == 0).sum() / (labels == 1).sum()).clamp(min=0.1, max=10.0) + loss_fct = nn.BCEWithLogitsLoss(pos_weight=pos_weight) + labels = labels.float() + loss = loss_fct(logits.view(-1), labels.view(-1)) + + # Add L2 regularization for classifier layers + l2_lambda = 0.01 + l2_reg = torch.tensor(0., device=logits.device) + for param in self.classifier.parameters(): + l2_reg += torch.norm(param) + loss += l2_lambda * l2_reg + + probs = torch.sigmoid(logits.detach()) + return {"loss": loss, "logits": probs} + + probs = torch.sigmoid(logits) + return {"logits": probs} \ No newline at end of file