feat(diarization): add the model card's low-latency streaming modes

AGENTS.md

@@ -108,7 +108,7 @@ tests/               ctest targets
                        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_streaming_diarization.cpp, streaming diarization == NeMo streaming, every latency mode (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)
@@ -277,9 +277,9 @@ 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_diarize_stream_begin / _begin_latency / _feed / _active / _time / _free / _chunk_samples
 parakeet_capi_free_diar_segments
-parakeet_capi_sas_stream_begin / _feed / _free
+parakeet_capi_sas_stream_begin / _begin_latency / _feed / _free
 ```
 
 `parakeet_capi_transcribe_path_json(ctx, wav, decoder)` returns malloc'd UTF-8
@@ -338,7 +338,8 @@ PARAKEET_TEST_DIAR_GGUF=/tmp/diar.gguf PARAKEET_TEST_BASELINE_DIAR=/tmp/diar_bas
 ```
 
 Quantized diarization GGUFs keep the same segments but move probabilities
-more; set `PARAKEET_TEST_DIAR_PROB_TOL=0.05` for Q8_0.
+more; set `PARAKEET_TEST_DIAR_PROB_TOL=0.15` for Q8_0. The baseline also
+holds each low-latency streaming mode (`--modes`, NeMo is slow on them).
 
 ASR:
 

docs/diarization.md

@@ -44,10 +44,30 @@ path agrees on 100%.
 path (NeMo peak-normalizes the waveform there) for short clips and for parity
 checks.
 
+## Live streaming and latency modes
+
 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.
+pieces of any size and computes the log-mel incrementally (bit-identical to
+the whole-clip mel). Segments that continue past a chunk boundary are not
+split.
+
+The checkpoint's own configuration processes 21.12 s chunks, too slow for live
+labels. The same checkpoint also supports the model card's low-latency
+configurations, selected with `parakeet_capi_diarize_stream_begin_latency`:
+
+| Mode | Input latency | Chunk | Look-ahead | FIFO | Cache | Update period |
+|---|---|---|---|---|---|---|
+| `PARAKEET_DIAR_LATENCY_MODEL` | 21.12 s | 264 | 0 | 0 | 264 | 264 |
+| `PARAKEET_DIAR_LATENCY_LOW` | 1.04 s | 9 | 4 | 264 | 264 | 222 |
+| `PARAKEET_DIAR_LATENCY_VERY_LOW` | 0.64 s | 6 | 2 | 264 | 264 | 222 |
+| `PARAKEET_DIAR_LATENCY_ULTRA_LOW` | 0.32 s | 3 | 1 | 264 | 264 | 222 |
+
+Sizes are 80 ms encoder frames; input latency is (chunk + look-ahead) x 80 ms.
+Each chunk is encoded together with its look-ahead, as NeMo's
+`streaming_feat_loader` does, and only the chunk itself is output and enters
+the FIFO and speaker cache. `parakeet_capi_diarize_stream_active` returns the
+segments still open ("who is speaking now") and
+`parakeet_capi_diarize_stream_time` how much audio has been diarized.
 
 ## Parity with NeMo
 
@@ -60,8 +80,18 @@ Measured against NeMo main (the reference must support the RoPE encoder), with
 | 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.
+Every latency mode matches NeMo in the same mode: on the 23.6 s fixture all
+four modes give NeMo's segments exactly (5, 5, 6 and 7 segments), and on the
+68.5 s clip the 1.04 s mode, which fills the FIFO and compresses the speaker
+cache, gives all 26 segments; probability max diff 0.01 or less throughout.
+
+Segment boundaries match to the 10 ms frame. F16 gives the same results. Q8_0
+keeps the same segments on the fixture and the 68.5 s clip; its probabilities
+move up to about 0.06 in the checkpoint mode and 0.12 in the low-latency modes
+there, where quantization noise changes which frames the cache keeps. Where a
+probability sits right at the 0.5 threshold, Q8_0 can also flip a frame or two:
+on a 31.5 s two-speaker clip that splits one segment into three short pieces
+(99.5% frame agreement). Use F16 when segment-exact output matters.
 
 After the speaker cache compresses, NeMo picks cache frames with
 `torch.topk`, whose order for tied scores is arbitrary. parakeet.cpp breaks ties
@@ -76,13 +106,31 @@ 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.
+`parakeet_capi_sas_stream_*` does the same live (`_begin_latency` takes a
+latency mode): once at least 4 s of diarized audio is uncommitted it is
+transcribed, and words that end at least 1 s before the diarized edge are
+committed with their speakers. Transcription resumes right after the last
+committed word, so no word is cut at the edge and audio the ASR skipped is
+heard again with more context. Offline ASR on shorter windows loses words, so
+committed text lags by about 3 to 5 s whatever the diarization latency.
 
 ## Speed
 
