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

Riva: refactor function_id and model_name

Browse Source
This commit is contained in:
vipyne
2025-05-02 11:34:40 -05:00
parent 6c428c303b
commit b2754117c8
3 changed files with 64 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
examples/foundational/07r-interruptible-riva-nim.py
View File

@@ -41,14 +41,14 @@ async def run_bot(webrtc_connection: SmallWebRTCConnection, _: argparse.Namespac
), ),
) )
stt = RivaSTTService(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")) stt = RivaSegmentedSTTService(api_key=os.getenv("NVIDIA_API_KEY")) # canary
# stt = ParakeetSTTService(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") 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 = FastPitchTTSService(api_key=os.getenv("NVIDIA_API_KEY"))
tts = RivaTTSService(api_key=os.getenv("NVIDIA_API_KEY")) # tts = RivaTTSService(api_key=os.getenv("NVIDIA_API_KEY"))
messages = [ messages = [
{ {

62
src/pipecat/services/riva/stt.py
View File

@@ -5,7 +5,7 @@
# #
import asyncio import asyncio
from typing import AsyncGenerator, List, Optional from typing import AsyncGenerator, List, Mapping, Optional
from loguru import logger from loguru import logger
from pydantic import BaseModel from pydantic import BaseModel
@@ -93,8 +93,10 @@ class RivaSTTService(STTService):
*, *,
api_key: str, api_key: str,
server: str = "grpc.nvcf.nvidia.com:443", server: str = "grpc.nvcf.nvidia.com:443",
function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081", model_function_map: Mapping[str, str] = {
model_name: str = "parakeet-1.1b-rnnt-multilingual-asr", "function_id": "1598d209-5e27-4d3c-8079-4751568b1081",
"model_name": "parakeet-ctc-1.1b-asr",
},
sample_rate: Optional[int] = None, sample_rate: Optional[int] = None,
params: InputParams = InputParams(), params: InputParams = InputParams(),
**kwargs, **kwargs,
@@ -114,11 +116,12 @@ class RivaSTTService(STTService):
self._stop_history_eou = -1 self._stop_history_eou = -1
self._stop_threshold_eou = -1.0 self._stop_threshold_eou = -1.0
self._custom_configuration = "" 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 = [ metadata = [
["function-id", function_id], ["function-id", self._function_id],
["authorization", f"Bearer {api_key}"], ["authorization", f"Bearer {api_key}"],
] ]
auth = riva.client.Auth(None, True, server, metadata) auth = riva.client.Auth(None, True, server, metadata)
@@ -133,6 +136,13 @@ class RivaSTTService(STTService):
def can_generate_metrics(self) -> bool: def can_generate_metrics(self) -> bool:
return False 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": "<UUID>", "model_name": "<model_name>"}
logger.warning(
f"{self.__class__.__name__}(api_key=<api_key>, model_function_map={example})"
)
async def start(self, frame: StartFrame): async def start(self, frame: StartFrame):
await super().start(frame) await super().start(frame)
@@ -253,25 +263,22 @@ class RivaSTTService(STTService):
class RivaSegmentedSTTService(SegmentedSTTService): 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 transcription on audio segments. It inherits from SegmentedSTTService to handle
audio buffering and speech detection. audio buffering and speech detection.
Args: Args:
api_key: NVIDIA API key for authentication api_key: NVIDIA API key for authentication
server: Riva server address (defaults to NVIDIA Cloud Function endpoint) server: Riva server address (defaults to NVIDIA Cloud Function endpoint)
function_id: NVIDIA Cloud Function ID for the Canary ASR service model_function_map: Mapping of model name and its corresponding NVIDIA Cloud Function ID
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 sample_rate: Audio sample rate in Hz. If not provided, uses the pipeline's rate
params: Additional configuration parameters for Riva params: Additional configuration parameters for Riva
**kwargs: Additional arguments passed to SegmentedSTTService **kwargs: Additional arguments passed to SegmentedSTTService
""" """
class InputParams(BaseModel): class InputParams(BaseModel):
"""Configuration parameters for Riva Canary ASR API."""
language: Optional[Language] = Language.EN_US language: Optional[Language] = Language.EN_US
profanity_filter: bool = False profanity_filter: bool = False
automatic_punctuation: bool = True automatic_punctuation: bool = True
@@ -284,8 +291,10 @@ class RivaSegmentedSTTService(SegmentedSTTService):
*, *,
api_key: str, api_key: str,
server: str = "grpc.nvcf.nvidia.com:443", server: str = "grpc.nvcf.nvidia.com:443",
function_id: str = "ee8dc628-76de-4acc-8595-1836e7e857bd", model_function_map: Mapping[str, str] = {
model_name: str = "canary-1b-asr", "function_id": "ee8dc628-76de-4acc-8595-1836e7e857bd",
"model_name": "canary-1b-asr",
},
sample_rate: Optional[int] = None, sample_rate: Optional[int] = None,
params: InputParams = InputParams(), params: InputParams = InputParams(),
**kwargs, **kwargs,
@@ -293,12 +302,13 @@ class RivaSegmentedSTTService(SegmentedSTTService):
super().__init__(sample_rate=sample_rate, **kwargs) super().__init__(sample_rate=sample_rate, **kwargs)
# Set model name # Set model name
self.set_model_name(model_name) self.set_model_name(model_function_map.get("model_name"))
# Initialize Riva settings # Initialize Riva settings
self._api_key = api_key self._api_key = api_key
self._server = server 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 # Store the language as a Language enum and as a string
self._language_enum = params.language or Language.EN_US 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) auth = riva.client.Auth(None, True, self._server, metadata)
self._asr_service = riva.client.ASRService(auth) 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): def _create_recognition_config(self):
"""Create the Riva ASR recognition configuration.""" """Create the Riva ASR recognition configuration."""
@@ -384,6 +394,13 @@ class RivaSegmentedSTTService(SegmentedSTTService):
"""Indicates whether this service can generate processing metrics.""" """Indicates whether this service can generate processing metrics."""
return True 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": "<UUID>", "model_name": "<model_name>"}
logger.warning(
f"{self.__class__.__name__}(api_key=<api_key>, model_function_map={example})"
)
async def start(self, frame: StartFrame): async def start(self, frame: StartFrame):
"""Initialize the service when the pipeline starts.""" """Initialize the service when the pipeline starts."""
await super().start(frame) await super().start(frame)
@@ -402,7 +419,7 @@ class RivaSegmentedSTTService(SegmentedSTTService):
self._config.language_code = self._language self._config.language_code = self._language
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
"""Transcribe an audio segment using Riva Canary ASR. """Transcribe an audio segment.
Args: Args:
audio: Raw audio bytes in WAV format (already converted by base class). audio: Raw audio bytes in WAV format (already converted by base class).
@@ -478,17 +495,18 @@ class ParakeetSTTService(RivaSTTService):
*, *,
api_key: str, api_key: str,
server: str = "grpc.nvcf.nvidia.com:443", server: str = "grpc.nvcf.nvidia.com:443",
function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081", model_function_map: Mapping[str, str] = {
model_name: str = "parakeet-ctc-1.1b-asr", "function_id": "1598d209-5e27-4d3c-8079-4751568b1081",
"model_name": "parakeet-ctc-1.1b-asr",
},
sample_rate: Optional[int] = None, sample_rate: Optional[int] = None,
params: RivaSTTService.InputParams, # Use parent class's type params: RivaSTTService.InputParams = RivaSTTService.InputParams(), # Use parent class's type
**kwargs, **kwargs,
): ):
super().__init__( super().__init__(
api_key=api_key, api_key=api_key,
server=server, server=server,
function_id=function_id, model_function_map=model_function_map,
model_name=model_name,
sample_rate=sample_rate, sample_rate=sample_rate,
params=params, params=params,
**kwargs, **kwargs,

29
src/pipecat/services/riva/tts.py
View File

@@ -6,7 +6,7 @@
import asyncio import asyncio
import os import os
from typing import AsyncGenerator, Optional from typing import AsyncGenerator, Mapping, Optional
# Suppress gRPC fork warnings # Suppress gRPC fork warnings
os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false"
@@ -46,8 +46,10 @@ class RivaTTSService(TTSService):
server: str = "grpc.nvcf.nvidia.com:443", server: str = "grpc.nvcf.nvidia.com:443",
voice_id: str = "Magpie-Multilingual.EN-US.Ray", voice_id: str = "Magpie-Multilingual.EN-US.Ray",
sample_rate: Optional[int] = None, sample_rate: Optional[int] = None,
function_id: str = "877104f7-e885-42b9-8de8-f6e4c6303969", model_function_map: Mapping[str, str] = {
model_name: str = "magpie-tts-multilingual", "function_id": "877104f7-e885-42b9-8de8-f6e4c6303969",
"model_name": "magpie-tts-multilingual",
},
params: InputParams = InputParams(), params: InputParams = InputParams(),
**kwargs, **kwargs,
): ):
@@ -56,12 +58,13 @@ class RivaTTSService(TTSService):
self._voice_id = voice_id self._voice_id = voice_id
self._language_code = params.language self._language_code = params.language
self._quality = params.quality 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) self.set_voice(voice_id)
metadata = [ metadata = [
["function-id", function_id], ["function-id", self._function_id],
["authorization", f"Bearer {api_key}"], ["authorization", f"Bearer {api_key}"],
] ]
auth = riva.client.Auth(None, True, server, metadata) auth = riva.client.Auth(None, True, server, metadata)
@@ -73,6 +76,13 @@ class RivaTTSService(TTSService):
riva.client.proto.riva_tts_pb2.RivaSynthesisConfigRequest() 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": "<UUID>", "model_name": "<model_name>"}
logger.warning(
f"{self.__class__.__name__}(api_key=<api_key>, model_function_map={example})"
)
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]: async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
def read_audio_responses(queue: asyncio.Queue): def read_audio_responses(queue: asyncio.Queue):
def add_response(r): def add_response(r):
@@ -134,8 +144,10 @@ class FastPitchTTSService(RivaTTSService):
server: str = "grpc.nvcf.nvidia.com:443", server: str = "grpc.nvcf.nvidia.com:443",
voice_id: str = "English-US.Female-1", voice_id: str = "English-US.Female-1",
sample_rate: Optional[int] = None, sample_rate: Optional[int] = None,
function_id: str = "0149dedb-2be8-4195-b9a0-e57e0e14f972", model_function_map: Mapping[str, str] = {
model_name: str = "fastpitch-hifigan-tts", "function_id": "0149dedb-2be8-4195-b9a0-e57e0e14f972",
"model_name": "fastpitch-hifigan-tts",
},
params: InputParams = InputParams(), params: InputParams = InputParams(),
**kwargs, **kwargs,
): ):
@@ -143,8 +155,7 @@ class FastPitchTTSService(RivaTTSService):
api_key=api_key, api_key=api_key,
voice_id=voice_id, voice_id=voice_id,
sample_rate=sample_rate, sample_rate=sample_rate,
function_id=function_id, model_function_map=model_function_map,
model_name=model_name,
params=params, params=params,
**kwargs, **kwargs,
) )