diar_streaming_sortformer_4spk-v2-hf

Sortformer 4-speaker streaming diarization (HF format)

This is a format conversion of NVIDIA NeMo's nvidia/diar_streaming_sortformer_4spk-v2 into the 🤗 Transformers format. The weights are bit-identical to the original NeMo checkpoint — only the serialization (.nemo → config.json + safetensors) and the surrounding inference code differ. No retraining, fine-tuning, quantization, or merging was performed.

Sortformer is an end-to-end neural speaker-diarization model. It emits, for every ~80 ms output frame, an independent probability for each of up to 4 speakers being active (multi-label / sigmoid), so it naturally handles overlapping speech.

Usage

Offline (full-sequence)

import torch
from transformers import AutoFeatureExtractor, AutoModelForAudioFrameClassification

model_id = "<this-repo>"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
model = AutoModelForAudioFrameClassification.from_pretrained(model_id).eval()

# `audio` is a 16 kHz mono waveform. The offline path peak-normalizes the waveform first:
import numpy as np
audio = audio / (np.abs(audio).max() + 1e-3)
inputs = feature_extractor(audio, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
    probs = model(**inputs).logits.sigmoid()  # (batch, num_frames, num_speakers)

Streaming (Arrival-Order Speaker Cache)

# NOTE: the streaming path expects features extracted WITHOUT peak-normalization.
inputs = feature_extractor(audio, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
    probs = model.diarize_streaming(inputs.input_features, attention_mask=inputs.attention_mask)

Runnable examples (including a live-microphone demo with a real-time plot) are in the Transformers repo under examples/pytorch/speaker-diarization/.

Conversion & parity

Converted with transformers/src/transformers/models/sortformer/convert_sortformer_nemo_to_hf.py. The HF implementation was validated to reproduce the original NeMo model to float precision:

License

Inherited from the source model: CC-BY-4.0. Please cite the original Sortformer work and the NVIDIA NeMo release.