From 75a536d809755da410b538a963316969059c98d1 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 28 Sep 2026 09:38:01 +0000 Subject: [PATCH 1/2] feat(diarization): add the model card's low-latency streaming modes Streaming diarization only ran the checkpoint's own configuration: 21.12 s chunks with no look-ahead, too slow for live speaker labels. The Nemotron-3-Diarization model card documents three low-latency configurations for the same checkpoint (1.04, 0.64 and 0.32 s input latency), which need a FIFO and per-chunk look-ahead. StreamingDiarization now takes a DiarStreamConfig (speaker cache, FIFO, chunk, left/right context, update period) with the presets as DiarLatency. It buffers mel frames itself and runs a chunk once its look-ahead has arrived, encoding the chunk with its context frames as NeMo's streaming_feat_loader does; only the chunk is output and enters the FIFO and speaker cache. Every mode matches NeMo in the same mode: all four give NeMo's segments exactly on the 23.6 s fixture, and the 1.04 s mode, which fills the FIFO and compresses the cache, gives all 26 segments on a 68.5 s clip. On CPU the 1.04 s mode runs 6.9x faster than real time. C-API, additive within ABI 7: - diarize_stream_begin_latency and sas_stream_begin_latency pick a PARAKEET_DIAR_LATENCY_* mode. - diarize_stream_active returns the segments still open ("who is speaking now"), diarize_stream_time the diarized time. - diarize_stream_chunk_samples now reports the input latency. Streaming SAS transcribed every time diarization advanced. At 1.04 s that meant 1 to 2 s ASR windows, which drop words, and the commit point skipped audio the ASR had missed. It now waits for 4 s of uncommitted audio and resumes right after the last committed word. gen_diar_baseline.py captures each mode (--modes) from one diarize() run per mode, and the diarize example takes --stream . Assisted-by: Claude:claude-opus-5-5 [Claude Code] --- examples/cli/diarize.cpp | 30 ++-- include/parakeet_capi.h | 47 ++++-- scripts/gen_diar_baseline.py | 61 ++++++-- src/diarization.cpp | 18 +-- src/diarization_streaming.cpp | 186 ++++++++++++++++++----- src/diarization_streaming.hpp | 85 ++++++++--- src/parakeet_capi.cpp | 159 ++++++++++++-------- tests/test_combined_offline.cpp | 57 ++++++- tests/test_streaming_diarization.cpp | 215 ++++++++++++++++----------- 9 files changed, 604 insertions(+), 254 deletions(-) diff --git a/examples/cli/diarize.cpp b/examples/cli/diarize.cpp index 1e7233d..a82301f 100644 --- a/examples/cli/diarize.cpp +++ b/examples/cli/diarize.cpp @@ -1,8 +1,9 @@ // Standalone diarize tool: loads a diarization GGUF and diarizes a WAV. -// Usage: diarize [--stream] +// Usage: diarize [--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 #include -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 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,10 +53,21 @@ static int diarize_stream(parakeet_ctx* ctx, const char* wav) { int main(int argc, char** argv) { if (argc < 3) { - std::fprintf(stderr, "usage: %s [--stream]\n", argv[0]); + std::fprintf(stderr, "usage: %s [--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]); @@ -63,7 +75,7 @@ int main(int argc, char** argv) { } 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) { diff --git a/include/parakeet_capi.h b/include/parakeet_capi.h index 31c23b3..789ef48 100644 --- a/include/parakeet_capi.h +++ b/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 diff --git a/scripts/gen_diar_baseline.py b/scripts/gen_diar_baseline.py index 52977f9..a0c6fcf 100644 --- a/scripts/gen_diar_baseline.py +++ b/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_`` / ``stream_segs_`` + 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,13 +51,35 @@ 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(): @@ -60,6 +87,9 @@ def main(): 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)) diff --git a/src/diarization.cpp b/src/diarization.cpp index 17d8915..64f9493 100644 --- a/src/diarization.cpp +++ b/src/diarization.cpp @@ -151,17 +151,17 @@ DiarizationResult DiarizationModel::run_streaming(const std::vector& samp if (T <= 0) return result; const int T_full = (int)(feats.size() / n_mels); - StreamingDiarization sd(loader_); - const int cm = sd.chunk_mel_frames(); - std::vector chunk; - for (int lo = 0; lo < T; lo += cm) { - const int n = std::min(cm, T - lo); - chunk.resize((size_t)n_mels * n); + // Whole clip in one feed: the diarizer chunks it with the checkpoint's + // streaming configuration, exactly as NeMo's diarize() does. + if ((int)(feats.size() / n_mels) != T) { + std::vector trimmed((size_t)n_mels * T); for (int m = 0; m < n_mels; ++m) - std::copy_n(feats.begin() + (size_t)m * T_full + lo, n, chunk.begin() + (size_t)m * n); - for (const auto& g : sd.feed_mel_chunk(chunk, n_mels, n, lo + n >= T)) - result.segments.push_back({g.speaker, g.start, g.end}); + std::copy_n(feats.begin() + (size_t)m * T_full, T, trimmed.begin() + (size_t)m * T); + feats.swap(trimmed); } + StreamingDiarization sd(loader_); + for (const auto& g : sd.feed_mel(feats, n_mels, T, /*is_last=*/true)) + result.segments.push_back({g.speaker, g.start, g.end}); std::sort(result.segments.begin(), result.segments.end(), [](const SpeakerSegment& a, const SpeakerSegment& b) { return a.start != b.start ? a.start < b.start : a.speaker < b.speaker; diff --git a/src/diarization_streaming.cpp b/src/diarization_streaming.cpp index 4cf82fb..c7b3288 100644 --- a/src/diarization_streaming.cpp +++ b/src/diarization_streaming.cpp @@ -30,8 +30,43 @@ float round2(float x) { return std::round(x * 100.0f) / 100.0f; } } // namespace +DiarStreamConfig DiarStreamConfig::from_model(const ParakeetConfig& cfg) { + const auto& d = cfg.diarization; + DiarStreamConfig c; + c.spkcache_len = d.spkcache_len; + c.fifo_len = d.fifo_len; + c.chunk_len = d.chunk_len; + c.left_context = 0; + c.right_context = 0; + c.update_period = d.spkcache_update_period; + return c; +} + +DiarStreamConfig diar_stream_config(DiarLatency latency, const ParakeetConfig& cfg) { + DiarStreamConfig c = DiarStreamConfig::from_model(cfg); + // Model-card presets: speaker cache 264, FIFO 264, update period 222. + auto preset = [&](int chunk, int right) { + c.spkcache_len = 264; + c.fifo_len = 264; + c.chunk_len = chunk; + c.left_context = 0; + c.right_context = right; + c.update_period = 222; + }; + switch (latency) { + case DiarLatency::Model: break; + case DiarLatency::Low: preset(9, 4); break; + case DiarLatency::VeryLow: preset(6, 2); break; + case DiarLatency::UltraLow: preset(3, 1); break; + } + return c; +} + StreamingDiarization::StreamingDiarization(const ModelLoader& ml) - : ml_(ml), encoder_(ml), head_(ml) { + : StreamingDiarization(ml, DiarStreamConfig::from_model(ml.config())) {} + +StreamingDiarization::StreamingDiarization(const ModelLoader& ml, const DiarStreamConfig& scfg) + : ml_(ml), encoder_(ml), head_(ml), scfg_(scfg) { const auto& cfg = ml.config(); const auto& d = cfg.diarization; d_model_ = (int)cfg.d_model; @@ -39,10 +74,6 @@ StreamingDiarization::StreamingDiarization(const ModelLoader& ml) subsampling_ = encoder_.subsampling(); upsample_ = (int)d.upsample_factor; n_mels_ = encoder_.n_mels(); - chunk_len_ = d.chunk_len; - spkcache_len_ = d.spkcache_len; - fifo_len_ = d.fifo_len; - update_period_ = d.spkcache_update_period; sil_frames_per_spk_ = d.spkcache_sil_frames_per_spk; frame_sec_ = d.frame_resolution_sec; onset_ = d.onset_threshold; @@ -54,8 +85,11 @@ StreamingDiarization::StreamingDiarization(const ModelLoader& ml) weak_boost_rate_ = d.weak_boost_rate; min_pos_scores_rate_ = d.min_pos_scores_rate; - if (chunk_len_ <= 0 || spkcache_len_ <= 0 || fifo_len_ < 0 || update_period_ <= 0 || - upsample_ != subsampling_ || spkcache_len_ / n_spk_ - sil_frames_per_spk_ <= 0) { + // NeMo _check_streaming_parameters: every size positive, and the update + // period between the chunk and the FIFO + chunk. + if (scfg_.chunk_len <= 0 || scfg_.spkcache_len <= 0 || scfg_.fifo_len < 0 || + scfg_.left_context < 0 || scfg_.right_context < 0 || scfg_.update_period <= 0 || + upsample_ != subsampling_ || scfg_.spkcache_len / n_spk_ - sil_frames_per_spk_ <= 0) { throw std::runtime_error("parakeet: invalid diarization streaming config"); } @@ -77,6 +111,10 @@ void StreamingDiarization::reset() { spkcache_compressed_ = false; mean_sil_emb_.assign(d_model_, 0.0f); n_sil_frames_ = 0; + buf_.clear(); + buf_start_ = 0; + frames_in_ = 0; + next_stt_ = 0; frames_done_ = 0; last_probs_.clear(); last_frames_ = 0; @@ -84,27 +122,95 @@ void StreamingDiarization::reset() { start_frame_.assign(n_spk_, 0); } -std::vector StreamingDiarization::feed_mel_chunk( +std::vector StreamingDiarization::feed_mel( const std::vector& mel, int n_mels, int n_frames, bool is_last) { - if (n_mels != n_mels_ || n_frames <= 0 || n_frames > chunk_mel_frames() || - mel.size() != (size_t)n_mels * n_frames) { - throw std::runtime_error("parakeet: bad streaming diarization chunk shape"); + if (n_mels != n_mels_ || n_frames < 0 || mel.size() != (size_t)n_mels * n_frames) + throw std::runtime_error("parakeet: bad streaming diarization mel shape"); + + // Buffer frame-major. + const size_t base = buf_.size(); + buf_.resize(base + (size_t)n_frames * n_mels_); + for (int m = 0; m < n_mels_; ++m) + for (int t = 0; t < n_frames; ++t) + buf_[base + (size_t)t * n_mels_ + m] = mel[(size_t)m * n_frames + t]; + frames_in_ += n_frames; + + // Run every chunk whose look-ahead has arrived (NeMo streaming_feat_loader: + // left/right context clipped at the stream edges). + std::vector closed; + std::vector> out(n_spk_); + last_frames_ = 0; + const long long cm = chunk_mel_frames(); + const long long lc = (long long)scfg_.left_context * subsampling_; + const long long rc = (long long)scfg_.right_context * subsampling_; + for (;;) { + const long long stt = next_stt_; + long long end = stt + cm; + if (is_last) { + if (stt >= frames_in_) break; + end = std::min(end, frames_in_); + } else if (frames_in_ < end + rc) { + break; + } + const int left = (int)std::min(lc, stt); + const int right = (int)std::min(rc, frames_in_ - end); + const bool final_chunk = is_last && end == frames_in_; + step(stt, end, left, right, final_chunk, closed); + // step() leaves this chunk's output in last_probs_ ([n_spk, n]). + const int n = (int)(end - stt); + for (int s = 0; s < n_spk_; ++s) + out[s].insert(out[s].end(), last_probs_.begin() + (size_t)s * n, + last_probs_.begin() + (size_t)(s + 1) * n); + last_frames_ += n; + next_stt_ = end; + // Keep only what the next chunk's left context needs. + const long long keep_from = std::max(buf_start_, next_stt_ - lc); + if (keep_from > buf_start_) { + buf_.erase(buf_.begin(), buf_.begin() + (size_t)(keep_from - buf_start_) * n_mels_); + buf_start_ = keep_from; + } } + if (is_last) { + // A stream that ended exactly on a chunk edge still closes its segments. + std::vector none; + track_segments(none, 0, true, closed); + } + last_probs_.clear(); + for (int s = 0; s < n_spk_; ++s) last_probs_.insert(last_probs_.end(), out[s].begin(), out[s].end()); + std::sort(closed.begin(), closed.end(), + [](const StreamingSpeakerSegment& a, const StreamingSpeakerSegment& b) { + return a.start != b.start ? a.start < b.start : a.speaker < b.speaker; + }); + return closed; +} - // 1. Pre-encode the chunk -> [cl, d]. - std::vector chunk_emb; - int cl = 0; - encoder_.pre_encode(mel, n_mels, n_frames, chunk_emb, cl); - - // 2. Transformer + head over [spkcache | fifo | chunk]. +// NeMo forward_streaming_step (sync mode) on mel frames [stt - left, end + right). +void StreamingDiarization::step(long long stt, long long end, int left, int right, + bool final_chunk, std::vector& closed) { + const int n_win = left + (int)(end - stt) + right; + std::vector win((size_t)n_mels_ * n_win); + const long long w0 = stt - left; + for (int t = 0; t < n_win; ++t) + for (int m = 0; m < n_mels_; ++m) + win[(size_t)m * n_win + t] = buf_[(size_t)(w0 - buf_start_ + t) * n_mels_ + m]; + + // 1. Pre-encode the window -> [lc + cl + rc, d] (pre-norm embeddings). + std::vector emb; + int T_c = 0; + encoder_.pre_encode(win, n_mels_, n_win, emb, T_c); + const int lc_enc = (int)std::lround((double)left / subsampling_); + const int rc_enc = (right + subsampling_ - 1) / subsampling_; + const int cl = T_c - lc_enc - rc_enc; + + // 2. Transformer + head over [spkcache | fifo | window]. const int S = (int)(spkcache_.size() / d_model_); const int F = (int)(fifo_.size() / d_model_); - const int total = S + F + cl; + const int total = S + F + T_c; std::vector seq; seq.reserve((size_t)total * d_model_); seq.insert(seq.end(), spkcache_.begin(), spkcache_.end()); seq.insert(seq.end(), fifo_.begin(), fifo_.end()); - seq.insert(seq.end(), chunk_emb.begin(), chunk_emb.end()); + seq.insert(seq.end(), emb.begin(), emb.end()); std::vector enc; encoder_.transformer_forward(seq, total, enc); @@ -122,35 +228,35 @@ std::vector StreamingDiarization::feed_mel_chunk( preds[(size_t)t * n_spk_ + s] = (float)(acc / upsample_); } - // 4. This chunk's high-resolution slice, trimmed to the real mel frames. - const int base = (S + F) * upsample_; - last_frames_ = n_frames; - last_probs_.resize((size_t)n_spk_ * n_frames); + // 4. The chunk's high-resolution slice (after the left context), trimmed + // to its real mel frames. + const int n = (int)(end - stt); + const size_t hr0 = (size_t)(S + F + lc_enc) * upsample_; + last_probs_.resize((size_t)n_spk_ * n); for (int s = 0; s < n_spk_; ++s) - std::copy_n(hp.begin() + (size_t)s * T_hr + base, n_frames, - last_probs_.begin() + (size_t)s * n_frames); + std::copy_n(hp.begin() + (size_t)s * T_hr + hr0, n, last_probs_.begin() + (size_t)s * n); // 5. Cache update for the next chunk. - if (!is_last) streaming_update(chunk_emb, cl, preds, S, F); + if (!final_chunk) streaming_update(emb, lc_enc, cl, preds, S, F); - std::vector out; - track_segments(last_probs_, n_frames, is_last, out); - return out; + track_segments(last_probs_, n, false, closed); } -// SortformerModules.streaming_update (sync mode, lc = rc = 0). -void StreamingDiarization::streaming_update(const std::vector& chunk_emb, - int cl, const std::vector& preds, +// SortformerModules.streaming_update (sync mode). emb is the pre-encoded +// window; the chunk is its rows [lc, lc + cl) and its predictions sit at +// S + F + lc in preds. +void StreamingDiarization::streaming_update(const std::vector& emb, int lc, int cl, + const std::vector& preds, int S, int F) { const int d = d_model_, ns = n_spk_; // FIFO predictions are refreshed from this step's output. fifo_preds_.assign(preds.begin() + (size_t)S * ns, preds.begin() + (size_t)(S + F) * ns); - fifo_.insert(fifo_.end(), chunk_emb.begin(), chunk_emb.end()); - append_rows(fifo_preds_, preds, ns, S + F, S + F + cl); + append_rows(fifo_, emb, d, lc, lc + cl); + append_rows(fifo_preds_, preds, ns, S + F + lc, S + F + lc + cl); - if (F + cl <= fifo_len_) return; + if (F + cl <= scfg_.fifo_len) return; - int pop = std::max(update_period_, cl - fifo_len_ + F); + int pop = std::max(scfg_.update_period, cl - scfg_.fifo_len + F); pop = std::min(pop, F + cl); if (!use_learnable_sil_emb_) { @@ -183,7 +289,7 @@ void StreamingDiarization::streaming_update(const std::vector& chunk_emb, drop_rows(fifo_, d, pop); drop_rows(fifo_preds_, ns, pop); - if ((int)(spkcache_.size() / d) > spkcache_len_) { + if ((int)(spkcache_.size() / d) > scfg_.spkcache_len) { compress_spkcache(); spkcache_compressed_ = true; } @@ -193,7 +299,7 @@ void StreamingDiarization::streaming_update(const std::vector& chunk_emb, void StreamingDiarization::compress_spkcache() { const int d = d_model_, ns = n_spk_; const int n = (int)(spkcache_preds_.size() / ns); - const int per_spk = spkcache_len_ / ns - sil_frames_per_spk_; + const int per_spk = scfg_.spkcache_len / ns - sil_frames_per_spk_; const int strong_k = (int)std::floor(per_spk * strong_boost_rate_); const int weak_k = (int)std::floor(per_spk * weak_boost_rate_); const int min_pos = (int)std::floor(per_spk * min_pos_scores_rate_); @@ -230,7 +336,7 @@ void StreamingDiarization::compress_spkcache() { // Boost frames newly added since the last compression. if (scores_boost_latest_ > 0.0f) - for (size_t i = (size_t)spkcache_len_ * ns; i < sc.size(); ++i) sc[i] += scores_boost_latest_; + for (size_t i = (size_t)scfg_.spkcache_len * ns; i < sc.size(); ++i) sc[i] += scores_boost_latest_; // _boost_topk_scores: add -scale*log(0.5) to each speaker's top-k frames. auto boost = [&](int k, float scale) { @@ -259,7 +365,7 @@ void StreamingDiarization::compress_spkcache() { for (int s = 0; s < ns; ++s) for (int t = 0; t < n_tot; ++t) flat[(size_t)s * n_tot + t] = {t < n ? sc[(size_t)t * ns + s] : kInf, s * n_tot + t}; - const int K = spkcache_len_; + const int K = scfg_.spkcache_len; std::nth_element(flat.begin(), flat.begin() + (K - 1), flat.end(), [](const std::pair& a, const std::pair& b) { return a.first != b.first ? a.first > b.first : a.second < b.second; diff --git a/src/diarization_streaming.hpp b/src/diarization_streaming.hpp index 21e8007..2b269dc 100644 --- a/src/diarization_streaming.hpp +++ b/src/diarization_streaming.hpp @@ -13,6 +13,31 @@ struct StreamingSpeakerSegment { float end; }; +// Streaming configuration, in ENCODER frames (80 ms), as in NeMo's +// SortformerModules. Input latency = (chunk_len + right_context) * 80 ms. +struct DiarStreamConfig { + int spkcache_len = 264; // speaker cache size + int fifo_len = 0; // FIFO of recent frames kept before the chunk + int chunk_len = 264; // frames diarized per step + int left_context = 0; // past frames re-encoded with the chunk + int right_context = 0; // look-ahead frames encoded with the chunk + int update_period = 264; // frames moved FIFO -> cache per update + + // The checkpoint's own configuration (NeMo diarize() default). + static DiarStreamConfig from_model(const ParakeetConfig& cfg); +}; + +// Streaming latency presets from the Nemotron-3-Diarization model card. +enum class DiarLatency { + Model = 0, // the checkpoint's configuration (21.12 s for Nemotron-3) + Low = 1, // 1.04 s: chunk 9, look-ahead 4, FIFO 264 + VeryLow = 2, // 0.64 s: chunk 6, look-ahead 2, FIFO 264 + UltraLow = 3, // 0.32 s: chunk 3, look-ahead 1, FIFO 264 +}; + +// The configuration for a preset. Model returns from_model(cfg). +DiarStreamConfig diar_stream_config(DiarLatency latency, const ParakeetConfig& cfg); + // StreamingDiarization — NeMo Sortformer cache-aware streaming ("AOSC"), // synchronous mode (SortformerEncLabelModel.forward_streaming_step + // SortformerModules.streaming_update), for nvidia/Nemotron-3-Diarization. @@ -24,43 +49,58 @@ struct StreamingSpeakerSegment { // FIFO -> speaker-cache update and the score-based cache compression that // keeps the speaker identities stable across chunks. // -// chunk_len / spkcache_len / fifo_len / spkcache_update_period come from the -// GGUF and are in ENCODER frames (80 ms), as in NeMo. With the Nemotron-3 -// config a chunk is 264 encoder frames = 2112 mel frames = 21.12 s. +// Each chunk is encoded together with `left_context` past and `right_context` +// future frames, like NeMo's streaming_feat_loader; only the chunk itself is +// output and enters the cache update. A chunk runs once its look-ahead has +// arrived, so the input latency is (chunk_len + right_context) * 80 ms. // // The mel must be the un-normalized log-mel of the stream (NeMo does not // peak-normalize in streaming mode), e.g. from pk::StreamingMel. class StreamingDiarization { public: + // Uses the checkpoint's own streaming configuration. explicit StreamingDiarization(const ModelLoader& ml); + StreamingDiarization(const ModelLoader& ml, const DiarStreamConfig& cfg); void reset(); - // Feed the next chunk: row-major [n_mels, n_frames] (mel[m*n_frames + t]), - // 0 < n_frames <= chunk_mel_frames(). Only the final chunk may be short. - // Returns the speaker segments that ENDED in this chunk; with is_last, - // every still-open segment is closed at the end of the stream. - std::vector feed_mel_chunk( + // Append mel frames, row-major [n_mels, n_frames] (mel[m*n_frames + t]), + // any count including 0, and diarize every chunk whose look-ahead is now + // available. is_last marks the end of the stream: the remaining frames + // are diarized with whatever look-ahead exists, and every open segment + // is closed. Returns the speaker segments that ENDED during this call. + std::vector feed_mel( const std::vector& mel, int n_mels, int n_frames, bool is_last); - // Speaker probabilities of the last fed chunk, speaker-major - // [n_speakers, last_chunk_frames()] (one frame per mel frame, 10 ms). - const std::vector& last_chunk_probs() const { return last_probs_; } - int last_chunk_frames() const { return last_frames_; } + // Speaker probabilities output by the last feed_mel call, speaker-major + // [n_speakers, last_frames()] (one frame per mel frame, 10 ms). + const std::vector& last_probs() const { return last_probs_; } + int last_frames() const { return last_frames_; } - // Segments that are still active at the current end of the stream, with - // `end` set to the stream time consumed so far. + // Segments still active at the end of the diarized audio, with `end` at + // the diarized time (frames_done()). std::vector open_segments() const; - int chunk_mel_frames() const { return chunk_len_ * subsampling_; } + const DiarStreamConfig& config() const { return scfg_; } + // Mel frames per chunk, and how many must arrive before a chunk runs. + int chunk_mel_frames() const { return scfg_.chunk_len * subsampling_; } + int latency_mel_frames() const { + return (scfg_.chunk_len + scfg_.right_context) * subsampling_; + } int n_mels() const { return n_mels_; } int n_speakers() const { return n_spk_; } float frame_sec() const { return frame_sec_; } - // Mel frames consumed so far. + // Mel frames diarized so far (output frames, excluding look-ahead). long long frames_done() const { return frames_done_; } + // Mel frames received so far. + long long frames_in() const { return frames_in_; } private: - void streaming_update(const std::vector& chunk_emb, int chunk_frames, + // One NeMo forward_streaming_step over the window [stt - left, end + right) + // of the stream; appends the chunk's output to last_probs_. + void step(long long stt, long long end, int left, int right, bool final_chunk, + std::vector& closed); + void streaming_update(const std::vector& emb, int lc, int cl, const std::vector& preds, int spkcache_frames, int fifo_frames); void compress_spkcache(); @@ -72,7 +112,8 @@ class StreamingDiarization { DiarizationHead head_; int d_model_, n_spk_, subsampling_, upsample_, n_mels_; - int chunk_len_, spkcache_len_, fifo_len_, update_period_, sil_frames_per_spk_; + DiarStreamConfig scfg_; + int sil_frames_per_spk_; float frame_sec_, onset_, offset_; float sil_threshold_, pred_score_threshold_, scores_boost_latest_; float strong_boost_rate_, weak_boost_rate_, min_pos_scores_rate_; @@ -87,8 +128,14 @@ class StreamingDiarization { std::vector mean_sil_emb_; // [d_model] long long n_sil_frames_ = 0; + // Input buffer: frame-major mel frames [buf_start_, frames_in_). + std::vector buf_; + long long buf_start_ = 0; + long long frames_in_ = 0; + long long next_stt_ = 0; // start frame of the next chunk + // Output state. - long long frames_done_ = 0; // mel frames consumed + long long frames_done_ = 0; // mel frames diarized std::vector last_probs_; int last_frames_ = 0; std::vector active_; // per speaker diff --git a/src/parakeet_capi.cpp b/src/parakeet_capi.cpp index 14fcc45..510a701 100644 --- a/src/parakeet_capi.cpp +++ b/src/parakeet_capi.cpp @@ -1124,76 +1124,85 @@ struct parakeet_diar_stream { parakeet_ctx* ctx = nullptr; std::unique_ptr sd; std::unique_ptr mel; - std::vector pending; // mel frames not yet diarized, frame-major [t][n_mels] long long samples_in = 0; // PCM samples fed so far bool finished = false; }; namespace { -// Append feat-major [n_mels, n] mel frames to a frame-major buffer. -void push_frames(std::vector& dst, const std::vector& fm, int n_mels, int n) { - const size_t base = dst.size(); - dst.resize(base + (size_t)n * n_mels); - for (int m = 0; m < n_mels; ++m) - for (int t = 0; t < n; ++t) dst[base + (size_t)t * n_mels + m] = fm[(size_t)m * n + t]; -} - -// Feed PCM to the stream's mel front end and run every full diarization chunk -// (and, with is_last, the tail). Closed segments are appended to `segs`. -// Returns the number of chunks run. -int diar_stream_advance(parakeet_diar_stream* s, const float* pcm, int n, bool is_last, - std::vector& segs) { +// Feed PCM to the stream's mel front end and the diarizer, which runs every +// chunk whose look-ahead has arrived (and, with is_last, the rest). Closed +// segments are appended to `segs`. Returns the mel frames diarized by this call. +long long diar_stream_advance(parakeet_diar_stream* s, const float* pcm, int n, bool is_last, + std::vector& segs) { const int n_mels = s->sd->n_mels(); + const long long done_before = s->sd->frames_done(); + std::vector mel; int nf = 0; if (n > 0) { - std::vector fm = s->mel->feed(pcm, n, nf); - push_frames(s->pending, fm, n_mels, nf); + mel = s->mel->feed(pcm, n, nf); s->samples_in += n; } if (is_last) { - std::vector fm = s->mel->finalize(nf); - push_frames(s->pending, fm, n_mels, nf); - // NeMo keeps floor(S / hop) frames; the centered STFT emits one more. + int nt = 0; + std::vector tail = s->mel->finalize(nt); + // Join the two feat-major blocks, then keep floor(S / hop) frames in + // total like NeMo (the centered STFT emits one more). const long long valid = s->samples_in / (long long)s->ctx->diar->config().hop_length; - const long long have = s->sd->frames_done() + (long long)(s->pending.size() / n_mels); - if (have > valid) s->pending.resize(s->pending.size() - (size_t)(have - valid) * n_mels); - } - const int cm = s->sd->chunk_mel_frames(); - int chunks = 0; - for (;;) { - const int avail = (int)(s->pending.size() / n_mels); - const bool last = is_last && avail <= cm; - if (avail < cm && !(last && avail > 0)) break; - const int take = std::min(avail, cm); - std::vector chunk((size_t)n_mels * take); - for (int t = 0; t < take; ++t) - for (int m = 0; m < n_mels; ++m) - chunk[(size_t)m * take + t] = s->pending[(size_t)t * n_mels + m]; - s->pending.erase(s->pending.begin(), s->pending.begin() + (size_t)take * n_mels); - auto closed = s->sd->feed_mel_chunk(chunk, n_mels, take, last); - segs.insert(segs.end(), closed.begin(), closed.end()); - ++chunks; - if (last) break; + const int keep = (int)std::max(0LL, std::min(nf + nt, valid - s->sd->frames_in())); + std::vector joined((size_t)n_mels * keep); + for (int m = 0; m < n_mels; ++m) + for (int t = 0; t < keep; ++t) + joined[(size_t)m * keep + t] = t < nf ? mel[(size_t)m * nf + t] + : tail[(size_t)m * nt + (t - nf)]; + mel.swap(joined); + nf = keep; + s->finished = true; } - if (is_last && chunks == 0 && s->sd->frames_done() > 0) { - // Stream length was an exact multiple of the chunk: close open segments. - auto open = s->sd->open_segments(); - segs.insert(segs.end(), open.begin(), open.end()); + auto closed = s->sd->feed_mel(mel, n_mels, nf, is_last); + segs.insert(segs.end(), closed.begin(), closed.end()); + return s->sd->frames_done() - done_before; +} + +pk::DiarLatency latency_from_int(int latency) { + switch (latency) { + case PARAKEET_DIAR_LATENCY_LOW: return pk::DiarLatency::Low; + case PARAKEET_DIAR_LATENCY_VERY_LOW: return pk::DiarLatency::VeryLow; + case PARAKEET_DIAR_LATENCY_ULTRA_LOW: return pk::DiarLatency::UltraLow; + default: return pk::DiarLatency::Model; } - if (is_last) s->finished = true; - return chunks; +} + +// Copy segments into a malloc'd C array (NULL when empty). False on OOM. +bool to_c_segments(const std::vector& segs, + parakeet_diar_segment** out, int* n_out) { + *out = nullptr; + *n_out = 0; + if (segs.empty()) return true; + auto* r = static_cast(std::malloc(segs.size() * sizeof(parakeet_diar_segment))); + if (!r) return false; + for (size_t i = 0; i < segs.size(); ++i) r[i] = {segs[i].speaker, segs[i].start, segs[i].end}; + *out = r; + *n_out = (int)segs.size(); + return true; } } // namespace -extern "C" parakeet_diar_stream* parakeet_capi_diarize_stream_begin(parakeet_ctx* diar_ctx) { +extern "C" parakeet_diar_stream* parakeet_capi_diarize_stream_begin_latency(parakeet_ctx* diar_ctx, + int latency) { if (!require_diar(diar_ctx)) return nullptr; + if (latency < PARAKEET_DIAR_LATENCY_MODEL || latency > PARAKEET_DIAR_LATENCY_ULTRA_LOW) { + diar_ctx->last_error = "unknown diarization latency mode"; + return nullptr; + } try { + const pk::ModelLoader& ml = diar_ctx->diar->loader(); auto* s = new parakeet_diar_stream(); s->ctx = diar_ctx; - s->sd = std::make_unique(diar_ctx->diar->loader()); - s->mel = std::make_unique(diar_ctx->diar->loader()); + s->sd = std::make_unique( + ml, pk::diar_stream_config(latency_from_int(latency), ml.config())); + s->mel = std::make_unique(ml); diar_ctx->last_error.clear(); return s; } catch (const std::exception& e) { @@ -1204,9 +1213,28 @@ extern "C" parakeet_diar_stream* parakeet_capi_diarize_stream_begin(parakeet_ctx return nullptr; } +extern "C" parakeet_diar_stream* parakeet_capi_diarize_stream_begin(parakeet_ctx* diar_ctx) { + return parakeet_capi_diarize_stream_begin_latency(diar_ctx, PARAKEET_DIAR_LATENCY_MODEL); +} + extern "C" int parakeet_capi_diarize_stream_chunk_samples(parakeet_diar_stream* s) { if (!s) return 0; - return s->sd->chunk_mel_frames() * (int)s->ctx->diar->config().hop_length; + return s->sd->latency_mel_frames() * (int)s->ctx->diar->config().hop_length; +} + +extern "C" float parakeet_capi_diarize_stream_time(parakeet_diar_stream* s) { + if (!s) return 0.0f; + return (float)(s->sd->frames_done() * s->sd->frame_sec()); +} + +extern "C" int parakeet_capi_diarize_stream_active(parakeet_diar_stream* s, + parakeet_diar_segment** out, int* n_out) { + if (!s || !out || !n_out) return 1; + if (!to_c_segments(s->sd->open_segments(), out, n_out)) { + s->ctx->last_error = "out of memory"; + return 1; + } + return 0; } extern "C" int parakeet_capi_diarize_stream_feed(parakeet_diar_stream* s, const float* pcm, @@ -1220,13 +1248,7 @@ extern "C" int parakeet_capi_diarize_stream_feed(parakeet_diar_stream* s, const try { std::vector segs; diar_stream_advance(s, pcm, n_samples, is_last != 0, segs); - if (!segs.empty()) { - auto* r = static_cast(std::malloc(segs.size() * sizeof(parakeet_diar_segment))); - if (!r) { s->ctx->last_error = "out of memory"; return 1; } - for (size_t i = 0; i < segs.size(); ++i) r[i] = {segs[i].speaker, segs[i].start, segs[i].end}; - *out = r; - *n_out = (int)segs.size(); - } + if (!to_c_segments(segs, out, n_out)) { s->ctx->last_error = "out of memory"; return 1; } s->ctx->last_error.clear(); return 0; } catch (const std::exception& e) { @@ -1269,10 +1291,11 @@ std::string word_key(const std::string& w) { } // namespace -extern "C" parakeet_sas_stream* parakeet_capi_sas_stream_begin(parakeet_ctx* asr_ctx, - parakeet_ctx* diar_ctx) { +extern "C" parakeet_sas_stream* parakeet_capi_sas_stream_begin_latency(parakeet_ctx* asr_ctx, + parakeet_ctx* diar_ctx, + int latency) { if (!require_asr(asr_ctx) || !require_diar(diar_ctx)) return nullptr; - parakeet_diar_stream* d = parakeet_capi_diarize_stream_begin(diar_ctx); + parakeet_diar_stream* d = parakeet_capi_diarize_stream_begin_latency(diar_ctx, latency); if (!d) return nullptr; auto* s = new (std::nothrow) parakeet_sas_stream(); if (!s) { parakeet_capi_diarize_stream_free(d); return nullptr; } @@ -1281,6 +1304,11 @@ extern "C" parakeet_sas_stream* parakeet_capi_sas_stream_begin(parakeet_ctx* asr return s; } +extern "C" parakeet_sas_stream* parakeet_capi_sas_stream_begin(parakeet_ctx* asr_ctx, + parakeet_ctx* diar_ctx) { + return parakeet_capi_sas_stream_begin_latency(asr_ctx, diar_ctx, PARAKEET_DIAR_LATENCY_MODEL); +} + extern "C" int parakeet_capi_sas_stream_feed(parakeet_sas_stream* s, const float* pcm, int n_samples, int is_last, parakeet_sas_result** out, int* n_out) { @@ -1292,10 +1320,15 @@ extern "C" int parakeet_capi_sas_stream_feed(parakeet_sas_stream* s, const float parakeet_ctx* failed = s->diar->ctx; try { std::vector closed; - const int chunks = diar_stream_advance(s->diar, pcm, n_samples, is_last != 0, closed); + const long long advanced = diar_stream_advance(s->diar, pcm, n_samples, is_last != 0, closed); for (const auto& c : closed) s->segs.push_back({c.speaker, c.start, c.end}); if (n_samples > 0) s->audio.insert(s->audio.end(), pcm, pcm + n_samples); - if (chunks == 0 && !is_last) return 0; + if (advanced == 0 && !is_last) return 0; + + // Offline ASR on a short window loses words, so wait until enough + // uncommitted, diarized audio has built up (it bounds how often the + // text commits, not the speaker latency). + constexpr double kSasMinWindowSec = 4.0; // Diarized audio ends at frames_done; transcribe the uncommitted span. const double hop_sec = (double)s->diar->ctx->diar->config().hop_length / 16000.0; @@ -1303,6 +1336,7 @@ extern "C" int parakeet_capi_sas_stream_feed(parakeet_sas_stream* s, const float size_t span = is_last ? s->audio.size() : std::min(s->audio.size(), (size_t)std::max(0.0, (diar_end - s->commit_sec) * 16000.0)); + if (!is_last && span < (size_t)(kSasMinWindowSec * 16000.0)) return 0; failed = s->asr; std::vector words; if (span > 0) { @@ -1319,8 +1353,9 @@ extern "C" int parakeet_capi_sas_stream_feed(parakeet_sas_stream* s, const float const double limit = (double)span / 16000.0 - kSasRightContextSec; keep = 0; while (keep < words.size() && words[keep].end <= limit) ++keep; - next_commit = s->commit_sec + (keep < words.size() ? words[keep].start - : std::max(0.0, limit)); + // Resume right after the last committed word: audio the ASR + // skipped this time is heard again with more context. + next_commit = s->commit_sec + (keep > 0 ? words[keep - 1].end : 0.0); } std::vector committed(words.begin(), words.begin() + keep); for (auto& w : committed) { w.start += (float)s->commit_sec; w.end += (float)s->commit_sec; } diff --git a/tests/test_combined_offline.cpp b/tests/test_combined_offline.cpp index 0e7fd54..3311d3e 100644 --- a/tests/test_combined_offline.cpp +++ b/tests/test_combined_offline.cpp @@ -295,9 +295,60 @@ int main() { } } - // 4. Streaming SAS. + // 3b. Low-latency streaming diarization: diarized time trails the audio by + // at most the mode's latency, and _active names the current speaker. { - parakeet_sas_stream* ss = parakeet_capi_sas_stream_begin(asr, diar); + parakeet_diar_stream* ds = parakeet_capi_diarize_stream_begin_latency(diar, PARAKEET_DIAR_LATENCY_LOW); + CHECK(ds != nullptr, "diarize_stream_begin_latency: %s", parakeet_capi_last_error(diar)); + const int latency = ds ? parakeet_capi_diarize_stream_chunk_samples(ds) : 0; + CHECK(latency == 16640, "LOW latency is %d samples, expected 16640 (1.04 s)", latency); + std::vector streamed; + int active_checks = 0; + float prev_t = 0.0f; + for (int lo = 0; ds && lo < n; lo += 1600) { + const int len = std::min(1600, n - lo); + parakeet_diar_segment* segs = nullptr; + int ns = 0; + const bool last = lo + len >= n; + CHECK(parakeet_capi_diarize_stream_feed(ds, pcm.data() + lo, len, last, &segs, &ns) == 0, + "feed: %s", parakeet_capi_last_error(diar)); + streamed.insert(streamed.end(), segs, segs + ns); + parakeet_capi_free_diar_segments(segs); + if (last) break; + const float fed = (lo + len) / 16000.0f; + const float t = parakeet_capi_diarize_stream_time(ds); + CHECK(fed - t <= 1.04f + 0.24f + 1e-3f, "diarized %.2f s of %.2f s fed", t, fed); + // Mid-utterance of each speaker the current speaker must be known, + // checked when the diarized time first passes each probe (it + // advances in 0.72 s steps). + for (const float probe : {3.0f, 9.0f, 16.0f, 22.0f}) { + if (prev_t < probe && t >= probe) { + parakeet_diar_segment* act = nullptr; + int na = 0; + CHECK(parakeet_capi_diarize_stream_active(ds, &act, &na) == 0, "active"); + const int want = (probe == 3.0f || probe == 16.0f) ? 0 : 1; + CHECK(na == 1 && act[0].speaker == want && act[0].end == t, + "at %.2f s: %d active, speaker %d (want %d)", t, na, + na ? act[0].speaker : -1, want); + parakeet_capi_free_diar_segments(act); + ++active_checks; + } + } + prev_t = t; + } + parakeet_capi_diarize_stream_free(ds); + std::vector spk; + std::sort(streamed.begin(), streamed.end(), [](const auto& a, const auto& b) { return a.start < b.start; }); + for (const auto& g : streamed) spk.push_back(g.speaker); + std::printf("low-latency diarization: %zu segments, turns %s, %d active checks\n", + streamed.size(), show(turns(spk)).c_str(), active_checks); + CHECK(turns(spk) == expected_turns, "low-latency turns %s", show(turns(spk)).c_str()); + CHECK(active_checks == 4, "only %d of 4 active checks ran", active_checks); + } + + // 4. Streaming SAS (checkpoint config, then 1.04 s latency). + for (const int latency : {PARAKEET_DIAR_LATENCY_MODEL, PARAKEET_DIAR_LATENCY_LOW}) { + parakeet_sas_stream* ss = parakeet_capi_sas_stream_begin_latency(asr, diar, latency); CHECK(ss != nullptr, "sas_stream_begin"); std::vector spk; int words = 0; @@ -317,7 +368,7 @@ int main() { parakeet_capi_free_sas_results(r, nr); } parakeet_capi_sas_stream_free(ss); - std::printf("streaming SAS: %d words, turns %s\n %s\n", words, + std::printf("streaming SAS (latency mode %d): %d words, turns %s\n %s\n", latency, words, show(turns(spk)).c_str(), text.c_str()); CHECK(turns(spk) == expected_turns, "streaming turns %s", show(turns(spk)).c_str()); CHECK(std::abs(words - n_words_offline) <= 3, "streaming words %d vs offline %d", diff --git a/tests/test_streaming_diarization.cpp b/tests/test_streaming_diarization.cpp index bff39c9..7beaf92 100644 --- a/tests/test_streaming_diarization.cpp +++ b/tests/test_streaming_diarization.cpp @@ -1,18 +1,19 @@ // Streaming diarization accuracy vs NeMo cache-aware streaming -// (SortformerEncLabelModel with streaming_mode=True, the model's own chunk and -// speaker-cache config). +// (SortformerEncLabelModel with streaming_mode=True). // -// Feeds the baseline clip through pk::StreamingDiarization two ways and checks -// both against NeMo's streaming output: -// A. mel chunks cut from the whole-clip mel (exactly what NeMo does) -// B. pk::StreamingMel fed with small PCM pieces (the live-audio path the -// C-API uses), re-chunked to the model's chunk size +// For the checkpoint's own configuration and every model-card latency mode +// present in the baseline (stream_probs_; scripts/gen_diar_baseline.py), +// feeds the clip through pk::StreamingDiarization two ways and checks both +// against NeMo: +// A. the whole-clip mel in one feed (NeMo computes the mel up front) +// B. pk::StreamingMel fed with 100 ms PCM pieces, each piece's frames fed as +// they arrive (the live path the C-API uses) // Checks: per-frame probabilities (max/mean abs diff) and the segments // (count, speaker, boundaries within 20 ms). // -// Env: PARAKEET_TEST_DIAR_GGUF + PARAKEET_TEST_BASELINE_DIAR -// (scripts/gen_diar_baseline.py); PARAKEET_TEST_DIAR_PROB_TOL as in -// test_diarization_accuracy. Skips (77) when unset. +// Env: PARAKEET_TEST_DIAR_GGUF + PARAKEET_TEST_BASELINE_DIAR; optional +// PARAKEET_TEST_DIAR_PROB_TOL (default 0.02) and PARAKEET_TEST_DIAR_VERBOSE +// (per-10 s max diff). Skips (77) when the required ones are unset. #include "diarization.hpp" #include "diarization_streaming.hpp" #include "mel.hpp" @@ -22,6 +23,7 @@ #include #include #include +#include #include namespace { @@ -31,30 +33,69 @@ struct Run { std::vector segs; }; -// Feed [n_mels, T] mel through the streaming diarizer in model-sized chunks. -Run stream_mel(pk::StreamingDiarization& sd, const std::vector& mel, int n_mels, int T) { - Run r; - const int ns = sd.n_speakers(), cm = sd.chunk_mel_frames(); - r.probs.assign((size_t)ns * T, 0.0f); - sd.reset(); - for (int lo = 0; lo < T; lo += cm) { - const int n = std::min(cm, T - lo); - std::vector chunk((size_t)n_mels * n); - for (int m = 0; m < n_mels; ++m) - std::copy_n(mel.begin() + (size_t)m * T + lo, n, chunk.begin() + (size_t)m * n); - auto segs = sd.feed_mel_chunk(chunk, n_mels, n, lo + n >= T); - r.segs.insert(r.segs.end(), segs.begin(), segs.end()); - for (int s = 0; s < ns; ++s) - std::copy_n(sd.last_chunk_probs().begin() + (size_t)s * n, n, - r.probs.begin() + (size_t)s * T + lo); - } +// Collect one feed_mel call's output into the run. +void collect(pk::StreamingDiarization& sd, Run& r, long long& done, int T, + const std::vector& segs) { + const int ns = sd.n_speakers(), n = sd.last_frames(); + for (int s = 0; s < ns; ++s) + for (int t = 0; t < n && done + t < T; ++t) + r.probs[(size_t)s * T + done + t] = sd.last_probs()[(size_t)s * n + t]; + done += n; + r.segs.insert(r.segs.end(), segs.begin(), segs.end()); +} + +void sort_segs(Run& r) { std::sort(r.segs.begin(), r.segs.end(), [](const auto& a, const auto& b) { return a.start != b.start ? a.start < b.start : a.speaker < b.speaker; }); +} + +// A: whole-clip mel [n_mels, T] in one feed. +Run whole_clip(pk::StreamingDiarization& sd, const std::vector& mel, int n_mels, int T) { + Run r; + r.probs.assign((size_t)sd.n_speakers() * T, 0.0f); + long long done = 0; + sd.reset(); + collect(sd, r, done, T, sd.feed_mel(mel, n_mels, T, true)); + sort_segs(r); + return r; +} + +// B: live, StreamingMel over 100 ms PCM pieces; frames beyond T are dropped +// like NeMo's floor(S / hop) length. +Run live(pk::StreamingDiarization& sd, const pk::ModelLoader& ml, + const std::vector& audio, int T) { + Run r; + r.probs.assign((size_t)sd.n_speakers() * T, 0.0f); + long long done = 0, fed = 0; + sd.reset(); + pk::StreamingMel sm(ml); + const int n_mels = sd.n_mels(); + auto feed = [&](std::vector fm, int nf, bool last) { + const int keep = (int)std::max(0LL, std::min(nf, T - fed)); + if (keep < nf) { + std::vector cut((size_t)n_mels * keep); + for (int m = 0; m < n_mels; ++m) + std::copy_n(fm.begin() + (size_t)m * nf, keep, cut.begin() + (size_t)m * keep); + fm.swap(cut); + } + fed += keep; + collect(sd, r, done, T, sd.feed_mel(fm, n_mels, keep, last)); + }; + for (size_t lo = 0; lo < audio.size(); lo += 1600) { + const int n = (int)std::min(1600, audio.size() - lo); + int nf = 0; + std::vector fm = sm.feed(audio.data() + lo, n, nf); + feed(std::move(fm), nf, false); + } + int nf = 0; + std::vector tail = sm.finalize(nf); + feed(std::move(tail), nf, true); + sort_segs(r); return r; } -int check(const char* label, const Run& r, const std::vector& ref_probs, +int check(const std::string& label, const Run& r, const std::vector& ref_probs, const std::vector& ref_segs, float tol) { int fails = 0; double max_d = 0.0, sum_d = 0.0; @@ -64,8 +105,8 @@ int check(const char* label, const Run& r, const std::vector& ref_probs, sum_d += d; } const double mean_d = sum_d / ref_probs.size(); + std::printf("[%s] probs: max_diff=%.5f mean_diff=%.6f\n", label.c_str(), max_d, mean_d); if (std::getenv("PARAKEET_TEST_DIAR_VERBOSE")) { - // Per-10 s max diff, to localize divergence (e.g. after cache compression). const size_t T = ref_probs.size() / 8; for (size_t lo = 0; lo < T; lo += 1000) { double m = 0.0; @@ -75,30 +116,32 @@ int check(const char* label, const Run& r, const std::vector& ref_probs, std::printf(" frames %zu..%zu max_diff %.4f\n", lo, std::min(T, lo + 1000), m); } } - std::printf("[%s] probs: max_diff=%.5f mean_diff=%.6f\n", label, max_d, mean_d); - if (max_d > tol || mean_d > 2e-3) { std::fprintf(stderr, "[%s] FAIL: probabilities\n", label); ++fails; } + if (max_d > tol || mean_d > 2e-3) { + std::fprintf(stderr, "[%s] FAIL: probabilities\n", label.c_str()); + ++fails; + } const size_t n_ref = ref_segs.size() / 3; - std::printf("[%s] segments: ours %zu, NeMo %zu\n", label, r.segs.size(), n_ref); - if (r.segs.size() != n_ref) { - std::fprintf(stderr, "[%s] FAIL: segment count\n", label); - return fails + 1; - } - // NeMo's rows are sorted by (start, speaker) the same way. + std::printf("[%s] segments: ours %zu, NeMo %zu\n", label.c_str(), r.segs.size(), n_ref); std::vector order(n_ref); for (size_t i = 0; i < n_ref; ++i) order[i] = i; std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { return ref_segs[a * 3 + 1] != ref_segs[b * 3 + 1] ? ref_segs[a * 3 + 1] < ref_segs[b * 3 + 1] : ref_segs[a * 3] < ref_segs[b * 3]; }); - for (size_t i = 0; i < n_ref; ++i) { - const float* g = &ref_segs[order[i] * 3]; - const auto& o = r.segs[i]; - const bool ok = o.speaker == (int)g[0] && std::fabs(o.start - g[1]) <= 0.02f && - std::fabs(o.end - g[2]) <= 0.02f; - std::printf(" %s spk%d %6.2f-%6.2f NeMo spk%d %6.2f-%6.2f\n", ok ? "ok " : "DIFF", - o.speaker, o.start, o.end, (int)g[0], g[1], g[2]); - if (!ok) ++fails; + for (size_t i = 0; i < std::max(n_ref, r.segs.size()); ++i) { + const bool hr = i < n_ref, ho = i < r.segs.size(); + const float* g = hr ? &ref_segs[order[i] * 3] : nullptr; + const bool ok = hr && ho && r.segs[i].speaker == (int)g[0] && + std::fabs(r.segs[i].start - g[1]) <= 0.02f && + std::fabs(r.segs[i].end - g[2]) <= 0.02f; + if (!ok) { + std::printf(" DIFF spk%d %6.2f-%6.2f NeMo spk%d %6.2f-%6.2f\n", + ho ? r.segs[i].speaker : -1, ho ? r.segs[i].start : 0.f, + ho ? r.segs[i].end : 0.f, hr ? (int)g[0] : -1, + hr ? g[1] : 0.f, hr ? g[2] : 0.f); + ++fails; + } } return fails; } @@ -118,58 +161,50 @@ int main() { auto m = pk::DiarizationModel::load(gguf); if (!m) { std::fprintf(stderr, "load failed: %s\n", gguf); return 1; } + const pk::ModelLoader& ml = m->loader(); - std::vector audio, ref_probs, ref_segs; + std::vector audio; std::vector shape; if (!pktest::load_baseline(base, "audio", audio, shape)) return 1; - if (!pktest::load_baseline(base, "stream_probs", ref_probs, shape)) return 1; - const int T = (int)shape[1]; - if (!pktest::load_baseline(base, "stream_segs", ref_segs, shape)) return 1; - pk::StreamingDiarization sd(m->loader()); - const int n_mels = sd.n_mels(); - int fails = 0; - - // A. Whole-clip mel (no peak normalization in streaming mode), trimmed to - // floor(S / hop) frames like NeMo. - { - std::vector mel; - int nm = 0, Tm = 0; - m->mel().compute(audio, mel, nm, Tm); - if (Tm < T) { std::fprintf(stderr, "mel too short: %d < %d\n", Tm, T); return 1; } - std::vector trimmed((size_t)nm * T); - for (int i = 0; i < nm; ++i) - std::copy_n(mel.begin() + (size_t)i * Tm, T, trimmed.begin() + (size_t)i * T); - fails += check("whole-clip mel", stream_mel(sd, trimmed, nm, T), ref_probs, ref_segs, tol); - } + // Whole-clip mel (no peak normalization in streaming mode), trimmed to + // floor(S / hop) frames like NeMo. + std::vector mel_full; + int n_mels = 0, Tm = 0; + m->mel().compute(audio, mel_full, n_mels, Tm); + const int T = std::min(Tm, (int)(audio.size() / ml.config().hop_length)); + std::vector mel((size_t)n_mels * T); + for (int i = 0; i < n_mels; ++i) + std::copy_n(mel_full.begin() + (size_t)i * Tm, T, mel.begin() + (size_t)i * T); - // B. Incremental mel from 100 ms PCM pieces. - { - pk::StreamingMel sm(m->loader()); - std::vector> cols; // per-frame mel columns - auto take = [&](const std::vector& fm, int n) { - for (int t = 0; t < n; ++t) { - std::vector col(n_mels); - for (int i = 0; i < n_mels; ++i) col[i] = fm[(size_t)i * n + t]; - cols.push_back(std::move(col)); - } - }; - for (size_t lo = 0; lo < audio.size(); lo += 1600) { - const int n = (int)std::min(1600, audio.size() - lo); - int nf = 0; - auto fm = sm.feed(audio.data() + lo, n, nf); - take(fm, nf); + const struct { const char* key; pk::DiarLatency latency; } modes[] = { + {"", pk::DiarLatency::Model}, + {"_low", pk::DiarLatency::Low}, + {"_very_low", pk::DiarLatency::VeryLow}, + {"_ultra_low", pk::DiarLatency::UltraLow}, + }; + int fails = 0, checked = 0; + for (const auto& mode : modes) { + std::vector ref_probs, ref_segs; + std::vector ps; + const std::string pk_ = std::string("stream_probs") + mode.key; + const std::string sk_ = std::string("stream_segs") + mode.key; + if (!pktest::load_baseline(base, pk_, ref_probs, ps)) { + std::printf("[%s] not in the baseline; skipped\n", pk_.c_str()); + continue; } - int nf = 0; - auto tail = sm.finalize(nf); - take(tail, nf); - if ((int)cols.size() < T) { std::fprintf(stderr, "stream mel too short\n"); return 1; } - std::vector mel((size_t)n_mels * T); - for (int t = 0; t < T; ++t) - for (int i = 0; i < n_mels; ++i) mel[(size_t)i * T + t] = cols[t][i]; - fails += check("StreamingMel", stream_mel(sd, mel, n_mels, T), ref_probs, ref_segs, tol); + if (!pktest::load_baseline(base, sk_, ref_segs, shape)) return 1; + if ((int)ps[1] != T) { + std::fprintf(stderr, "[%s] length %lld != %d\n", pk_.c_str(), (long long)ps[1], T); + return 1; + } + pk::StreamingDiarization sd(ml, pk::diar_stream_config(mode.latency, ml.config())); + const std::string name = mode.key[0] ? std::string(mode.key + 1) : "model"; + fails += check(name + " whole-clip", whole_clip(sd, mel, n_mels, T), ref_probs, ref_segs, tol); + fails += check(name + " live", live(sd, ml, audio, T), ref_probs, ref_segs, tol); + ++checked; } - + if (checked == 0) { std::fprintf(stderr, "no streaming baseline found\n"); return 1; } std::printf(fails ? "test_streaming_diarization: FAIL\n" : "test_streaming_diarization: PASS\n"); return fails ? 1 : 0; } From 48c43aa249ddcc90796165c217b195accf0a425a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 28 Sep 2026 09:38:01 +0000 Subject: [PATCH 2/2] docs(diarization): document the latency modes and correct the Q8_0 note Adds the latency modes, their parity with NeMo and their speed on a 12.3 min clip. Q8_0 was described as segment-exact; on a 31.5 s clip a frame at the 0.5 threshold splits one segment (99.5% frame agreement), so the note now says so and the suggested test tolerance is 0.15. Assisted-by: Claude:claude-opus-5-5 [Claude Code] --- AGENTS.md | 9 +++--- docs/diarization.md | 70 ++++++++++++++++++++++++++++++++++++++------- 2 files changed, 65 insertions(+), 14 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 4e3ab2e..8f89fd1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -108,7 +108,7 @@ tests/ ctest targets test_streaming_eou_reset.cpp, multi-utterance streaming: decoder resets on , 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: diff --git a/docs/diarization.md b/docs/diarization.md index 6ddefaf..85a21af 100644 --- a/docs/diarization.md +++ b/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.