nemo_voice_agent.pipecat.services.nemo.streaming_diar

Module Contents

Classes

Name	Description
`DiarizationConfig`	Diarization configuration parameters for inference.
`NeMoStreamingDiarService`	-

API

class nemo_voice_agent.pipecat.services.nemo.streaming_diar.DiarizationConfig(
    model_path: str = 'nvidia/diar_streaming_sort...,
    device: str = 'cuda',
    log: bool = False,
    max_num_speakers: int = 4,
    spkcache_len: int = 188,
    spkcache_refresh_rate: int = 144,
    fifo_len: int = 188,
    chunk_len: int = 6,
    chunk_left_context: int = 1,
    chunk_right_context: int = 7
)

Dataclass

Diarization configuration parameters for inference.

chunk_left_context

int = 1

chunk_len

int = 6

chunk_right_context

int = 7

device

str = 'cuda'

fifo_len

int = 188

log

bool = False

max_num_speakers

int = 4

model_path

str = 'nvidia/diar_streaming_sortformer_4spk-v2'

spkcache_len

int = 188

spkcache_refresh_rate

int = 144

class nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService(
    cfg: nemo_voice_agent.pipecat.services.nemo.streaming_diar.DiarizationConfig,
    model: str,
    frame_len_in_secs: float = 0.08,
    sample_rate: int = 16000,
    left_offset: int = 8,
    right_offset: int = 8,
    use_amp: bool = False,
    compute_dtype: torch.dtype = torch.float32
)

buffer_size_in_secs

chunk_size

= self.cfg.chunk_len

device

= cfg.device

diarizer

= self.build_diarizer()

feature_bufferer

max_num_speakers

= self.cfg.max_num_speakers

streaming_state

= self.init_streaming_state(batch_size=1)

total_preds

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.build_diarizer()

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.diarize(
    audio: bytes,
    stream_id: str = 'default'
) -> str

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.init_streaming_state(
    batch_size: int = 1
) -> nemo.collections.asr.modules.sortformer_modules.StreamingSortformerState

Initialize the streaming state for the diarization model.

Parameters:

batch_size

int" default="1

The batch size to use.

Returns: StreamingSortformerState

The initialized streaming state.

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.print_diar_result(
    diar_result: numpy.ndarray
)

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.reset_state(
    stream_id: str = 'default'
)

nemo_voice_agent.pipecat.services.nemo.streaming_diar.NeMoStreamingDiarService.stream_step(
    processed_signal: torch.Tensor,
    processed_signal_length: torch.Tensor,
    streaming_state: nemo.collections.asr.modules.sortformer_modules.StreamingSortformerState,
    total_preds: torch.Tensor,
    left_offset: int = 0,
    right_offset: int = 0
) -> typing.Tuple[nemo.collections.asr.modules.sortformer_modules.StreamingSortformerState, torch.Tensor]

Execute a single streaming step for diarization.

Parameters:

processed_signal

Tensor

The processed audio signal.

processed_signal_length

Tensor

The length of the processed signal.

streaming_state

StreamingSortformerState

The current streaming state.

total_preds

Tensor

The total predictions so far.

left_offset

int" default="0

The left offset for the current chunk.

right_offset

int" default="0

The right offset for the current chunk.

Returns: Tuple[StreamingSortformerState, Tensor]

Tuple[SortformerStreamingState, Tensor]: The updated streaming state and predictions.