From daf9d47e58269a9f1338415843294462ed7eaf07 Mon Sep 17 00:00:00 2001 From: Mark Backman Date: Fri, 2 May 2025 09:52:23 -0400 Subject: [PATCH] Update RivaSegmentedSTTService --- .../07r-interruptible-riva-nim.py | 10 +- src/pipecat/services/riva/stt.py | 306 ++++++++++++------ src/pipecat/services/riva/tts.py | 5 +- 3 files changed, 219 insertions(+), 102 deletions(-) diff --git a/examples/foundational/07r-interruptible-riva-nim.py b/examples/foundational/07r-interruptible-riva-nim.py index ca74dfd7a..070467844 100644 --- a/examples/foundational/07r-interruptible-riva-nim.py +++ b/examples/foundational/07r-interruptible-riva-nim.py @@ -16,8 +16,12 @@ from pipecat.pipeline.runner import PipelineRunner from pipecat.pipeline.task import PipelineParams, PipelineTask from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext from pipecat.services.nim.llm import NimLLMService -from pipecat.services.riva.stt import ParakeetSTTService, RivaOfflineSTTService -from pipecat.services.riva.tts import RivaTTSService, FastPitchTTSService +from pipecat.services.riva.stt import ( + ParakeetSTTService, + RivaSegmentedSTTService, + RivaSTTService, +) +from pipecat.services.riva.tts import FastPitchTTSService, RivaTTSService from pipecat.transports.base_transport import TransportParams from pipecat.transports.network.small_webrtc import SmallWebRTCTransport from pipecat.transports.network.webrtc_connection import SmallWebRTCConnection @@ -37,7 +41,7 @@ async def run_bot(webrtc_connection: SmallWebRTCConnection, _: argparse.Namespac ), ) - stt = RivaOfflineSTTService(api_key=os.getenv("NVIDIA_API_KEY")) + stt = RivaSegmentedSTTService(api_key=os.getenv("NVIDIA_API_KEY")) # stt = ParakeetSTTService(api_key=os.getenv("NVIDIA_API_KEY")) llm = NimLLMService(api_key=os.getenv("NVIDIA_API_KEY"), model="meta/llama-3.1-405b-instruct") diff --git a/src/pipecat/services/riva/stt.py b/src/pipecat/services/riva/stt.py index 3676befa3..d3b826060 100644 --- a/src/pipecat/services/riva/stt.py +++ b/src/pipecat/services/riva/stt.py @@ -5,7 +5,7 @@ # import asyncio -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, List, Optional from loguru import logger from pydantic import BaseModel @@ -19,8 +19,7 @@ from pipecat.frames.frames import ( StartFrame, TranscriptionFrame, ) -from pipecat.services.stt_service import SegmentedSTTService -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import SegmentedSTTService, STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 @@ -29,10 +28,62 @@ try: except ModuleNotFoundError as e: logger.error(f"Exception: {e}") - logger.error("In order to use NVIDIA Riva STT, you need to `pip install pipecat-ai[riva]`. Also set NVIDIA_API_KEY env var.") + logger.error("In order to use NVIDIA Riva STT, you need to `pip install pipecat-ai[riva]`.") raise Exception(f"Missing module: {e}") +def language_to_riva_language(language: Language) -> Optional[str]: + """Maps Language enum to Riva ASR language codes. + + Source: + https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-riva-build-table.html?highlight=fr%20fr + + Args: + language: Language enum value. + + Returns: + Optional[str]: Riva language code or None if not supported. + """ + language_map = { + # Arabic + Language.AR: "ar-AR", + # English + Language.EN: "en-US", # Default to US + Language.EN_US: "en-US", + Language.EN_GB: "en-GB", + # French + Language.FR: "fr-FR", + Language.FR_FR: "fr-FR", + # German + Language.DE: "de-DE", + Language.DE_DE: "de-DE", + # Hindi + Language.HI: "hi-IN", + Language.HI_IN: "hi-IN", + # Italian + Language.IT: "it-IT", + Language.IT_IT: "it-IT", + # Japanese + Language.JA: "ja-JP", + Language.JA_JP: "ja-JP", + # Korean + Language.KO: "ko-KR", + Language.KO_KR: "ko-KR", + # Portuguese + Language.PT: "pt-BR", # Default to Brazilian + Language.PT_BR: "pt-BR", + # Russian + Language.RU: "ru-RU", + Language.RU_RU: "ru-RU", + # Spanish + Language.ES: "es-ES", # Default to Spain + Language.ES_ES: "es-ES", + Language.ES_US: "es-US", # US Spanish + } + + return language_map.get(language) + + class RivaSTTService(STTService): class InputParams(BaseModel): language: Optional[Language] = Language.EN_US @@ -200,38 +251,38 @@ class RivaSTTService(STTService): def __iter__(self): return self -class RivaOfflineSTTService(SegmentedSTTService): - """Speech-to-text service using Fal's Wizper API. - This service uses Fal's Wizper API to perform speech-to-text transcription on audio - segments. It inherits from SegmentedSTTService to handle audio buffering and speech detection. +class RivaSegmentedSTTService(SegmentedSTTService): + """Speech-to-text service using NVIDIA Riva Canary ASR API. + + This service uses NVIDIA's Riva Canary ASR API to perform speech-to-text + transcription on audio segments. It inherits from SegmentedSTTService to handle + audio buffering and speech detection. Args: - api_key: NVIDIA_API_KEY. - sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate. - params: Configuration parameters for Riva. - **kwargs: Additional arguments passed to SegmentedSTTService. + api_key: NVIDIA API key for authentication + server: Riva server address (defaults to NVIDIA Cloud Function endpoint) + function_id: NVIDIA Cloud Function ID for the Canary ASR service + model_name: Name of the Canary ASR model to use + sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate + params: Additional configuration parameters for Riva + **kwargs: Additional arguments passed to SegmentedSTTService """ class InputParams(BaseModel): - """Configuration parameters for Fal's Wizper API. + """Configuration parameters for Riva Canary ASR API.""" - Attributes: - language: Language of the audio input. Defaults to English. - task: Task to perform ('transcribe' or 'translate'). Defaults to 'transcribe'. - chunk_level: Level of chunking ('segment'). Defaults to 'segment'. - version: Version of Wizper model to use. Defaults to '3'. - """ - - language: Optional[Language] = Language.EN - task: str = "transcribe" - chunk_level: str = "segment" - version: str = "3" + language: Optional[Language] = Language.EN_US + profanity_filter: bool = False + automatic_punctuation: bool = True + verbatim_transcripts: bool = False + boosted_lm_words: Optional[List[str]] = None + boosted_lm_score: float = 4.0 def __init__( self, *, - api_key: str = None, + api_key: str, server: str = "grpc.nvcf.nvidia.com:443", function_id: str = "ee8dc628-76de-4acc-8595-1836e7e857bd", model_name: str = "canary-1b-asr", @@ -240,13 +291,27 @@ class RivaOfflineSTTService(SegmentedSTTService): **kwargs, ): super().__init__(sample_rate=sample_rate, **kwargs) + + # Set model name + self.set_model_name(model_name) + + # Initialize Riva settings self._api_key = api_key - self._profanity_filter = False - self._automatic_punctuation = False - self._no_verbatim_transcripts = False - self._language_code = params.language - self._boosted_lm_words = None - self._boosted_lm_score = 4.0 + self._server = server + self._function_id = function_id + + # Store the language as a Language enum and as a string + self._language_enum = params.language or Language.EN_US + self._language = self.language_to_service_language(self._language_enum) or "en-US" + + # Configure transcription parameters + self._profanity_filter = params.profanity_filter + self._automatic_punctuation = params.automatic_punctuation + self._verbatim_transcripts = params.verbatim_transcripts + self._boosted_lm_words = params.boosted_lm_words + self._boosted_lm_score = params.boosted_lm_score + + # Voice activity detection thresholds (use Riva defaults) self._start_history = -1 self._start_threshold = -1.0 self._stop_history = -1 @@ -255,50 +320,50 @@ class RivaOfflineSTTService(SegmentedSTTService): self._stop_threshold_eou = -1.0 self._custom_configuration = "" - self.set_model_name(model_name) - - metadata = [ - ["function-id", function_id], - ["authorization", f"Bearer {api_key}"], - ] - auth = riva.client.Auth(None, True, server, metadata) - - self._asr_service = riva.client.ASRService(auth) - - self._queue = asyncio.Queue() + # Create Riva client self._config = None - self._thread_task = None - self._response_task = None + self._asr_service = None + self._settings = {"language": self._language_enum} - def can_generate_metrics(self) -> bool: - return False + def language_to_service_language(self, language: Language) -> Optional[str]: + """Convert pipecat Language enum to Riva's language code.""" + return language_to_riva_language(language) - async def start(self, frame: StartFrame): - await super().start(frame) - - if self._config: + def _initialize_client(self): + """Initialize the Riva ASR client with authentication metadata.""" + if self._asr_service is not None: return - # config = riva.client.StreamingRecognitionConfig( - config=riva.client.RecognitionConfig( - # encoding=riva.client.AudioEncoding.LINEAR_PCM, - language_code=self._language_code, - # model="", + # Set up authentication metadata for NVIDIA Cloud Functions + metadata = [ + ["function-id", self._function_id], + ["authorization", f"Bearer {self._api_key}"], + ] + + # Create authenticated client + auth = riva.client.Auth(None, True, self._server, metadata) + self._asr_service = riva.client.ASRService(auth) + + logger.info(f"Initialized Riva Canary ASR service with model: {self.model_name}") + + def _create_recognition_config(self): + """Create the Riva ASR recognition configuration.""" + # Create base configuration + config = riva.client.RecognitionConfig( + language_code=self._language, # Now using the string, not a tuple max_alternatives=1, profanity_filter=self._profanity_filter, enable_automatic_punctuation=self._automatic_punctuation, - verbatim_transcripts=not self._no_verbatim_transcripts, - # sample_rate_hertz=self.sample_rate, - # audio_channel_count=1, - # enable_word_time_offsets=args.word_time_offsets or args.speaker_diarization,?? - # ), - # interim_results=True, + verbatim_transcripts=self._verbatim_transcripts, ) - riva.client.add_word_boosting_to_config( - config, self._boosted_lm_words, self._boosted_lm_score - ) + # Add word boosting if specified + if self._boosted_lm_words: + riva.client.add_word_boosting_to_config( + config, self._boosted_lm_words, self._boosted_lm_score + ) + # Add voice activity detection parameters riva.client.add_endpoint_parameters_to_config( config, self._start_history, @@ -308,56 +373,102 @@ class RivaOfflineSTTService(SegmentedSTTService): self._stop_threshold, self._stop_threshold_eou, ) - riva.client.add_custom_configuration_to_config(config, self._custom_configuration) - self._config = config + # Add any custom configuration + if self._custom_configuration: + riva.client.add_custom_configuration_to_config(config, self._custom_configuration) + return config + def can_generate_metrics(self) -> bool: + """Indicates whether this service can generate processing metrics.""" + return True + + async def start(self, frame: StartFrame): + """Initialize the service when the pipeline starts.""" + await super().start(frame) + self._initialize_client() + self._config = self._create_recognition_config() + + async def set_language(self, language: Language): + """Set the language for the STT service.""" + logger.info(f"Switching STT language to: [{language}]") + self._language_enum = language + self._language = self.language_to_service_language(language) or "en-US" + self._settings["language"] = language + + # Update configuration with new language + if self._config: + self._config.language_code = self._language async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: - """Transcribe an audio segment + """Transcribe an audio segment using Riva Canary ASR. Args: audio: Raw audio bytes in WAV format (already converted by base class). Yields: Frame: TranscriptionFrame containing the transcribed text. - - Note: - The audio is already in WAV format from the SegmentedSTTService. - Only non-empty transcriptions are yielded. """ try: - response = self._asr_service.offline_recognize(audio, self._config) - # response = riva.client.print_offline(response=self._asr_service.offline_recognize(audio, self._config)) - print(f"_____stt.py * response: {response}") - # # Send to Fal directly (audio is already in WAV format from base class) - # data_uri = fal_client.encode(audio, "audio/x-wav") - # response = await self._fal_client.run( - # "fal-ai/wizper", - # arguments={"audio_url": data_uri, **self._settings}, - # ) + await self.start_processing_metrics() + await self.start_ttfb_metrics() - if response and "text" in response: - text = response["text"].strip() - if text: # Only yield non-empty text - logger.debug(f"Transcription: [{text}]") - yield TranscriptionFrame( - text, "", time_now_iso8601(), Language(self._settings["language"]) - ) + # Make sure the client is initialized + if self._asr_service is None: + self._initialize_client() + + # Make sure the config is created + if self._config is None: + self._config = self._create_recognition_config() + + # Type assertion to satisfy the IDE + assert self._asr_service is not None, "ASR service not initialized" + assert self._config is not None, "Recognition config not created" + + # Process audio with Riva ASR - explicitly request non-future response + raw_response = self._asr_service.offline_recognize(audio, self._config, future=False) + + await self.stop_ttfb_metrics() + await self.stop_processing_metrics() + + # Process the response - handle different possible return types + try: + # If it's a future-like object, get the result + if hasattr(raw_response, "result"): + response = raw_response.result() + else: + response = raw_response + + # Process transcription results + transcription_found = False + + # Now we can safely check results + # Type hint for the IDE + results = getattr(response, "results", []) + + for result in results: + alternatives = getattr(result, "alternatives", []) + if alternatives: + text = alternatives[0].transcript.strip() + if text: + logger.debug(f"Transcription: [{text}]") + yield TranscriptionFrame( + text, "", time_now_iso8601(), self._language_enum + ) + transcription_found = True + + if not transcription_found: + logger.debug("No transcription results found in Riva response") + + except AttributeError as ae: + logger.error(f"Unexpected response structure from Riva: {ae}") + yield ErrorFrame(f"Unexpected Riva response format: {str(ae)}") except Exception as e: - logger.error(f"Riva Offline STT error: {e}") - yield ErrorFrame(f"Riva Offline STT error: {str(e)}") + logger.exception(f"Riva Canary ASR error: {e}") + yield ErrorFrame(f"Riva Canary ASR error: {str(e)}") - def __next__(self) -> bytes: - if not self._thread_running: - raise StopIteration - future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop()) - return future.result() - - def __iter__(self): - return self class ParakeetSTTService(RivaSTTService): class InputParams(BaseModel): @@ -391,4 +502,3 @@ class ParakeetSTTService(RivaSTTService): "`ParakeetSTTService` is deprecated, use `RivaSTTService` instead.", DeprecationWarning, ) - diff --git a/src/pipecat/services/riva/tts.py b/src/pipecat/services/riva/tts.py index f37064acb..bfa732678 100644 --- a/src/pipecat/services/riva/tts.py +++ b/src/pipecat/services/riva/tts.py @@ -5,8 +5,12 @@ # import asyncio +import os from typing import AsyncGenerator, Optional +# Suppress gRPC fork warnings +os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" + from loguru import logger from pydantic import BaseModel @@ -152,4 +156,3 @@ class FastPitchTTSService(RivaTTSService): "`FastPitchTTSService` is deprecated, use `RivaTTSService` instead.", DeprecationWarning, ) -