feat: Nemotron-3-Diarization support + speaker-attributed ASR

AGENTS.md

@@ -76,12 +76,18 @@ src/                 libparakeet implementation
                        tdt.cpp / rnnt.cpp , TDT / RNNT greedy loops
                        streaming_encoder.hpp/cpp, cache-aware streaming FastConformer encoder
                        streaming.hpp/cpp  , pk::StreamingSession (carried RNN-T + EOU events) + run_stream_over_pcm
+                       diarization.hpp/cpp, pk::DiarizationModel: offline speaker diarization (Sortformer)
+                       diarization_encoder/head.*, RoPE Transformer encoder + speaker head
+                       diarization_streaming.*, NeMo cache-aware streaming diarization (speaker cache + FIFO)
+                       sas_merge.hpp/cpp  , ASR words x speaker segments -> speaker-attributed utterances
 examples/cli/        parakeet-cli binary
                        subcommands: info, transcribe (+ --stream), quantize
+                     diarize binary: diarize <diar.gguf> <wav> [--stream]
 scripts/             Python tooling
                        convert_parakeet_to_gguf.py, .nemo/.hf -> GGUF (--dtype f32|f16|q8_0)
                        gen_nemo_baseline.py        , NeMo intermediates -> baseline.gguf
                        gen_stream_baseline.py      , NeMo cache-aware streaming encode+decode -> stream baseline.gguf
+                       gen_diar_baseline.py        , NeMo offline + streaming diarization -> diar baseline.gguf
                        validate_vs_nemo.py         , WER parity gate vs NeMo
                        publish_hf.py               , convert+quantize -> HF upload (dry-run default)
                        requirements.txt            , nemo_toolkit[asr] + gguf