+Streaming modes on the 12.3 min, 3 speaker clip (model load included; agreement
+is with NeMo's default `diarize()`, so it shows what the lower latency costs):
+
+| Mode | F16 | x real time | Agreement | Q8_0 | Agreement |
+|---|---|---|---|---|---|
+| checkpoint (21.12 s) | 5.7 s | 130x | 100% | 6.0 s | 99.9% |
+| low (1.04 s) | 107 s | 6.9x | 99.4% | 115 s | 99.4% |
+| very low (0.64 s) | 166 s | 4.4x | 99.1% | 172 s | 99.2% |
+| ultra low (0.32 s) | 295 s | 2.5x | 98.3% | 332 s | 98.4% |
+
+A low-latency chunk costs about 0.1 s on this CPU, so labels arrive about
+0.1 s after the input latency.
+
+Whole-file diarization:
+
 End to end with the `diarize` example (model load included), AMD Ryzen 9
 9950X3D, CPU:
 
@@ -110,4 +158,6 @@ PARAKEET_TEST_GGUF=asr.gguf ctest --test-dir build -R "diar|sas|combined"
   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.
+Set `PARAKEET_TEST_DIAR_PROB_TOL=0.15` for Q8_0 (the segment checks stay
+exact). The baseline includes every latency mode by default; NeMo runs the
+low-latency modes slowly on CPU, so `--modes low` limits a long clip to one.

examples/cli/diarize.cpp

@@ -1,8 +1,9 @@
 // Standalone diarize tool: loads a diarization GGUF and diarizes a WAV.
-// Usage: diarize <gguf> <wav> [--stream]
+// Usage: diarize <gguf> <wav> [--stream [model|low|very_low|ultra_low]]
 // 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.
+// --stream feeds the audio through the streaming C-API in 100 ms pieces (NeMo
+// cache-aware streaming) instead of the whole-file path, optionally in one of
+// the model card's latency modes (default: the checkpoint's configuration).
 #include "parakeet_capi.h"
 #include "audio_io.hpp"
 
@@ -11,21 +12,21 @@
 #include <cstring>
 #include <vector>
 
-static int diarize_stream(parakeet_ctx* ctx, const char* wav) {
+static int diarize_stream(parakeet_ctx* ctx, const char* wav, int latency) {
     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);
+    parakeet_diar_stream* s = parakeet_capi_diarize_stream_begin_latency(ctx, latency);
     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);
+    for (int lo = 0; lo < n || lo == 0; lo += 1600) {
+        const int len = std::min(1600, n - lo);
         parakeet_diar_segment* segs = nullptr;
         int ns = 0;
         if (parakeet_capi_diarize_stream_feed(s, audio.samples.data() + lo, len,
@@ -52,18 +53,29 @@ static int diarize_stream(parakeet_ctx* ctx, const char* wav) {
 
 int main(int argc, char** argv) {
     if (argc < 3) {
-        std::fprintf(stderr, "usage: %s <gguf> <wav> [--stream]\n", argv[0]);
+        std::fprintf(stderr, "usage: %s <gguf> <wav> [--stream [model|low|very_low|ultra_low]]\n", argv[0]);
         return 1;
     }
     const bool stream = argc > 3 && std::strcmp(argv[3], "--stream") == 0;
+    int latency = PARAKEET_DIAR_LATENCY_MODEL;
+    if (stream && argc > 4) {
+        const char* names[] = {"model", "low", "very_low", "ultra_low"};
+        latency = -1;
+        for (int i = 0; i < 4; ++i)
+            if (std::strcmp(argv[4], names[i]) == 0) latency = i;
+        if (latency < 0) {
+            std::fprintf(stderr, "unknown latency %s (model, low, very_low, ultra_low)\n", argv[4]);
+            return 1;
+        }
+    }
     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]);
+        rc = diarize_stream(ctx, argv[2], latency);
     } else {
         char* json = parakeet_capi_diarize_path(ctx, argv[2]);
         if (!json) {

include/parakeet_capi.h

@@ -52,6 +52,9 @@ typedef struct parakeet_ctx parakeet_ctx;
 //     for nvidia/Nemotron-3-Diarization. A parakeet_ctx now holds either an
 //     ASR or a diarization model; parakeet_capi_load detects which. No
 //     existing signatures changed.
+//     Additive, same ABI: parakeet_capi_diarize_stream_begin_latency /
+//     _time / _active and parakeet_capi_sas_stream_begin_latency (the model
+//     card's 1.04 / 0.64 / 0.32 s streaming modes).
 int parakeet_capi_abi_version(void);
 
 // Load a GGUF model. Returns an owning context, or NULL on failure.
@@ -399,11 +402,16 @@ char* parakeet_capi_transcribe_and_diarize_json(parakeet_ctx* asr_ctx, parakeet_
 
 // --- Streaming diarization -------------------------------------------------
 // NeMo cache-aware streaming (speaker cache + FIFO) over live 16 kHz mono
-// float PCM. Audio is processed in the model's chunks
-// (parakeet_capi_diarize_stream_chunk_samples; 21.12 s for
-// Nemotron-3-Diarization), so segments arrive once per chunk. Speaker indices
-// stay consistent across chunks. The stream borrows `diar_ctx`: free the
-// stream first, and do not use one context from two threads at once.
+// float PCM. Speaker indices stay consistent across chunks. The stream
+// borrows `diar_ctx`: free the stream first, and do not use one context from
+// two threads at once.
+//
+// Latency modes (the Nemotron-3-Diarization model card presets). Latency is
+// the audio buffered before a chunk runs: (chunk + look-ahead) x 80 ms.
+#define PARAKEET_DIAR_LATENCY_MODEL     0  // checkpoint config: 21.12 s chunks
+#define PARAKEET_DIAR_LATENCY_LOW       1  // 1.04 s
+#define PARAKEET_DIAR_LATENCY_VERY_LOW  2  // 0.64 s
+#define PARAKEET_DIAR_LATENCY_ULTRA_LOW 3  // 0.32 s
 
 typedef struct parakeet_diar_segment {
     int   speaker;
@@ -413,12 +421,26 @@ typedef struct parakeet_diar_segment {
 
 typedef struct parakeet_diar_stream parakeet_diar_stream;
 
-// NULL on error (last_error on diar_ctx).
+// Begin a stream in the checkpoint's own configuration
+// (PARAKEET_DIAR_LATENCY_MODEL). NULL on error (last_error on diar_ctx).
 parakeet_diar_stream* parakeet_capi_diarize_stream_begin(parakeet_ctx* diar_ctx);
 
-// Samples per processing chunk (the segment latency). 0 on NULL.
+// Begin a stream in one of the PARAKEET_DIAR_LATENCY_* modes.
+parakeet_diar_stream* parakeet_capi_diarize_stream_begin_latency(parakeet_ctx* diar_ctx,
+                                                                 int latency);
+
+// Samples buffered before a chunk runs (the input latency). 0 on NULL.
 int parakeet_capi_diarize_stream_chunk_samples(parakeet_diar_stream* s);
 
+// Seconds of audio diarized so far (trails the audio fed by up to the latency).
+float parakeet_capi_diarize_stream_time(parakeet_diar_stream* s);
+
+// Segments still open at the diarized time ("who is speaking now"), with `end`
+// at parakeet_capi_diarize_stream_time. Same ownership as
+// parakeet_capi_diarize_stream_feed. Returns 0, or non-zero on error.
+int parakeet_capi_diarize_stream_active(parakeet_diar_stream* s,
+                                        parakeet_diar_segment** out, int* n_out);
+
 // Feed PCM; `is_last` flushes the tail and closes open segments. Returns 0 and
 // sets *out / *n_out to the segments that ENDED since the previous call
 // (free with parakeet_capi_free_diar_segments; *out may be NULL when
@@ -431,17 +453,22 @@ void parakeet_capi_free_diar_segments(parakeet_diar_segment* segs);
 void parakeet_capi_diarize_stream_free(parakeet_diar_stream* s);
 
 // --- Streaming speaker-attributed ASR ---------------------------------------
-// Streaming diarization plus ASR over the same live 16 kHz PCM. Each time a
-// diarization chunk completes, the not-yet-committed audio is transcribed;
+// Streaming diarization plus ASR over the same live 16 kHz PCM. Each time
+// diarization advances, the not-yet-committed audio is transcribed;
 // all words but the last (which may still be cut by the chunk edge) are
 // committed with their speakers, and the rest is carried into the next
 // chunk. `is_last` commits everything. Borrows both contexts.
 
 typedef struct parakeet_sas_stream parakeet_sas_stream;
 
-// NULL on error (last_error on the context that failed).
+// NULL on error (last_error on the context that failed). _begin uses the
+// checkpoint's diarization config; _begin_latency takes a
+// PARAKEET_DIAR_LATENCY_* mode, which also sets how often words commit.
 parakeet_sas_stream* parakeet_capi_sas_stream_begin(parakeet_ctx* asr_ctx,
                                                     parakeet_ctx* diar_ctx);
+parakeet_sas_stream* parakeet_capi_sas_stream_begin_latency(parakeet_ctx* asr_ctx,
+                                                            parakeet_ctx* diar_ctx,
+                                                            int latency);
 
 // Returns 0 and sets *out / *n_out to the utterances committed by this call
 // (free with parakeet_capi_free_sas_results(*out, *n_out)). Consecutive calls

scripts/gen_diar_baseline.py

@@ -21,6 +21,11 @@
                                        (``streaming_mode=True``, the model's
                                        own chunk / speaker-cache config)
 * ``stream_segs``    ``[N, 3]``        streaming diarize() segments
+* ``stream_probs_<mode>`` / ``stream_segs_<mode>``
+                                       the same for each low-latency streaming
+                                       mode of the model card (LATENCY_MODES)
+                                       listed in --modes (default: all). NeMo
+                                       runs these slowly on CPU.
 
 ``dither`` is forced to 0 so the mel is deterministic (the C++ side has no
 dither).
@@ -46,20 +51,45 @@
     sys.exit(2)
 
 
-def _segments(model, path):
-    out = model.diarize(audio=[path], batch_size=1)
+# Streaming configurations from the Nemotron-3-Diarization model card, in 80 ms
+# encoder frames: (spkcache_len, fifo_len, chunk_len, chunk_right_context,
+# spkcache_update_period). Latency = (chunk_len + right_context) * 80 ms.
+LATENCY_MODES = {
+    "low": (264, 264, 9, 4, 222),         # 1.04 s
+    "very_low": (264, 264, 6, 2, 222),    # 0.64 s
+    "ultra_low": (264, 264, 3, 1, 222),   # 0.32 s
+}
+
+
+def _set_streaming(model, spkcache, fifo, chunk, right, update):
+    sm = model.sortformer_modules
+    sm.spkcache_len = spkcache
+    sm.fifo_len = fifo
+    sm.chunk_len = chunk
+    sm.chunk_right_context = right
+    sm.spkcache_update_period = update
+    model._check_streaming_parameters()
+
+
+def _diarize(model, path):
+    """One diarize() run: (probs [n_spk, T], segments [N, 3]) as NeMo returns them."""
+    lines, preds = model.diarize(audio=[path], batch_size=1, include_tensor_outputs=True)
     rows = []
-    for s in out[0]:
+    for s in lines[0]:
         start, end, spk = s.split()
         rows.append([float(spk.split("_")[-1]), float(start), float(end)])
-    return np.asarray(rows, dtype=np.float32).reshape(-1, 3)
+    probs = preds[0][0].detach().cpu().numpy().T.copy()
+    return probs, np.asarray(rows, dtype=np.float32).reshape(-1, 3)
 
 
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--model", required=True, help=".nemo path or HF id")
     ap.add_argument("--audio", required=True, help="16 kHz mono wav")
     ap.add_argument("--output", required=True)
+    ap.add_argument("--modes", default=",".join(LATENCY_MODES),
+                    help="comma-separated low-latency modes to capture "
+                         f"({', '.join(LATENCY_MODES)}); empty for none")
     args = ap.parse_args()
 
     if args.model.endswith(".nemo"):
@@ -72,18 +102,25 @@ def main():
     y, sr = sf.read(args.audio, dtype="float32")
     if y.ndim != 1 or sr != 16000:
         sys.exit(f"gen_diar_baseline: {args.audio} must be 16 kHz mono (got sr={sr}, shape={y.shape})")
-    x = torch.from_numpy(y)[None]
-    n = torch.tensor([len(y)])
+
+    sm = m.sortformer_modules
+    saved = (sm.spkcache_len, sm.fifo_len, sm.chunk_len, sm.chunk_right_context,
+             sm.spkcache_update_period)
+    runs = [("offline", False, None), ("stream", True, saved)]
+    for k in filter(None, args.modes.split(",")):
+        runs.append((f"stream_{k}", True, LATENCY_MODES[k]))
 
     results = {}
-    for name, streaming in (("offline", False), ("stream", True)):
+    for name, streaming, cfg in runs:
         m.streaming_mode = streaming
+        if cfg is not None:
+            _set_streaming(m, *cfg)
+        # Keys: offline_probs, stream_probs, stream_probs_low, ...
+        pk, sk = (f"{name}_probs", f"{name}_segs") if "_" not in name else \
+            (name.replace("stream_", "stream_probs_"), name.replace("stream_", "stream_segs_"))
         with torch.no_grad():
-            preds = m.forward(x, n)  # [1, T, n_spk]
-        results[f"{name}_probs"] = preds[0].numpy().T.copy()  # [n_spk, T]
-        results[f"{name}_segs"] = _segments(m, args.audio)
-        print(f"{name}: probs {results[name + '_probs'].shape}, "
-              f"{len(results[name + '_segs'])} segments")
+            results[pk], results[sk] = _diarize(m, args.audio)
+        print(f"{name}: probs {results[pk].shape}, {len(results[sk])} segments")
 
     w = gguf.GGUFWriter(args.output, "parakeet-diar-baseline")
     w.add_tensor("audio", np.ascontiguousarray(y, dtype=np.float32))