# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/nemotron3_diarization/modular_nemotron3_diarization.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_nemotron3_diarization.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig, SubConfigSpec from ...utils import auto_docstring from ...utils.type_validators import interval, positive_int_field @auto_docstring(checkpoint="nvidia/Nemotron-3-Diarization") @strict class Nemotron3DiarizationAudioConfig(PreTrainedConfig): r""" subsampling_factor (`int`, *optional*, defaults to 8): Number of consecutive spectrogram frames stacked into one encoder frame. The classifier upsamples its outputs by the same factor, so speaker activity is predicted at the spectrogram frame rate. """ model_type = "nemotron3_diarization_audio" hidden_size: int = 512 intermediate_size: int = 2048 num_hidden_layers: int = 31 num_attention_heads: int = 8 num_key_value_heads: int | None = None hidden_act: str = "gelu" max_position_embeddings: int = 5000 initializer_range: float = 0.02 rope_parameters: dict | None = None attention_dropout: float | int = 0.0 num_mel_bins: int = 128 base_config_key = "audio_config" subsampling_factor: int = 8 def __post_init__(self, **kwargs): kwargs.setdefault("partial_rotary_factor", 1.0) if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads kwargs.setdefault("partial_rotary_factor", 0.5) super().__post_init__(**kwargs) def validate_architecture(self): if self.hidden_size % self.num_attention_heads != 0: raise ValueError( f"`hidden_size` ({self.hidden_size}) must be divisible by `num_attention_heads` " f"({self.num_attention_heads})." ) @auto_docstring(checkpoint="nvidia/Nemotron-3-Diarization") @strict class Nemotron3DiarizationHeadConfig(PreTrainedConfig): r""" hidden_size (`int`, *optional*, defaults to 192): Hidden size of the speaker head: the encoder output is projected to it, and the upsampler and the classifier keep it. num_speakers (`int`, *optional*, defaults to 8): Maximum number of speakers, i.e. the number of per-frame activity outputs. Speakers are ordered by their first arrival in the audio. audio_hidden_size (`int`, *optional*, defaults to 512): Hidden size of the encoder output the head projects from. Must match `Nemotron3DiarizationAudioConfig.hidden_size`. subsampling_factor (`int`, *optional*, defaults to 8): Upsampling factor of the head, back to the spectrogram frame rate. Must match `Nemotron3DiarizationAudioConfig.subsampling_factor`. """ base_config_key = "head_config" hidden_size: int = positive_int_field(default=192) num_speakers: int = positive_int_field(default=8) audio_hidden_size: int = positive_int_field(default=512) subsampling_factor: int = positive_int_field(default=8) @auto_docstring(checkpoint="nvidia/Nemotron-3-Diarization") @strict class Nemotron3DiarizationStreamingConfig(PreTrainedConfig): r""" fifo_length (`int`, *optional*, defaults to 264): Capacity of the FIFO queue of the most recent encoder frames in streaming mode (offline mode uses `Nemotron3DiarizationConfig.fifo_length`). speaker_cache_update_period (`int`, *optional*, defaults to 222): Number of encoder frames moved from the FIFO queue to the speaker cache when the queue overflows, in streaming mode (offline mode uses `Nemotron3DiarizationConfig.speaker_cache_update_period`). speaker_cache_length (`int`, *optional*, defaults to 264): Capacity of the Arrival-Order Speaker Cache. Must be at least `(1 + speaker_cache_silence_frames_per_speaker) * num_speakers`. speaker_cache_silence_frames_per_speaker (`int`, *optional*, defaults to 1): Number of speaker-cache slots per speaker reserved for the learned silence embedding when the cache is compressed. prediction_score_threshold (`float`, *optional*, defaults to 0.25): Lower clamp of the speaker probabilities before taking their log in the speaker-cache frame scores. latest_frames_score_boost (`float`, *optional*, defaults to 0.05): Score bonus given to the frames newly added to the speaker cache when it is compressed. strong_boost_rate (`float`, *optional*, defaults to 0.75): Fraction of the per-speaker cache budget whose best frames get a strong score boost, so that every speaker keeps at least that many frames in the cache. weak_boost_rate (`float`, *optional*, defaults to 1.5): Fraction of the per-speaker cache budget whose best frames get a weak score boost, which prevents one speaker from dominating the cache. min_positive_scores_rate (`float`, *optional*, defaults to 0.5): Fraction of the per-speaker cache budget: a speaker with at least that many positively scored frames has its non-positive (overlapped speech) frames excluded from the cache. num_speakers (`int`, *optional*, defaults to 8): Number of speakers tracked by the speaker cache. Must match `Nemotron3DiarizationHeadConfig.num_speakers`. subsampling_factor (`int`, *optional*, defaults to 8): Number of speaker-probability frames per encoder frame. Must match `Nemotron3DiarizationAudioConfig.subsampling_factor`. """ fifo_length: int = positive_int_field(default=264) speaker_cache_update_period: int = positive_int_field(default=222) speaker_cache_length: int = positive_int_field(default=264) speaker_cache_silence_frames_per_speaker: int = interval(min=0)(default=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 = positive_int_field(default=8) subsampling_factor: int = positive_int_field(default=8) @auto_docstring(checkpoint="nvidia/Nemotron-3-Diarization") @strict class Nemotron3DiarizationConfig(PreTrainedConfig): r""" audio_config (`Nemotron3DiarizationAudioConfig` or `dict`, *optional*): Configuration of the transformer audio encoder. Defaults to `Nemotron3DiarizationAudioConfig()`. head_config (`Nemotron3DiarizationHeadConfig` or `dict`, *optional*): Configuration of the speaker head. Defaults to `Nemotron3DiarizationHeadConfig()`. streaming_config (`Nemotron3DiarizationStreamingConfig` or `dict`, *optional*): Speaker-cache policy, and the FIFO sizes of streaming mode. Defaults to `Nemotron3DiarizationStreamingConfig()`. chunk_length (`int`, *optional*, defaults to 340): Offline mode: number of encoder frames per chunk when a whole recording is diarized in one forward. In streaming mode the chunk is the input of each forward. chunk_right_context (`int`, *optional*, defaults to 40): Offline mode: number of look-ahead encoder frames each chunk takes from the following ones. In streaming mode the look-ahead is `num_lookahead_frames` of each forward. fifo_length (`int`, *optional*, defaults to 40): Offline mode: capacity of the FIFO queue of the most recent encoder frames. Streaming mode uses `streaming_config.fifo_length`. speaker_cache_update_period (`int`, *optional*, defaults to 300): Offline mode: number of encoder frames moved from the FIFO queue to the speaker cache when the queue overflows. Streaming mode uses `streaming_config.speaker_cache_update_period`. """ model_type = "nemotron3_diarization" sub_configs_defaults = { "audio_config": SubConfigSpec(config_class=Nemotron3DiarizationAudioConfig), "head_config": SubConfigSpec(config_class=Nemotron3DiarizationHeadConfig), "streaming_config": SubConfigSpec(config_class=Nemotron3DiarizationStreamingConfig), } audio_config: Nemotron3DiarizationAudioConfig | dict | None = None head_config: Nemotron3DiarizationHeadConfig | dict | None = None streaming_config: Nemotron3DiarizationStreamingConfig | dict | None = None chunk_length: int = 340 chunk_right_context: int = 40 fifo_length: int = 40 speaker_cache_update_period: int = 300 initializer_range: float = 0.02 def validate_architecture(self): if self.chunk_right_context >= self.chunk_length: raise ValueError( f"`chunk_right_context` ({self.chunk_right_context}) must be smaller than " f"`chunk_length` ({self.chunk_length})." ) if self.head_config.audio_hidden_size != self.audio_config.hidden_size: raise ValueError( f"`head_config.audio_hidden_size` ({self.head_config.audio_hidden_size}) must match " f"`audio_config.hidden_size` ({self.audio_config.hidden_size})." ) if self.head_config.subsampling_factor != self.audio_config.subsampling_factor: raise ValueError( f"`head_config.subsampling_factor` ({self.head_config.subsampling_factor}) must match " f"`audio_config.subsampling_factor` ({self.audio_config.subsampling_factor})." ) if self.streaming_config.num_speakers != self.head_config.num_speakers: raise ValueError( f"`streaming_config.num_speakers` ({self.streaming_config.num_speakers}) must match " f"`head_config.num_speakers` ({self.head_config.num_speakers})." ) if self.streaming_config.subsampling_factor != self.audio_config.subsampling_factor: raise ValueError( f"`streaming_config.subsampling_factor` ({self.streaming_config.subsampling_factor}) must match " f"`audio_config.subsampling_factor` ({self.audio_config.subsampling_factor})." ) silence_frames = self.streaming_config.speaker_cache_silence_frames_per_speaker min_speaker_cache_length = (1 + silence_frames) * self.streaming_config.num_speakers if self.streaming_config.speaker_cache_length < min_speaker_cache_length: raise ValueError( f"`streaming_config.speaker_cache_length` ({self.streaming_config.speaker_cache_length}) must be at " "least `(1 + streaming_config.speaker_cache_silence_frames_per_speaker) * " f"streaming_config.num_speakers` ({min_speaker_cache_length})." ) __all__ = [ "Nemotron3DiarizationAudioConfig", "Nemotron3DiarizationConfig", "Nemotron3DiarizationHeadConfig", "Nemotron3DiarizationStreamingConfig", ]