Add Transformers weights, config and usage section

README.md

@@ -185,6 +185,126 @@ python ${NEMO_ROOT}/examples/speaker_tasks/diarization/neural_diarizer/e2e_diari
 
 More details on the evaluation can be found in the [Diarization Evaluation](diarization_evaluation.md) subcard.
 
+### 🤗 Transformers usage
+
+This model is supported natively in [🤗 Transformers](https://github.com/huggingface/transformers)!
+Install from source:
+
+```bash
+pip install git+https://github.com/huggingface/transformers
+```
+
+For more details about usage, please refer to the [Transformers documentation](https://huggingface.co/docs/transformers/en/model_doc/nemotron3_diarization).
+
+<details>
+  <summary>➡️ Offline diarization</summary>
+
+```python
+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")
+```
+</details>
+
+<details>
+  <summary>➡️ Streaming diarization</summary>
+
+Audio arrives chunk by chunk, and each forward takes one chunk: the processor cuts it for its `streaming_mode` and
+adds `num_lookahead_frames`, the number of trailing look-ahead frames the model attends to but does not score, since
+they open the next chunk. The forward returns the `speaker_cache` to pass to the next call. The last chunk of a
+session is extracted with `is_last_audio_chunk=True`: it has no look-ahead, so every remaining frame is scored.
+
+
+| `streaming_mode`          | Latency¹ |
+| ------------------------- | -------- |
+| `"low_latency"` (default) | 1.04 s   |
+| `"very_low_latency"`      | 0.64 s   |
+| `"ultra_low_latency"`     | 0.32 s   |
+
+
+¹ Audio to wait for before the model runs on a chunk: the chunk plus its look-ahead, excluding compute time.
+
+```python
+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}, ...]
+```
+</details>
+
 ### Integration with Streaming ASR
 
 Please refer to the [ASR Integration Guide](ASR_INTEGRATION_GUIDE.md) for the detailed instructions on integration of Nemotron Diarization with Streaming ASR.

config.json

@@ -0,0 +1,54 @@
+{
+  "architectures": [
+    "Nemotron3DiarizationForAudioFrameClassification"
+  ],
+  "audio_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "gelu",
+    "hidden_size": 512,
+    "initializer_range": 0.02,
+    "intermediate_size": 2048,
+    "max_position_embeddings": 5000,
+    "model_type": "nemotron3_diarization_audio",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 31,
+    "num_key_value_heads": 8,
+    "num_mel_bins": 128,
+    "partial_rotary_factor": 1.0,
+    "rope_parameters": {
+      "partial_rotary_factor": 1.0,
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    },
+    "subsampling_factor": 8
+  },
+  "chunk_length": 340,
+  "chunk_right_context": 40,
+  "dtype": "float32",
+  "fifo_length": 40,
+  "head_config": {
+    "audio_hidden_size": 512,
+    "hidden_size": 192,
+    "model_type": "",
+    "num_speakers": 8,
+    "subsampling_factor": 8
+  },
+  "initializer_range": 0.02,
+  "model_type": "nemotron3_diarization",
+  "speaker_cache_update_period": 300,
+  "streaming_config": {
+    "fifo_length": 264,
+    "latest_frames_score_boost": 0.05,
+    "min_positive_scores_rate": 0.5,
+    "model_type": "",
+    "num_speakers": 8,
+    "prediction_score_threshold": 0.25,
+    "speaker_cache_length": 264,
+    "speaker_cache_silence_frames_per_speaker": 1,
+    "speaker_cache_update_period": 222,
+    "strong_boost_rate": 0.75,
+    "subsampling_factor": 8,
+    "weak_boost_rate": 1.5
+  },
+  "transformers_version": "5.18.0.dev0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c074d86335b3b794f8fa5edc25594558f128bdb3914d27806a3a5a2e44963cb6
+size 396954592

processor_config.json

@@ -0,0 +1,31 @@
+{
+  "feature_extractor": {
+    "feature_extractor_type": "NemotronAsrStreamingFeatureExtractor",
+    "feature_size": 128,
+    "hop_length": 160,
+    "n_fft": 512,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "preemphasis": 0.97,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "win_length": 400
+  },
+  "processor_class": "Nemotron3DiarizationProcessor",
+  "streaming_mode": "low_latency",
+  "streaming_modes": {
+    "low_latency": [
+      9,
+      4
+    ],
+    "ultra_low_latency": [
+      3,
+      1
+    ],
+    "very_low_latency": [
+      6,
+      2
+    ]
+  },
+  "subsampling_factor": 8
+}