@@ -101,10 +107,15 @@ tests/               ctest targets
                        test_streaming_decode.cpp , streaming RNN-T tokens == NeMo cache-aware streaming
                        test_streaming_eou_reset.cpp, multi-utterance streaming: decoder resets on <EOU>, transcript == NeMo reset-on-EOU (issue #13; PARAKEET_TEST_BASELINE_EOU_RESET)
                        test_capi_stream.cpp    , streaming C-API transcript == NeMo streaming (PARAKEET_TEST_BASELINE_EOU_STREAM)
+                       test_diarization_accuracy.cpp, offline diarization == NeMo (PARAKEET_TEST_BASELINE_DIAR)
+                       test_streaming_diarization.cpp, streaming diarization == NeMo streaming (same baseline)
+                       test_combined_offline.cpp, SAS + streaming diarization/SAS through the C-API
+                       test_sas_merge.cpp      , SAS merge/grouping (model-independent)
                        python/check_convert.py , converter round-trip (model-dependent)
                        python/check_baseline.py, baseline dumper (model-dependent)
                        fixtures/clip.wav       , 2 s 16 kHz mono WAV for stage parity tests
                        fixtures/speech.wav     , LibriSpeech 2086-149220-0033, ~7.4 s
+                       fixtures/two_speakers.wav, LibriSpeech 1272 + 2086 alternating A-B-A-B, 23.6 s
 third_party/         vendored deps
                        ggml/     , submodule pinned at v0.13.0
                        dr_wav.h  , vendored single header
@@ -114,6 +125,7 @@ docs/
   conversion.md     , GGUF schema reference
   quantization.md   , quantization allowlist, policy, measured size + WER per type
   parity.md         , full model coverage matrix + per-stage tensor parity
+  diarization.md    , speaker diarization + speaker-attributed ASR: parity, C-API, speed
 .github/workflows/
   ci.yml            , build job (per-push) + closed-loop job (pull_request + dispatch)
 ```
@@ -258,6 +270,18 @@ parakeet_capi_stream_finalize   # flush the end-of-stream tail
 parakeet_capi_stream_free
 ```
 
+Speaker diarization (ABI v7, additive; not used by LocalAI yet). A
+diarization GGUF loads into its own `parakeet_ctx`; see `docs/diarization.md`:
+
+```
+parakeet_capi_diarize_path / _pcm              # offline, JSON segments
+parakeet_capi_transcribe_and_diarize(_json)    # speaker-attributed ASR (two contexts)
+parakeet_capi_free_sas_results                 # frees the array and every .text
+parakeet_capi_diarize_stream_begin / _feed / _free / _chunk_samples
+parakeet_capi_free_diar_segments
+parakeet_capi_sas_stream_begin / _feed / _free
+```
+
 `parakeet_capi_transcribe_path_json(ctx, wav, decoder)` returns malloc'd UTF-8
 JSON `{"text":..,"words":[{"w","start","end","conf"}],"tokens":[{"id","t","conf"}]}`
 (times in seconds, conf in `(0,1]`), built from
@@ -300,6 +324,24 @@ drain, a word finalizes when the next `▁`-token arrives, the last word on
 
 ## Dumping NeMo baselines
 
+Diarization (needs NeMo main / >= 3.1: NeMo 3.0 cannot load the RoPE
+encoder of nvidia/Nemotron-3-Diarization):
+
+```
+.venv/bin/python scripts/convert_parakeet_to_gguf.py \
+    --model nvidia/Nemotron-3-Diarization --output /tmp/diar.gguf
+.venv/bin/python scripts/gen_diar_baseline.py \
+    --model nvidia/Nemotron-3-Diarization \
+    --audio tests/fixtures/two_speakers.wav --output /tmp/diar_baseline.gguf
+PARAKEET_TEST_DIAR_GGUF=/tmp/diar.gguf PARAKEET_TEST_BASELINE_DIAR=/tmp/diar_baseline.gguf \
+    ctest --test-dir build -R diar --output-on-failure
+```
+
+Quantized diarization GGUFs keep the same segments but move probabilities
+more; set `PARAKEET_TEST_DIAR_PROB_TOL=0.05` for Q8_0.
+
+ASR:
+
 Used by Phase 1 parity tests.  Requires the venv and a 16 kHz mono WAV.
 
 ```

CMakeLists.txt

@@ -92,7 +92,12 @@ set(PARAKEET_SRC
     src/transducer_batch.cpp
     src/tokenizer.cpp
     src/search.cpp
-    src/transcription.cpp)
+    src/transcription.cpp
+    src/diarization.cpp
+    src/diarization_encoder.cpp
+    src/diarization_head.cpp
+    src/sas_merge.cpp
+    src/diarization_streaming.cpp)
 
 if(PARAKEET_SHARED)
   add_library(parakeet SHARED ${PARAKEET_SRC})

docs/diarization.md

@@ -0,0 +1,113 @@
+# Speaker diarization
+
+parakeet.cpp runs [nvidia/Nemotron-3-Diarization](https://huggingface.co/nvidia/Nemotron-3-Diarization),
+a Sortformer model that answers "who spoke when" for up to 8 speakers, and
+combines it with any Parakeet ASR model for speaker-attributed transcripts
+("who said what").
+
+## Model
+
+- Encoder: FeatureStacking (8 mel frames stacked, 80 ms per step) and a
+  31-layer pre-norm Transformer with RoPE attention (d_model 512, 8 heads).
+- Head: projection to 192, a subpixel Conv1d that upsamples 8x back to 10 ms
+  frames, two linear layers and a sigmoid per speaker.
+- Output: per-speaker activity probabilities every 10 ms. Segments come from
+  thresholding at 0.5.
+- Speakers are numbered in order of first appearance.
+
+## Converting
+
+```
+.venv/bin/python scripts/convert_parakeet_to_gguf.py \
+    --model nvidia/Nemotron-3-Diarization --dtype q8_0 \
+    --output models/nemotron-3-diarization.q8_0.gguf
+```
+
+`--model` also takes a local `.nemo`. The converter reads the checkpoint
+directly, so it works with a NeMo that cannot instantiate the model (NeMo 3.0
+has no RoPE Transformer encoder). Sizes: F32 397 MB, F16 201 MB, Q8_0 109 MB.
+
+## Offline and streaming inference
+
+NeMo's `diarize()` for this checkpoint runs cache-aware streaming inference
+(`streaming_mode: true` in the model config), and `parakeet_capi_diarize_*`
+does the same: audio is processed in 264-step chunks (21.12 s) with a
+speaker cache of 264 steps that keeps speaker identities stable across chunks.
+This is also what keeps long recordings correct: the model is trained on
+sessions of about 105 s, and attending over a whole long clip at once is
+outside that range. On a 12 minute, 3 speaker recording the offline path
+(which matches NeMo offline on short clips) agrees with NeMo's `diarize()` on
+22% of speech frames and puts almost everything on one speaker; the streaming
+path agrees on 100%.
+
+`DiarizationModel::run_offline` / `speaker_probs` still implement the offline
+path (NeMo peak-normalizes the waveform there) for short clips and for parity
+checks.
+
+The live streaming API (`parakeet_capi_diarize_stream_*`) takes 16 kHz PCM in
+pieces of any size, computes the log-mel incrementally (bit-identical to the
+whole-clip mel) and returns segments once per 21.12 s chunk. Segments that
+continue past a chunk boundary are not split.
+
+## Parity with NeMo
+
+Measured against NeMo main (the reference must support the RoPE encoder), with
+`scripts/gen_diar_baseline.py`:
+
+| Clip | Speakers | Offline prob max diff | Offline segments | Streaming segments |
+|---|---|---|---|---|
+| `tests/fixtures/two_speakers.wav`, 23.6 s | 2 | 0.004 | 5 / 5 identical | 5 / 5 identical |
+| synthetic two-voice dialogue (VibeVoice sample), 68.5 s | 2 | 0.004 | 26 / 26 identical | 26 / 26 identical |
+| synthetic three-voice dialogue (VibeVoice sample), 12.3 min | 3 | | | 100% frame agreement |
+
+Segment boundaries match to the 10 ms frame. F16 gives the same results; Q8_0
+keeps the same segments with a probability max diff of about 0.03.
+
+After the speaker cache compresses, NeMo picks cache frames with
+`torch.topk`, whose order for tied scores is arbitrary. parakeet.cpp breaks ties
+toward the earlier frame, so streaming probabilities can drift by up to about
+0.02 on long clips while the segments stay the same.
+
+## Speaker-attributed ASR
+
+`parakeet_capi_transcribe_and_diarize(_json)` runs an ASR context and a
+diarization context on the same PCM and assigns every ASR word to the speaker
+whose segments overlap it most. A word that overlaps no segment takes the
+nearest segment's speaker if that segment is within 0.5 s, otherwise -1.
+Consecutive words from one speaker (gaps up to 0.5 s) form an utterance.
+
+`parakeet_capi_sas_stream_*` does the same live: when a diarization chunk
+completes, the audio not yet committed is transcribed, and words that end at
+least 1 s before the chunk edge are committed with their speakers. The rest is
+transcribed again with the next chunk, so no word is cut at the edge.
+
+## Speed
+
+End to end with the `diarize` example (model load included), AMD Ryzen 9
+9950X3D, CPU:
+
+| Audio | F32 | Q8_0 |
+|---|---|---|
+| 23.6 s | 0.43 s | 0.25 s |
+| 68.5 s | 0.80 s | 0.60 s |
+| 12.3 min | 6.8 s | 6.0 s |
+
+Streaming cost grows linearly with length (each chunk attends over at most
+528 steps), so long recordings run at about 110x real time.
+
+## Tests
+
+```
+PARAKEET_TEST_DIAR_GGUF=diar.gguf PARAKEET_TEST_BASELINE_DIAR=diar_baseline.gguf \
+PARAKEET_TEST_GGUF=asr.gguf ctest --test-dir build -R "diar|sas|combined"
+```
+
+- `test_diarization_accuracy`: offline probabilities and segments, and the
+  default `diarize_pcm` segments, against NeMo.
+- `test_streaming_diarization`: streaming probabilities and segments against
+  NeMo streaming, from a whole-clip mel and from live 100 ms PCM pieces.
+- `test_combined_offline`: speaker-attributed ASR and the streaming C-API on
+  the two-speaker fixture.
+- `test_sas_merge`: word to speaker assignment (no model needed).
+
+Set `PARAKEET_TEST_DIAR_PROB_TOL=0.05` for Q8_0.

examples/cli/CMakeLists.txt

@@ -1,3 +1,7 @@
 add_executable(parakeet-cli main.cpp)
 target_link_libraries(parakeet-cli PRIVATE parakeet)
 target_include_directories(parakeet-cli PRIVATE ${CMAKE_SOURCE_DIR}/src)
+
+add_executable(diarize diarize.cpp)
+target_link_libraries(diarize PRIVATE parakeet)
+target_include_directories(diarize PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/src)

examples/cli/diarize.cpp

@@ -0,0 +1,79 @@
+// Standalone diarize tool: loads a diarization GGUF and diarizes a WAV.
+// Usage: diarize <gguf> <wav> [--stream]
+// Prints {"speakers":N,"segments":[{"speaker","start","end"}, ...]}.
+// --stream feeds the audio through the streaming C-API in 1 s pieces (NeMo
+// cache-aware streaming) instead of the offline path.
+#include "parakeet_capi.h"
+#include "audio_io.hpp"
+
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+static int diarize_stream(parakeet_ctx* ctx, const char* wav) {
+    pk::Audio audio;
+    if (!pk::load_audio_16k_mono(wav, audio)) {
+        std::fprintf(stderr, "cannot read %s\n", wav);
+        return 1;
+    }
+    parakeet_diar_stream* s = parakeet_capi_diarize_stream_begin(ctx);
+    if (!s) {
+        std::fprintf(stderr, "stream_begin failed: %s\n", parakeet_capi_last_error(ctx));
+        return 1;
+    }
+    std::vector<parakeet_diar_segment> all;
+    const int n = (int)audio.samples.size();
+    for (int lo = 0; lo < n || lo == 0; lo += 16000) {
+        const int len = std::min(16000, n - lo);
+        parakeet_diar_segment* segs = nullptr;
+        int ns = 0;
+        if (parakeet_capi_diarize_stream_feed(s, audio.samples.data() + lo, len,
+                                              lo + len >= n, &segs, &ns) != 0) {
+            std::fprintf(stderr, "stream_feed failed: %s\n", parakeet_capi_last_error(ctx));
+            parakeet_capi_diarize_stream_free(s);
+            return 1;
+        }
+        all.insert(all.end(), segs, segs + ns);
+        parakeet_capi_free_diar_segments(segs);
+        if (lo + len >= n) break;
+    }
+    parakeet_capi_diarize_stream_free(s);
+    std::sort(all.begin(), all.end(), [](const auto& a, const auto& b) {
+        return a.start != b.start ? a.start < b.start : a.speaker < b.speaker;
+    });
+    std::printf("{\"speakers\":8,\"segments\":[");
+    for (size_t i = 0; i < all.size(); ++i)
+        std::printf("%s{\"speaker\":%d,\"start\":%.2f,\"end\":%.2f}", i ? "," : "",
+                    all[i].speaker, all[i].start, all[i].end);
+    std::printf("]}\n");
+    return 0;
+}
+
+int main(int argc, char** argv) {
+    if (argc < 3) {
+        std::fprintf(stderr, "usage: %s <gguf> <wav> [--stream]\n", argv[0]);
+        return 1;
+    }
+    const bool stream = argc > 3 && std::strcmp(argv[3], "--stream") == 0;
+    parakeet_ctx* ctx = parakeet_capi_load(argv[1]);
+    if (!ctx) {
+        std::fprintf(stderr, "failed to load %s\n", argv[1]);
+        return 1;
+    }
+    int rc = 0;
+    if (stream) {
+        rc = diarize_stream(ctx, argv[2]);
+    } else {
+        char* json = parakeet_capi_diarize_path(ctx, argv[2]);
+        if (!json) {
+            std::fprintf(stderr, "diarize failed: %s\n", parakeet_capi_last_error(ctx));
+            rc = 1;
+        } else {
+            std::printf("%s\n", json);
+            parakeet_capi_free_string(json);
+        }
+    }
+    parakeet_capi_free(ctx);
+    return rc;
+}