From b2754117c859417e513c3227f8da3ee46981d44b Mon Sep 17 00:00:00 2001 From: vipyne Date: Fri, 2 May 2025 11:34:40 -0500 Subject: [PATCH] Riva: refactor function_id and model_name --- .../07r-interruptible-riva-nim.py | 8 +-- src/pipecat/services/riva/stt.py | 62 ++++++++++++------- src/pipecat/services/riva/tts.py | 29 ++++++--- 3 files changed, 64 insertions(+), 35 deletions(-) diff --git a/examples/foundational/07r-interruptible-riva-nim.py b/examples/foundational/07r-interruptible-riva-nim.py index 7b82871ab..efe9b5479 100644 --- a/examples/foundational/07r-interruptible-riva-nim.py +++ b/examples/foundational/07r-interruptible-riva-nim.py @@ -41,14 +41,14 @@ async def run_bot(webrtc_connection: SmallWebRTCConnection, _: argparse.Namespac ), ) - stt = RivaSTTService(api_key=os.getenv("NVIDIA_API_KEY")) - # stt = RivaSegmentedSTTService(api_key=os.getenv("NVIDIA_API_KEY")) + # stt = RivaSTTService(api_key=os.getenv("NVIDIA_API_KEY")) # parakeet og + stt = RivaSegmentedSTTService(api_key=os.getenv("NVIDIA_API_KEY")) # canary # 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") - # tts = FastPitchTTSService(api_key=os.getenv("NVIDIA_API_KEY")) - tts = RivaTTSService(api_key=os.getenv("NVIDIA_API_KEY")) + tts = FastPitchTTSService(api_key=os.getenv("NVIDIA_API_KEY")) + # tts = RivaTTSService(api_key=os.getenv("NVIDIA_API_KEY")) messages = [ { diff --git a/src/pipecat/services/riva/stt.py b/src/pipecat/services/riva/stt.py index 0d15ab0af..2b263f03e 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, List, Optional +from typing import AsyncGenerator, List, Mapping, Optional from loguru import logger from pydantic import BaseModel @@ -93,8 +93,10 @@ class RivaSTTService(STTService): *, api_key: str, server: str = "grpc.nvcf.nvidia.com:443", - function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081", - model_name: str = "parakeet-1.1b-rnnt-multilingual-asr", + model_function_map: Mapping[str, str] = { + "function_id": "1598d209-5e27-4d3c-8079-4751568b1081", + "model_name": "parakeet-ctc-1.1b-asr", + }, sample_rate: Optional[int] = None, params: InputParams = InputParams(), **kwargs, @@ -114,11 +116,12 @@ class RivaSTTService(STTService): self._stop_history_eou = -1 self._stop_threshold_eou = -1.0 self._custom_configuration = "" + self._function_id = model_function_map.get("function_id") - self.set_model_name(model_name) + self.set_model_name(model_function_map.get("model_name")) metadata = [ - ["function-id", function_id], + ["function-id", self._function_id], ["authorization", f"Bearer {api_key}"], ] auth = riva.client.Auth(None, True, server, metadata) @@ -133,6 +136,13 @@ class RivaSTTService(STTService): def can_generate_metrics(self) -> bool: return False + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + async def start(self, frame: StartFrame): await super().start(frame) @@ -253,25 +263,22 @@ class RivaSTTService(STTService): class RivaSegmentedSTTService(SegmentedSTTService): - """Speech-to-text service using NVIDIA Riva Canary ASR API. + """Speech-to-text service using NVIDIA Riva's offline/batch models. - This service uses NVIDIA's Riva Canary ASR API to perform speech-to-text + By default, his 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 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 + model_function_map: Mapping of model name and its corresponding NVIDIA Cloud Function ID 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 Riva Canary ASR API.""" - language: Optional[Language] = Language.EN_US profanity_filter: bool = False automatic_punctuation: bool = True @@ -284,8 +291,10 @@ class RivaSegmentedSTTService(SegmentedSTTService): *, api_key: str, server: str = "grpc.nvcf.nvidia.com:443", - function_id: str = "ee8dc628-76de-4acc-8595-1836e7e857bd", - model_name: str = "canary-1b-asr", + model_function_map: Mapping[str, str] = { + "function_id": "ee8dc628-76de-4acc-8595-1836e7e857bd", + "model_name": "canary-1b-asr", + }, sample_rate: Optional[int] = None, params: InputParams = InputParams(), **kwargs, @@ -293,12 +302,13 @@ class RivaSegmentedSTTService(SegmentedSTTService): super().__init__(sample_rate=sample_rate, **kwargs) # Set model name - self.set_model_name(model_name) + self.set_model_name(model_function_map.get("model_name")) # Initialize Riva settings self._api_key = api_key self._server = server - self._function_id = function_id + self._function_id = model_function_map.get("function_id") + self._model_name = model_function_map.get("model_name") # Store the language as a Language enum and as a string self._language_enum = params.language or Language.EN_US @@ -344,7 +354,7 @@ class RivaSegmentedSTTService(SegmentedSTTService): 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}") + logger.info(f"Initialized RivaSegmentedSTTService with model: {self.model_name}") def _create_recognition_config(self): """Create the Riva ASR recognition configuration.""" @@ -384,6 +394,13 @@ class RivaSegmentedSTTService(SegmentedSTTService): """Indicates whether this service can generate processing metrics.""" return True + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + async def start(self, frame: StartFrame): """Initialize the service when the pipeline starts.""" await super().start(frame) @@ -402,7 +419,7 @@ class RivaSegmentedSTTService(SegmentedSTTService): self._config.language_code = self._language async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: - """Transcribe an audio segment using Riva Canary ASR. + """Transcribe an audio segment. Args: audio: Raw audio bytes in WAV format (already converted by base class). @@ -478,17 +495,18 @@ class ParakeetSTTService(RivaSTTService): *, api_key: str, server: str = "grpc.nvcf.nvidia.com:443", - function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081", - model_name: str = "parakeet-ctc-1.1b-asr", + model_function_map: Mapping[str, str] = { + "function_id": "1598d209-5e27-4d3c-8079-4751568b1081", + "model_name": "parakeet-ctc-1.1b-asr", + }, sample_rate: Optional[int] = None, - params: RivaSTTService.InputParams, # Use parent class's type + params: RivaSTTService.InputParams = RivaSTTService.InputParams(), # Use parent class's type **kwargs, ): super().__init__( api_key=api_key, server=server, - function_id=function_id, - model_name=model_name, + model_function_map=model_function_map, sample_rate=sample_rate, params=params, **kwargs, diff --git a/src/pipecat/services/riva/tts.py b/src/pipecat/services/riva/tts.py index 5ab0a0eaa..4fc8b8489 100644 --- a/src/pipecat/services/riva/tts.py +++ b/src/pipecat/services/riva/tts.py @@ -6,7 +6,7 @@ import asyncio import os -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, Mapping, Optional # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -46,8 +46,10 @@ class RivaTTSService(TTSService): server: str = "grpc.nvcf.nvidia.com:443", voice_id: str = "Magpie-Multilingual.EN-US.Ray", sample_rate: Optional[int] = None, - function_id: str = "877104f7-e885-42b9-8de8-f6e4c6303969", - model_name: str = "magpie-tts-multilingual", + model_function_map: Mapping[str, str] = { + "function_id": "877104f7-e885-42b9-8de8-f6e4c6303969", + "model_name": "magpie-tts-multilingual", + }, params: InputParams = InputParams(), **kwargs, ): @@ -56,12 +58,13 @@ class RivaTTSService(TTSService): self._voice_id = voice_id self._language_code = params.language self._quality = params.quality + self._function_id = model_function_map.get("function_id") - self.set_model_name(model_name) + self.set_model_name(model_function_map.get("model_name")) self.set_voice(voice_id) metadata = [ - ["function-id", function_id], + ["function-id", self._function_id], ["authorization", f"Bearer {api_key}"], ] auth = riva.client.Auth(None, True, server, metadata) @@ -73,6 +76,13 @@ class RivaTTSService(TTSService): riva.client.proto.riva_tts_pb2.RivaSynthesisConfigRequest() ) + async def set_model(self, model: str): + logger.warning(f"Cannot set model after initialization. Set model and function id like so:") + example = {"function_id": "", "model_name": ""} + logger.warning( + f"{self.__class__.__name__}(api_key=, model_function_map={example})" + ) + async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]: def read_audio_responses(queue: asyncio.Queue): def add_response(r): @@ -134,8 +144,10 @@ class FastPitchTTSService(RivaTTSService): server: str = "grpc.nvcf.nvidia.com:443", voice_id: str = "English-US.Female-1", sample_rate: Optional[int] = None, - function_id: str = "0149dedb-2be8-4195-b9a0-e57e0e14f972", - model_name: str = "fastpitch-hifigan-tts", + model_function_map: Mapping[str, str] = { + "function_id": "0149dedb-2be8-4195-b9a0-e57e0e14f972", + "model_name": "fastpitch-hifigan-tts", + }, params: InputParams = InputParams(), **kwargs, ): @@ -143,8 +155,7 @@ class FastPitchTTSService(RivaTTSService): api_key=api_key, voice_id=voice_id, sample_rate=sample_rate, - function_id=function_id, - model_name=model_name, + model_function_map=model_function_map, params=params, **kwargs, )