phPEE-canary-enc-1b-v2-sortformer-v2.1

Parallel Expert Encoder (PEE) — Canary 1B v2 + Sortformer v2.1

A Parallel Expert speech Encoder that runs two experts on the same mel input and fuses their outputs:

• ASR expert — a Canary 1B v2 Conformer encoder.
• Speaker expert — a Sortformer v2.1 streaming speaker-diarization model.

The two branches run in parallel on the same (un-normalised) mel features; the ASR branch internally re-applies normalize_batch. The speaker-activity predictions are fused into the ASR hidden states via LayerNorm → sinusoidal speaker kernel → ADD. The module is a drop-in replacement for ConformerEncoder — identical forward(audio_signal, length) I/O (output shape (B, D, T)), so it can replace the encoder in a SALM/Canary perception stack.

Highlights

• Drop-in with ConformerEncoder (same d_model, subsampling_factor, pre_encode).
• Optional speaker override: pass spk_targets (RTTM/oracle (B, T, n_spk)) to skip running Sortformer.
• Long-form / online inference: when input exceeds one window and Sortformer runs, it processes the audio in fixed windows with overlap-and-trim ASR context while the streaming diarizer carries its speaker-cache/FIFO state across windows.
• Self-contained .nemo bundle: model_config.yaml carries inline asr_encoder_cfg
  • diarization_model_cfg. By default the diarizer is frozen (freeze_diar=True) and the ASR encoder is trainable (freeze_asr=False).

Usage

Implemented by ParallelExpertEncoder / ParallelExpertEncoderPT in NeMo: nemo/collections/asr/modules/parallel_expert_encoder.py.

from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoderPT

# Load this self-contained PE bundle and get the inner encoder
encoder = ParallelExpertEncoderPT.load_from_nemo("path/to/bundle.nemo")

# Drop-in ConformerEncoder I/O: audio_signal = (B, feat_in, n_frames) un-normalised mels
outputs, encoded_lengths = encoder(audio_signal=mels, length=lengths)

Key config

Note: This is an experimental module and a placeholder PEE checkpoint.