wehub-resource-sync
335 lines
12 KiB
Python
335 lines
12 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
import queue
|
|
import threading
|
|
from collections import Counter
|
|
from collections.abc import Generator
|
|
from dataclasses import dataclass
|
|
|
|
import riva.client
|
|
from riva.client.proto.riva_asr_pb2 import SpeakerDiarizationConfig
|
|
|
|
from livekit import rtc
|
|
from livekit.agents import (
|
|
DEFAULT_API_CONNECT_OPTIONS,
|
|
APIConnectOptions,
|
|
LanguageCode,
|
|
stt,
|
|
)
|
|
from livekit.agents.types import NOT_GIVEN, NotGivenOr
|
|
from livekit.agents.utils import AudioBuffer, is_given
|
|
from livekit.agents.voice.io import TimedString
|
|
|
|
from . import auth
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class STTOptions:
|
|
model: str
|
|
function_id: str
|
|
punctuate: bool
|
|
language_code: LanguageCode
|
|
sample_rate: int
|
|
use_ssl: bool
|
|
server: str
|
|
enable_diarization: bool
|
|
max_speaker_count: int
|
|
|
|
|
|
class STT(stt.STT):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: str = "parakeet-1.1b-en-US-asr-streaming-silero-vad-sortformer",
|
|
function_id: str = "1598d209-5e27-4d3c-8079-4751568b1081",
|
|
punctuate: bool = True,
|
|
language_code: str = "en-US",
|
|
sample_rate: int = 16000,
|
|
server: str = "grpc.nvcf.nvidia.com:443",
|
|
use_ssl: bool = True,
|
|
api_key: NotGivenOr[str] = NOT_GIVEN,
|
|
enable_diarization: bool = False,
|
|
max_speaker_count: int = 0,
|
|
):
|
|
super().__init__(
|
|
capabilities=stt.STTCapabilities(
|
|
streaming=True,
|
|
interim_results=True,
|
|
diarization=enable_diarization,
|
|
aligned_transcript="word",
|
|
offline_recognize=False,
|
|
),
|
|
)
|
|
|
|
if enable_diarization and "sortformer" not in model:
|
|
logger.warning(
|
|
"Speaker diarization is enabled but model '%s' may not support it. "
|
|
"Diarization requires a Sortformer-based model.",
|
|
model,
|
|
)
|
|
|
|
if is_given(api_key):
|
|
self.nvidia_api_key = api_key
|
|
else:
|
|
self.nvidia_api_key = os.getenv("NVIDIA_API_KEY", "")
|
|
if use_ssl and not self.nvidia_api_key:
|
|
raise ValueError(
|
|
"NVIDIA_API_KEY is not set while using SSL. Either pass api_key parameter, set NVIDIA_API_KEY environment variable "
|
|
+ "or disable SSL and use a locally hosted Riva NIM service."
|
|
)
|
|
|
|
logger.info(f"Initializing NVIDIA STT with model: {model}, server: {server}")
|
|
logger.debug(
|
|
f"Function ID: {function_id}, LanguageCode: {language_code}, Sample rate: {sample_rate}"
|
|
)
|
|
|
|
self._opts = STTOptions(
|
|
model=model,
|
|
function_id=function_id,
|
|
punctuate=punctuate,
|
|
language_code=LanguageCode(language_code),
|
|
sample_rate=sample_rate,
|
|
server=server,
|
|
use_ssl=use_ssl,
|
|
enable_diarization=enable_diarization,
|
|
max_speaker_count=max_speaker_count,
|
|
)
|
|
|
|
def _recognize_impl(
|
|
self,
|
|
buffer: AudioBuffer,
|
|
*,
|
|
language: NotGivenOr[str] = NOT_GIVEN,
|
|
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
|
|
) -> stt.SpeechEvent:
|
|
raise NotImplementedError("Not implemented")
|
|
|
|
def stream(
|
|
self,
|
|
*,
|
|
language: NotGivenOr[str] = NOT_GIVEN,
|
|
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
|
|
) -> stt.RecognizeStream:
|
|
effective_language = (
|
|
LanguageCode(language) if is_given(language) else self._opts.language_code
|
|
)
|
|
return SpeechStream(stt=self, conn_options=conn_options, language=effective_language)
|
|
|
|
def log_asr_models(self, asr_service: riva.client.ASRService) -> dict:
|
|
config_response = asr_service.stub.GetRivaSpeechRecognitionConfig(
|
|
riva.client.RivaSpeechRecognitionConfigRequest()
|
|
)
|
|
|
|
asr_models = {}
|
|
for model_config in config_response.model_config:
|
|
if model_config.parameters.get("type") == "online":
|
|
language_code = model_config.parameters["language_code"]
|
|
model = {"model": [model_config.model_name]}
|
|
if language_code in asr_models:
|
|
asr_models[language_code].append(model)
|
|
else:
|
|
asr_models[language_code] = [model]
|
|
|
|
asr_models = dict(sorted(asr_models.items()))
|
|
return asr_models
|
|
|
|
|
|
class SpeechStream(stt.SpeechStream):
|
|
def __init__(self, *, stt: STT, conn_options: APIConnectOptions, language: str):
|
|
super().__init__(stt=stt, conn_options=conn_options, sample_rate=stt._opts.sample_rate)
|
|
self._stt = stt
|
|
self._language = language
|
|
|
|
self._audio_queue = queue.Queue()
|
|
self._shutdown_event = threading.Event()
|
|
self._recognition_thread = None
|
|
|
|
self._speaking = False
|
|
self._request_id = ""
|
|
|
|
self._auth = auth.create_riva_auth(
|
|
api_key=self._stt.nvidia_api_key,
|
|
function_id=self._stt._opts.function_id,
|
|
server=stt._opts.server,
|
|
use_ssl=stt._opts.use_ssl,
|
|
)
|
|
self._asr_service = riva.client.ASRService(self._auth)
|
|
|
|
self._event_loop = asyncio.get_running_loop()
|
|
self._done_fut = asyncio.Future()
|
|
|
|
async def _run(self) -> None:
|
|
config = self._create_streaming_config()
|
|
|
|
self._recognition_thread = threading.Thread(
|
|
target=self._recognition_worker,
|
|
args=(config,),
|
|
name="nvidia-asr-recognition",
|
|
daemon=True,
|
|
)
|
|
self._recognition_thread.start()
|
|
|
|
try:
|
|
await self._collect_audio()
|
|
|
|
finally:
|
|
self._audio_queue.put(None)
|
|
await self._done_fut
|
|
|
|
def _create_streaming_config(self) -> riva.client.StreamingRecognitionConfig:
|
|
recognition_config = riva.client.RecognitionConfig(
|
|
encoding=riva.client.AudioEncoding.LINEAR_PCM,
|
|
language_code=self._language,
|
|
model=self._stt._opts.model,
|
|
max_alternatives=1,
|
|
enable_automatic_punctuation=self._stt._opts.punctuate,
|
|
sample_rate_hertz=self._stt._opts.sample_rate,
|
|
audio_channel_count=1,
|
|
enable_word_time_offsets=True,
|
|
)
|
|
|
|
if self._stt._opts.enable_diarization:
|
|
diarization_config = SpeakerDiarizationConfig(
|
|
enable_speaker_diarization=True,
|
|
max_speaker_count=self._stt._opts.max_speaker_count,
|
|
)
|
|
recognition_config.diarization_config.CopyFrom(diarization_config)
|
|
|
|
return riva.client.StreamingRecognitionConfig(
|
|
config=recognition_config,
|
|
interim_results=True,
|
|
)
|
|
|
|
async def _collect_audio(self) -> None:
|
|
async for data in self._input_ch:
|
|
if isinstance(data, rtc.AudioFrame):
|
|
audio_bytes = data.data.tobytes()
|
|
if audio_bytes:
|
|
self._audio_queue.put(audio_bytes)
|
|
elif isinstance(data, self._FlushSentinel):
|
|
break
|
|
|
|
def _recognition_worker(self, config: riva.client.StreamingRecognitionConfig) -> None:
|
|
try:
|
|
audio_generator = self._audio_chunk_generator()
|
|
|
|
response_generator = self._asr_service.streaming_response_generator(
|
|
audio_generator, config
|
|
)
|
|
|
|
for response in response_generator:
|
|
self._handle_response(response)
|
|
|
|
except Exception:
|
|
logger.exception("Error in NVIDIA recognition thread")
|
|
finally:
|
|
self._event_loop.call_soon_threadsafe(self._done_fut.set_result, None)
|
|
|
|
def _audio_chunk_generator(self) -> Generator[bytes, None, None]:
|
|
"""
|
|
The nvidia riva SDK requires a generator for realtime STT - so we have to
|
|
wrap the
|
|
"""
|
|
while True:
|
|
audio_chunk = self._audio_queue.get()
|
|
|
|
if not audio_chunk:
|
|
break
|
|
|
|
yield audio_chunk
|
|
|
|
def _handle_response(self, response) -> None:
|
|
try:
|
|
if not hasattr(response, "results") or not response.results:
|
|
return
|
|
|
|
for result in response.results:
|
|
if not hasattr(result, "alternatives") or not result.alternatives:
|
|
continue
|
|
|
|
alternative = result.alternatives[0]
|
|
transcript = getattr(alternative, "transcript", "")
|
|
is_final = getattr(result, "is_final", False)
|
|
|
|
if not transcript.strip():
|
|
continue
|
|
|
|
self._request_id = f"nvidia-{id(response)}"
|
|
|
|
if not self._speaking and transcript.strip():
|
|
self._speaking = True
|
|
self._event_loop.call_soon_threadsafe(
|
|
self._event_ch.send_nowait,
|
|
stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH),
|
|
)
|
|
|
|
speech_data = self._convert_to_speech_data(alternative, is_final=is_final)
|
|
|
|
if is_final:
|
|
self._event_loop.call_soon_threadsafe(
|
|
self._event_ch.send_nowait,
|
|
stt.SpeechEvent(
|
|
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
|
|
request_id=self._request_id,
|
|
alternatives=[speech_data],
|
|
),
|
|
)
|
|
|
|
if self._speaking:
|
|
self._event_loop.call_soon_threadsafe(
|
|
self._event_ch.send_nowait,
|
|
stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH),
|
|
)
|
|
else:
|
|
self._event_loop.call_soon_threadsafe(
|
|
self._event_ch.send_nowait,
|
|
stt.SpeechEvent(
|
|
type=stt.SpeechEventType.INTERIM_TRANSCRIPT,
|
|
request_id=self._request_id,
|
|
alternatives=[speech_data],
|
|
),
|
|
)
|
|
|
|
except Exception:
|
|
logger.exception("Error handling response")
|
|
|
|
def _convert_to_speech_data(self, alternative, *, is_final: bool) -> stt.SpeechData:
|
|
transcript = getattr(alternative, "transcript", "")
|
|
confidence = getattr(alternative, "confidence", 0.0)
|
|
words = getattr(alternative, "words", [])
|
|
|
|
start_time = 0.0
|
|
end_time = 0.0
|
|
speaker_id: str | None = None
|
|
if words:
|
|
start_time = getattr(words[0], "start_time", 0) / 1000.0 + self.start_time_offset
|
|
end_time = getattr(words[-1], "end_time", 0) / 1000.0 + self.start_time_offset
|
|
|
|
if self._stt._opts.enable_diarization and is_final:
|
|
speaker_tags = [getattr(word, "speaker_tag", 0) for word in words]
|
|
if speaker_tags:
|
|
speaker = Counter(speaker_tags).most_common(1)[0][0]
|
|
speaker_id = f"S{speaker}"
|
|
|
|
return stt.SpeechData(
|
|
language=LanguageCode(self._language),
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
confidence=confidence,
|
|
text=transcript,
|
|
speaker_id=speaker_id,
|
|
words=[
|
|
TimedString(
|
|
text=getattr(word, "word", ""),
|
|
start_time=getattr(word, "start_time", 0) + self.start_time_offset,
|
|
end_time=getattr(word, "end_time", 0) + self.start_time_offset,
|
|
)
|
|
for word in words
|
|
]
|
|
if words
|
|
else None,
|
|
)
|