Transformers documentation
Nemotron 3 说话人分离
该模型于 2026-09-23 贡献至 Hugging Face Transformers。
Nemotron 3 说话人分离
概述
Nemotron 3 Diarization 是一款开放权重的流式说话人分离模型，旨在确定真实音频中“谁在何时说话”。它同时支持流式和离线推理，可处理多达八位说话人，并按每位说话人在输入音频中首次出现的时间顺序排列说话人输出。
该模型采用了为 Streaming Sortformer 1、2 引入的 Arrival-Order Speaker Cache (AOSC) 1 和 FIFO 队列。单个检查点支持可配置的延迟配置，从 80 ms 的输入缓冲区到 30.4 s 的离线式缓冲区，并支持以 10 ms 为倍数进行配置的输出帧分辨率。通过分块推理，音频的最大时长不受限制。
用法
离线
import torch
from transformers import AutoModelForAudioFrameClassification, AutoProcessor
from transformers.audio_utils import load_audio
model_id = "nvidia/Nemotron-3-Diarization"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForAudioFrameClassification.from_pretrained(model_id, device_map="auto")
sampling_rate = processor.feature_extractor.sampling_rate
audio = load_audio(
"https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/diarization_example.mp3",
sampling_rate=sampling_rate,
)
inputs = processor(audio, sampling_rate=sampling_rate).to(model.device, dtype=model.dtype)
with torch.inference_mode():
logits = model(**inputs).logits # (1, num_frames, 8), one frame every 10 ms
segments = processor.extract_speaker_dict(logits, inputs.attention_mask)[0]
for segment in segments:
print(f"speaker_{segment['Speaker']}: {segment['Start']:.2f}s - {segment['End']:.2f}s")流式
音频以分块（chunk）的形式到达，每次前向传播处理一个分块：处理器会根据其 streaming_mode 对音频进行切分，并添加 num_lookahead_frames 个尾部前瞻帧。模型会关注这些前瞻帧但不对它们进行评分，因为它们属于下一个分块的开头。前向传播会返回 speaker_cache，以便传递给下一次调用。会话的最后一个分块通过设置 is_last_audio_chunk=True 来提取：由于没有前瞻帧，因此会对所有剩余帧进行评分。
streaming_mode | 延迟¹ |
|---|---|
"low_latency"（默认） | 1.04 s |
"very_low_latency" | 0.64 s |
"ultra_low_latency" | 0.32 s |
¹ 模型处理某个数据块之前需要等待的音频时长：即该数据块及其前瞻部分，不包括计算时间。
import torch
from transformers import AutoModelForAudioFrameClassification, AutoProcessor
from transformers.audio_utils import load_audio
model_id = "nvidia/Nemotron-3-Diarization"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForAudioFrameClassification.from_pretrained(model_id, device_map="auto")
processor.set_streaming_mode("low_latency") # the default, can also be "very_low_latency" and "ultra_low_latency"
print(f"Streaming latency: {processor.streaming_latency_ms} ms")
sampling_rate = processor.feature_extractor.sampling_rate
audio = load_audio(
"https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/diarization_example.mp3",
sampling_rate=sampling_rate,
)
def inputs_generator():
"""Yields the processor outputs of each chunk."""
yield processor(
audio[: processor.num_samples_first_audio_chunk],
sampling_rate=sampling_rate,
is_streaming=True,
is_first_audio_chunk=True,
)
mel_frame_idx = processor.num_mel_frames_per_step
start_idx = processor.audio_chunk_start(mel_frame_idx)
while (end_idx := start_idx + processor.num_samples_per_audio_chunk) <= audio.shape[0]:
yield processor(
audio[start_idx:end_idx],
sampling_rate=sampling_rate,
is_streaming=True,
is_first_audio_chunk=False,
)
mel_frame_idx += processor.num_mel_frames_per_step
start_idx = processor.audio_chunk_start(mel_frame_idx)
# the audio ended: the frames left in the buffer are the last ones of the session
yield processor(
audio[start_idx:],
sampling_rate=sampling_rate,
is_streaming=True,
is_first_audio_chunk=False,
is_last_audio_chunk=True,
)
speaker_cache, logits = None, []
with torch.inference_mode():
for inputs in inputs_generator():
inputs = inputs.to(model.device, dtype=model.dtype)
# `inputs` carries `num_lookahead_frames` for every chunk but the last, `speaker_cache` links the chunks
outputs = model(**inputs, speaker_cache=speaker_cache)
logits.append(outputs.logits) # the chunk's frames, without its look-ahead
speaker_cache = outputs.speaker_cache
logits = torch.cat(logits, dim=1) # (1, num_frames, 8), one frame every 10 ms
segments = processor.extract_speaker_dict(logits)[0] # [{"Start": 0.0, "End": 15.43, "Speaker": 0}, ...]让它飞速运转
流式步骤的编码器输入为 [speaker cache | FIFO | chunk]，其长度会随着缓存和 FIFO 的填充与缩减而变化：torch.compile 会在每个会话中重新编译约一百次。将每个步骤填充到该模式的最大窗口长度即可固定形状。每个 chunk 的位置都从零开始，因此右侧填充不会改变有效帧：
import torch.nn.functional as F
chunk_length, chunk_right_context = processor.streaming_modes[processor.streaming_mode]
max_window = (
model.config.streaming_config.speaker_cache_length
+ model.config.streaming_config.fifo_length
+ chunk_length
+ chunk_right_context
)
encoder = model.model
compiled_forward = torch.compile(encoder.forward, mode="reduce-overhead", fullgraph=True, dynamic=False)
def padded_forward(inputs_embeds, attention_mask=None, position_ids=None, **kwargs):
batch_size, num_frames, _ = inputs_embeds.shape
if attention_mask is None:
attention_mask = inputs_embeds.new_ones(batch_size, num_frames, dtype=torch.bool)
padding = max_window - num_frames
hidden_states = compiled_forward(
inputs_embeds=F.pad(inputs_embeds, (0, 0, 0, padding)),
attention_mask=F.pad(attention_mask.bool(), (0, padding), value=False),
position_ids=torch.arange(max_window, device=inputs_embeds.device)[None, :],
**kwargs,
)
return hidden_states[:, :num_frames].clone() # CUDA graphs reuse the output buffer
encoder.forward = padded_forward
# warm up before the session: compiles, then records the CUDA graph, so the first real chunk runs at full speed
with torch.inference_mode():
for _ in range(3):
hidden_size = model.config.audio_config.hidden_size
padded_forward(torch.zeros(1, max_window, hidden_size, device=model.device, dtype=model.dtype))上述流式循环只会编译一次。离线前向传播采用相同的分块方式，因此可以使用 max_window 中的 config.fifo_length、config.chunk_left_context、config.chunk_length 和 config.chunk_right_context 应用相同的封装器。
| 与 eager 模式相比的加速比（A100，批处理大小为 1） | float32 | bfloat16 |
|---|---|---|
| 流式，逐步 | 1.2x | 4.4x |
| 离线，488 秒录音 | 1.3x | 2.8x |
Nemotron3DiarizationConfig
class transformers.Nemotron3DiarizationConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Noneaudio_config: transformers.models.nemotron3_diarization.configuration_nemotron3_diarization.Nemotron3DiarizationAudioConfig | dict | None = Nonehead_config: transformers.models.nemotron3_diarization.configuration_nemotron3_diarization.Nemotron3DiarizationHeadConfig | dict | None = Nonestreaming_config: transformers.models.nemotron3_diarization.configuration_nemotron3_diarization.Nemotron3DiarizationStreamingConfig | dict | None = Nonechunk_length: int = 340chunk_right_context: int = 40chunk_left_context: int = 0fifo_length: int = 40speaker_cache_update_period: int = 300initializer_range: float = 0.02 )
参数
- audio_config (
Nemotron3DiarizationAudioConfig或dict，可选) — Transformer 音频编码器的配置。默认为Nemotron3DiarizationAudioConfig()。 - head_config (
Nemotron3DiarizationHeadConfig或dict，可选) — 说话人头的配置。默认为Nemotron3DiarizationHeadConfig()。 - streaming_config（
Nemotron3DiarizationStreamingConfig或dict，可选）— 说话人缓存策略以及流式模式的 FIFO 大小。默认值为Nemotron3DiarizationStreamingConfig()。 - chunk_length (
int, 可选，默认为 340) — 离线模式：当整段录音在一次前向传播中进行说话人分离时，每个块包含的编码器帧数。在流式模式下，该块是每次前向传播的输入。 - chunk_right_context (
int，可选，默认值为 40) — 离线模式：每个 chunk 从后续 chunk 中获取的前瞻编码器帧数。在流式模式下，前瞻量为每次前向传播的num_lookahead_frames。 - chunk_left_context (
int，可选，默认为 0) — 离线模式：每个 chunk 从前序 chunk 中再次取用的回看编码器帧数，这些帧会被插入到缓存帧与当前 chunk 之间。与回看帧类似，它们会被关注，但既不会被评分，也不会被推入 FIFO 队列。在流式模式下，回看帧数为每次前向传播的num_lookback_frames。 - fifo_length (
int，可选，默认为 40) — 离线模式：最近编码器帧的 FIFO 队列容量。流式模式使用streaming_config.fifo_length。 - speaker_cache_update_period (
int，可选，默认值为 300) — 离线模式：当 FIFO 队列溢出时，从该队列移至说话人缓存的编码器帧数。流式模式使用streaming_config.speaker_cache_update_period。 - initializer_range (
float，可选，默认为0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
这是用于存储 Nemotron3DiarizationModel 配置的类。它根据指定的参数实例化 Nemotron3 Diarization 模型，并定义模型架构。使用默认值实例化配置将得到与 nvidia/Nemotron-3-Diarization 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。
Nemotron3DiarizationAudioConfig
class transformers.Nemotron3DiarizationAudioConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonehidden_size: int = 512intermediate_size: int = 2048num_hidden_layers: int = 31num_attention_heads: int = 8num_key_value_heads: int | None = Nonehidden_act: str = 'gelu'max_position_embeddings: int = 5000initializer_range: float = 0.02rope_parameters: dict | None = Noneattention_dropout: float | int = 0.0num_mel_bins: int = 128subsampling_factor: int = 8 )
参数
- hidden_size (
int，可选，默认为512) — 隐藏表示的维度。 - intermediate_size (
int，可选，默认为2048) — MLP 表示的维度。 - num_hidden_layers (
int，可选，默认为31) — Transformer 解码器中的隐藏层数量。 - num_attention_heads (
int，可选，默认为8) — Transformer 解码器中每个注意力层的注意力头数量。 - num_key_value_heads (
int, 可选) — 这是用于实现 Grouped Query Attention 的 key_value 头数。如果num_key_value_heads=num_attention_heads，模型将使用 Multi Head Attention (MHA)；如果num_key_value_heads=1，模型将使用 Multi Query Attention (MQA)；否则使用 GQA。在将多头检查点转换为 GQA 检查点时，每个组的 key 和 value 头应通过对该组内所有原始头进行平均池化来构建。更多详情请参阅 这篇论文。如果未指定，则默认为num_attention_heads。 - hidden_act (
str，可选，默认为gelu) — 解码器中的非线性激活函数（函数或字符串）。例如，"gelu"、"relu"、"silu"等。 - max_position_embeddings (
int，可选，默认为5000) — 该模型可能使用的最大序列长度。 - initializer_range (
float，可选，默认为0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。 - rope_parameters (
dict, 可选) — 包含 RoPE 嵌入配置参数的字典。该字典应包含rope_theta的值，并且如果希望使用 RoPE 处理更长的max_position_embeddings，则可选地包含用于缩放的参数。 - attention_dropout (
Union[float, int]，可选，默认为0.0) — 注意力概率的 dropout 比率。 - num_mel_bins (
int, 可选，默认为128) — 每个输入帧使用的 mel 特征数量。应与AutoFeatureExtractor类中使用的值相对应。 - subsampling_factor (
int，可选，默认为 8) — 堆叠到单个编码器帧中的连续频谱帧数。分类器会以相同的因子对其输出进行上采样，因此说话人活动预测的速率与频谱帧率一致。
这是用于存储 Nemotron3DiarizationModel 配置的类。它根据指定的参数实例化 Nemotron3 Diarization 模型，并定义模型架构。使用默认值实例化配置将得到与 nvidia/Nemotron-3-Diarization 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。
Nemotron3DiarizationHeadConfig
class transformers.Nemotron3DiarizationHeadConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonehidden_size: int = 192num_speakers: int = 8audio_hidden_size: int = 512subsampling_factor: int = 8 )
参数
- hidden_size (
int, 可选，默认为 192) — 说话人头的隐藏层大小：编码器输出会投影到该维度，且上采样器和分类器会保留该维度。 - num_speakers (
int, 可选，默认为 8) — 说话人的最大数量，即每帧活动输出的数量。说话人按其在音频中首次出现的时间顺序排列。 - audio_hidden_size (
int，可选，默认值为 512) — 头部所投影的编码器输出的隐藏层大小。必须与Nemotron3DiarizationAudioConfig.hidden_size相匹配。 - subsampling_factor (
int，可选，默认为 8) — 头部的上采样因子，用于恢复到频谱图的帧率。必须与Nemotron3DiarizationAudioConfig.subsampling_factor匹配。
这是用于存储 Nemotron3DiarizationModel 配置的类。它根据指定的参数实例化 Nemotron3 Diarization 模型，并定义模型架构。使用默认值实例化配置将得到与 nvidia/Nemotron-3-Diarization 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。
Nemotron3DiarizationStreamingConfig
class transformers.Nemotron3DiarizationStreamingConfig
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonefifo_length: int = 264speaker_cache_update_period: int = 222speaker_cache_length: int = 264speaker_cache_silence_frames_per_speaker: int = 1prediction_score_threshold: float = 0.25latest_frames_score_boost: float = 0.05strong_boost_rate: float = 0.75weak_boost_rate: float = 1.5min_positive_scores_rate: float = 0.5num_speakers: int = 8subsampling_factor: int = 8 )
参数
- fifo_length (
int，可选，默认为 264) — 流式模式下最近编码器帧的 FIFO 队列容量（离线模式使用Nemotron3DiarizationConfig.fifo_length）。 - speaker_cache_update_period (
int，可选，默认值为 222) — 在流式模式下，当 FIFO 队列溢出时，从该队列移动到说话人缓存中的编码器帧数（离线模式使用Nemotron3DiarizationConfig.speaker_cache_update_period）。 - speaker_cache_length (
int, 可选，默认为 264) — 到达顺序说话人缓存的容量。必须至少为(1 + speaker_cache_silence_frames_per_speaker) * num_speakers。 - speaker_cache_silence_frames_per_speaker (
int，可选，默认为 1) — 当缓存被压缩时，为每个说话人保留的用于学习静音嵌入的说话人缓存槽位数量。 - prediction_score_threshold (
float，可选，默认值为 0.25) — 在对说话人缓存帧分数取对数之前，对说话人概率施加的下限截断值。 - latest_frames_score_boost（
float，可选，默认值为 0.05）—— 当说话人缓存被压缩时，新添加到缓存中的帧所获得的分数加成。 - strong_boost_rate (
float，可选，默认值为 0.75) — 每个说话人缓存预算中，获得强分数加成的最佳帧所占的比例，以确保每个说话人在缓存中至少保留相应数量的帧。 - weak_boost_rate (
float，可选，默认值为 1.5) — 每个说话人缓存预算中，获得弱分数加权的最佳帧所占的比例。该加权机制可防止某个说话人独占缓存。 - min_positive_scores_rate (
float，可选，默认值为 0.5) — 每个说话人缓存预算的比例：如果某个说话人拥有至少该数量的正分帧，则其非正分（重叠语音）帧将被排除在缓存之外。 - num_speakers (
int，可选，默认为 8) — 说话人缓存所跟踪的说话人数量。必须与Nemotron3DiarizationHeadConfig.num_speakers一致。 - subsampling_factor (
int，可选，默认为 8) — 每个编码器帧对应的说话人概率帧数。必须与Nemotron3DiarizationAudioConfig.subsampling_factor一致。
这是用于存储 Nemotron3DiarizationModel 配置的类。它根据指定的参数实例化 Nemotron3 Diarization 模型，并定义模型架构。使用默认值实例化配置将得到与 nvidia/Nemotron-3-Diarization 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请参阅 PreTrainedConfig 的文档。
Nemotron3DiarizationAudioModel
class transformers.Nemotron3DiarizationAudioModel
< source >( config: Nemotron3DiarizationAudioConfig )
参数
- config (Nemotron3DiarizationAudioConfig) — 模型配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
未添加任何特定头部的裸 Nemotron3 Diarization Model，直接输出原始 hidden-states。
该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。
该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。
forward
< source >( input_features: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutput or tuple(torch.Tensor)
参数
- input_features (
torch.Tensor，形状为(batch_size, sequence_length, feature_dim)，可选) — 与输入音频特征对应的张量。可以使用 NemotronAsrStreamingFeatureExtractor 获取音频特征。有关详细信息，请参阅NemotronAsrStreamingFeatureExtractor.__call__()（Nemotron3DiarizationProcessor 使用 NemotronAsrStreamingFeatureExtractor 处理音频）。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
- inputs_embeds（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入input_ids。如果您希望对input_ids索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。 - position_ids (
torch.Tensor形状为(batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为[0, config.n_positions - 1].
返回值
BaseModelOutput 或 tuple(torch.Tensor)
一个 BaseModelOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（Nemotron3DiarizationConfig）和输入。
Nemotron3DiarizationAudioModel 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
last_hidden_state（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor）— 模型最后一层输出的隐藏状态序列。hidden_states (
tuple(torch.Tensor)，可选，当传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为(batch_size, sequence_length, hidden_size)。模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (
tuple(torch.Tensor)，可选，当传入output_attentions=True或config.output_attentions=True时返回) — 一个torch.Tensor元组（每层一个），形状为(batch_size, num_heads, sequence_length, sequence_length)。注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。
Nemotron3DiarizationModel
class transformers.Nemotron3DiarizationModel
< source >( config: Nemotron3DiarizationConfig )
参数
- config (Nemotron3DiarizationConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
不含说话人分类器的 Nemotron3Diarization 模型：对单个块（及其缓存帧）进行编码，并将编码器帧上采样回频谱图帧率。该模型不持有任何流式状态，分块前向传播请参见 Nemotron3DiarizationForAudioFrameClassification。
该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。
该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。
forward
< source >( input_features: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.Tensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutput or tuple(torch.Tensor)
参数
- input_features (
torch.Tensor，形状为(batch_size, sequence_length, feature_dim)，可选) — 与输入音频特征对应的张量。可以使用 NemotronAsrStreamingFeatureExtractor 获取音频特征。有关详细信息，请参阅NemotronAsrStreamingFeatureExtractor.__call__()（Nemotron3DiarizationProcessor 使用 NemotronAsrStreamingFeatureExtractor 处理音频）。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
- inputs_embeds（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入input_ids。如果您希望对input_ids索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权，这将非常有用。 - position_ids (
torch.Tensor形状为(batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为[0, config.n_positions - 1].
返回值
BaseModelOutput 或 tuple(torch.Tensor)
一个 BaseModelOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（Nemotron3DiarizationConfig）和输入。
Nemotron3DiarizationModel 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
last_hidden_state（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor）— 模型最后一层输出的隐藏状态序列。hidden_states (
tuple(torch.Tensor)，可选，当传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组（如果模型具有嵌入层，则包含嵌入层的输出，以及每一层的输出），形状为(batch_size, sequence_length, hidden_size)。模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
attentions (
tuple(torch.Tensor)，可选，当传入output_attentions=True或config.output_attentions=True时返回) — 一个torch.Tensor元组（每层一个），形状为(batch_size, num_heads, sequence_length, sequence_length)。注意力权重，即经过注意力 softmax 之后的权重，用于在自注意力头中计算加权平均值。
Nemotron3DiarizationProcessor
class transformers.Nemotron3DiarizationProcessor
< source >( feature_extractorsubsampling_factor = 8streaming_modes = Nonestreaming_mode = 'low_latency' )
参数
- feature_extractor（
NemotronAsrStreamingFeatureExtractor）— 特征提取器是必需的输入。 - subsampling_factor (
int，可选，默认为 8) — 每个编码器帧对应的 mel 帧数，与Nemotron3DiarizationAudioConfig.subsampling_factor保持一致。 - streaming_modes (
dict[str, tuple[int, int]]，可选) — 检查点支持的流式模式，将名称映射为编码器帧中的(chunk_length, chunk_right_context)。 处理器是该集合的唯一可信来源：set_streaming_mode()会对其进行验证。默认值为模型卡中定义的模式："low_latency"(9, 4)、"very_low_latency"(6, 2) 和"ultra_low_latency"(3, 1)。 - streaming_mode (
str，可选，默认为"low_latency") — 会话的流式模式，取值为streaming_modes之一；可通过set_streaming_mode()进行修改。离线使用时会忽略此参数。
构建一个 Nemotron3DiarizationProcessor，它将特征提取器封装为单个处理器。
Nemotron3DiarizationProcessor 提供了 NemotronAsrStreamingFeatureExtractor 的所有功能。有关更多信息，请参阅 ~NemotronAsrStreamingFeatureExtractor。
__call__
< source >( audio: typing.Union[numpy.ndarray, ForwardRef('torch.Tensor'), collections.abc.Sequence[numpy.ndarray], collections.abc.Sequence['torch.Tensor']]sampling_rate: int | None = Noneis_streaming: bool = Falseis_first_audio_chunk: bool = Trueis_last_audio_chunk: bool = False**kwargs: *<class 'transformers.models.nemotron3_diarization.processing_nemotron3_diarization.Nemotron3DiarizationProcessorKwargs'> ) → BatchFeature
参数
- audio (
Union[numpy.ndarray, torch.Tensor, collections.abc.Sequence[numpy.ndarray], collections.abc.Sequence[torch.Tensor]]) — 要准备处理的音频或音频批次。每个音频可以是 NumPy 数组或 PyTorch 张量。 如果为 NumPy 数组/PyTorch 张量，每个音频的形状应为 (C, T)，其中 C 为通道数，T 为音频的采样长度。 - sampling_rate (
int, 可选) — 输入音频的采样率，单位为 Hz。如果提供该值，将与特征提取器预期的采样率（16000 Hz）进行校验。 - is_streaming (
bool，可选，默认为False) — 指示音频是否为流式会话中的一个数据块，is_first_audio_chunk和is_last_audio_chunk用于区分首块、末块与其他块。数据块的大小由streaming_mode指定，可通过set_streaming_mode()进行修改。除最后一个数据块外，第一个数据块必须恰好包含num_samples_first_audio_chunk个音频采样点，后续数据块必须恰好包含num_samples_per_audio_chunk个采样点。 - is_first_audio_chunk (
bool，可选，默认为True) — 指示这是否为流式会话的第一个数据块。对于第一个数据块以及离线使用场景，特征提取器会对分析窗口进行居中处理（center=True）；而对于后续数据块则不进行居中处理（center=False），从而确保每个数据块的频谱图能够逐帧复现单次完整语音处理的结果。当is_streaming=False时，此参数必须为True。 - is_last_audio_chunk (
bool，可选，默认为False) — 指示该块是否结束流式会话。会话中的每个块都以chunk_right_context个前瞻编码器帧结尾，这些帧仅在下一步进行评分，并且是下一个块的开头。最后一个块没有下一步，因此其所有帧都会被评分，无论数量多少，并且其末尾会进行零填充，就像完整语句处理中对音频末尾进行填充一样，从而输出该语句的最后几帧。当is_streaming=False时，必须为False。 - return_tensors (
str或 TensorType, 可选) — 如果设置，将返回特定框架的张量。可接受的值为：'pt'：返回 PyTorchtorch.Tensor对象。'np'：返回 NumPynp.ndarray对象。
- **kwargs (ProcessingKwargs，可选) — 每种模态（文本、图像、视频、音频）的额外处理选项。模型特定参数已在上方列出；有关支持参数的完整列表，请参阅 TypedDict 类。
返回值
特征提取器的输出为 input_features 和 attention_mask。在流式模式下，分析窗口延伸到最后一个块之外的尾部帧会被丢弃，因此 input_features 恰好包含该块的帧，可以直接传递给模型，并且除最后一个块外的每个块都携带 num_lookahead_frames，即其尾部前瞻编码器帧的数量，从而将模型置于流式模式。
Nemotron3DiarizationSpeakerCache
class transformers.Nemotron3DiarizationSpeakerCache
< source >( config: Nemotron3DiarizationStreamingConfigfifo_length: int | None = Nonespeaker_cache_update_period: int | None = None )
Nemotron3DiarizationForAudioFrameClassification 的流式状态：按到达顺序排列的说话人缓存，以及每个 chunk 都会关注到的最近编码器帧的 FIFO 队列。
更新
< source >( chunk_input_embeds: Tensorchunk_logits: Tensorsilence_embeds: Tensornum_chunk_frames: intmask: typing.Optional[torch.Tensor] = Nonenum_lookback_frames: int = 0 )
参数
- chunk_input_embeds（形状为
(batch_size, num_input_frames, hidden_size)的torch.Tensor）— 该步骤的编码器输入：由get_embeds返回的缓存帧、回看帧、当前块及其前看帧。 - chunk_logits（形状为
(batch_size, num_input_frames * subsampling_factor, num_speakers)的torch.Tensor）— 该步骤的说话人 logits，用于在压缩说话人缓存时对帧进行评分。 - silence_embeds（形状为
(hidden_size,)的torch.Tensor）— 学习得到的静音嵌入，用于填充压缩缓存中预留的静音槽位。 - num_chunk_frames (
int) —chunk_input_embeds中位于缓存帧和回看帧之后的分块帧数量。只有这些帧会加入 FIFO 队列：其后的前瞻帧将在下一步再次输入。 - mask (
torch.Tensor，形状为(batch_size, num_input_frames)，可选) —chunk_input_embeds中的有效帧，其填充帧的说话人概率被设为零。 - num_lookback_frames (
int, 可选，默认为 0) —chunk_input_embeds中缓存帧与当前块之间的回看帧数：即上一块中再次输入、且已随该块一起推送的帧。
将处理后的数据块推入 FIFO 队列，当队列溢出时，将其最旧的帧移至扬声器缓存。
Nemotron3DiarizationOutput
class transformers.Nemotron3DiarizationOutput
< source >( logits: typing.Optional[torch.Tensor] = Nonehidden_states: tuple[torch.Tensor, ...] | None = Noneattentions: tuple[torch.Tensor, ...] | None = Nonespeaker_cache: transformers.models.nemotron3_diarization.modeling_nemotron3_diarization.Nemotron3DiarizationSpeakerCache | None = None )
参数
- logits (
torch.Tensor，形状为(batch_size, num_frames, config.head_config.num_speakers)) — 以频谱帧率为单位的逐帧说话人活动 logits。logits.sigmoid()给出每个说话人在每一帧处于活动状态的概率；说话人按其在音频中首次出现的时间顺序排列。 - hidden_states (
tuple[torch.Tensor, ...]，可选，当output_hidden_states=True时返回) — 每个 chunk 的编码器隐藏状态，按 chunk 顺序排列：编码器对每个 chunk 运行一次，因此元组中每个 chunk 包含config.audio_config.num_hidden_layers + 1个张量。它们的序列长度为该 chunk 的长度，包含缓存和前瞻帧。 - attentions (
tuple[torch.Tensor, ...]，可选，当output_attentions=True时返回) — 每个块的编码器注意力权重，按块顺序排列，每个块包含config.audio_config.num_hidden_layers个张量。 - speaker_cache (
Nemotron3DiarizationSpeakerCache，可选，在流式模式下返回) — 更新后的流式状态，用于传递给同一音频流的下一个音频块的 forward。
Nemotron3DiarizationForAudioFrameClassification
class transformers.Nemotron3DiarizationForAudioFrameClassification
< source >( config: Nemotron3DiarizationConfig )
参数
- config (Nemotron3DiarizationConfig) — 模型的配置类，包含模型的所有参数。使用配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
Streaming Sortformer 说话人分离模型：针对每个频谱图帧，预测按首次出现顺序排列的至多 config.head_config.num_speakers 位说话人的活动状态。音频按块（chunk）处理，每个块会关注若干前瞻帧，以及保存在 Nemotron3DiarizationSpeakerCache 中的 Arrival-Order Speaker Cache 和 FIFO 队列。整段录音由 forward 本身进行分块（离线模式）；流式输入则每次 forward 送入一个块（流式模式）。
该模型继承自 PreTrainedModel。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。
该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。
forward
< source >( input_features: Tensorattention_mask: typing.Optional[torch.Tensor] = Nonespeaker_cache: transformers.models.nemotron3_diarization.modeling_nemotron3_diarization.Nemotron3DiarizationSpeakerCache | None = Nonenum_lookahead_frames: int | None = Nonenum_lookback_frames: int | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → Nemotron3DiarizationOutput or tuple(torch.Tensor)
参数
- input_features (
torch.Tensor，形状为(batch_size, sequence_length, feature_dim)) — 与输入音频特征对应的张量。可以使用 NemotronAsrStreamingFeatureExtractor 获取音频特征。有关详细信息，请参阅NemotronAsrStreamingFeatureExtractor.__call__()（Nemotron3DiarizationProcessor 使用 NemotronAsrStreamingFeatureExtractor 来处理音频）。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ，可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
- speaker_cache (
Nemotron3DiarizationSpeakerCache，可选) — 由同一音频流的前一个数据块的 forward 返回的流式状态。 - num_lookahead_frames (
int, 可选) — 流式模式：输入中仅用于前瞻的尾部编码器帧数量。这些帧会被关注，但其 logits 不会返回，也不会加入 FIFO 队列，因为它们用于开启下一个块。 Nemotron3DiarizationProcessor 会为会话中除最后一个块以外的每个块设置此值。这些参数与
num_lookback_frames共同决定运行模式。流式模式，每次前向处理一个数据块： 提供了num_lookahead_frames或num_lookback_frames（第一个数据块会创建speaker_cache，后续数据块会接收该缓存），或者仅提供了speaker_cache（会话的最后一个数据块，无前瞻）。输入减去其回看帧和前瞻帧后即为一个数据块，无论其长度如何，都会作为一个整体推入由config.streaming_config确定大小的 FIFO 队列。离线模式，两者均未提供：输入为完整录音，前向过程将其分割为config.chunk_length个编码器帧的数据块，每个数据块最多从前一个数据块获取config.chunk_left_context个回看帧，从下一个数据块获取config.chunk_right_context个前瞻帧，并使用由config.fifo_length确定大小的 FIFO 队列；不返回缓存。 - num_lookback_frames (
int, 可选) — 流式模式：输入中仅用于回看的编码器前导帧数量，即在speaker_cache的帧之上，再次提供的前一个块的末尾部分。这些帧会被关注，但不会返回其 logits，也不会再次加入 FIFO 队列。默认值为 0。
返回值
Nemotron3DiarizationOutput 或 tuple(torch.Tensor)
一个 Nemotron3DiarizationOutput 或 torch.Tensor 元组（如果传入了 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（Nemotron3DiarizationConfig）和输入。
Nemotron3DiarizationForAudioFrameClassification 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
- logits（形状为
(batch_size, num_frames, config.head_config.num_speakers)的torch.Tensor）— 以频谱图帧率为单位的逐帧说话人活动 logits。logits.sigmoid()给出每个说话人在每一帧处于活动状态的概率；说话人按其在音频中首次出现的时间顺序排列。 - hidden_states (
tuple[torch.Tensor, ...]，可选，当output_hidden_states=True时返回) — 按块顺序排列的每个块的编码器隐藏状态：编码器对每个块运行一次，因此元组中每个块包含config.audio_config.num_hidden_layers + 1个张量。它们的序列长度是该块的长度，包含缓存和前瞻帧。 - attentions (
tuple[torch.Tensor, ...]，可选，当output_attentions=True时返回) — 每个 chunk 的编码器注意力权重，按 chunk 顺序排列，每个 chunk 包含config.audio_config.num_hidden_layers个张量。 - speaker_cache（
Nemotron3DiarizationSpeakerCache，可选，在流式模式下返回）— 更新后的流式状态，用于传递给同一音频流的下一个音频块的 forward。
示例：
>>> from transformers import AutoModelForAudioFrameClassification, AutoProcessor
>>> from transformers.audio_utils import load_audio
>>> model_id = "nvidia/Nemotron-3-Diarization"
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> model = AutoModelForAudioFrameClassification.from_pretrained(model_id, device_map="auto")
>>> sampling_rate = processor.feature_extractor.sampling_rate
>>> audio = load_audio(
... "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/en-Alice_woman.wav",
... sampling_rate=sampling_rate,
... )
>>> inputs = processor(audio, sampling_rate=sampling_rate).to(model.device, dtype=model.dtype)
>>> probabilities = model(**inputs).logits.sigmoid() # (1, num_frames, 8), one frame every 10 ms