fix(diar): keep frames for words not yet tagged

src/asr/diar/diar_pipeline.cpp

@@ -116,6 +116,8 @@ DiarStream::reset() {
     probs_.clear();
     probs_base_ = 0;
     frozen_segs_.clear();
+    if (hold_from_ != std::numeric_limits<int64_t>::max())
+        hold_from_ = 0;
 }
 
 void
@@ -307,10 +309,9 @@ DiarStream::speaker_for_frames(int64_t start_frame, int64_t end_frame) const {
     if (n == 0 || probs_.empty())
         return -1;
     // riva extrapolates words past the diarized frontier with the last frame;
-    // symmetrically, ranges before the retained window (only reachable when a
-    // caller tags words hours after they were spoken) clamp to its first
-    // frame. Word tagging happens near the frontier, so the clamp is inert in
-    // practice.
+    // symmetrically, ranges before the retained window clamp to its first
+    // frame. Callers that tag late hold their frames (hold_frames_from), so
+    // the clamp only affects ranges nobody held.
     start_frame = std::min(std::max<int64_t>(start_frame, probs_base_), n - 1);
     end_frame = std::max(start_frame + 1, std::min(end_frame, n));
     std::vector<double> mean(n_spk_, 0.0);
@@ -422,7 +423,7 @@ DiarStream::maybe_compact() {
     const float thr = std::min(cfg.onset, cfg.offset);
     const int64_t gap_need = static_cast<int64_t>(
         std::ceil((cfg.min_duration_off + cfg.pad_onset + cfg.pad_offset + 0.5) / sec_per_frame_));
-    const int64_t cut_limit = retained - compact_retain_frames_;
+    const int64_t cut_limit = std::min(retained - compact_retain_frames_, hold_from_ - probs_base_);
     int64_t run = 0, best_cut = -1;
     for (int64_t f = 0; f < cut_limit; f++) {
         bool silent = true;

src/asr/diar/diar_pipeline.h

@@ -13,6 +13,8 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdint>
+#include <limits>
 #include <memory>
 #include <string>
 #include <vector>
@@ -173,6 +175,12 @@ class DiarStream {
     // retained tail. If no such gap exists yet, compaction retries later.
     // Defaults retain roughly the latest 10 minutes after a 20-minute trigger.
     void set_compaction(int64_t trigger_frames, int64_t retain_frames);
+    // Compaction never drops frames at or after `frame`. A caller that tags
+    // words only once their result is final holds the first frame it may
+    // still tag, so those words read raw probabilities instead of the clamped
+    // first retained frame. Unset (the default), only the retain window
+    // limits the cut. reset() rewinds a set hold to frame 0.
+    void hold_frames_from(int64_t frame) { hold_from_ = frame; }
 
     // Speaker id (0-based) for a [start, end) frame range: mean probability
     // over the range, argmax. Frames beyond the emitted timeline use the last
@@ -225,6 +233,7 @@ class DiarStream {
     // Initialized to ~20 min trigger / ~10 min retained at native cadence.
     int64_t compact_trigger_frames_ = 0;
     int64_t compact_retain_frames_ = 0;
+    int64_t hold_from_ = std::numeric_limits<int64_t>::max();
 };
 
 }  // namespace nemo_speech::asr

src/asr/recognizer.cpp

@@ -340,6 +340,9 @@ RecognitionStream::RecognitionStream(
         }
         diar_ = std::make_unique<DiarStream>(
             *recognizer_->diar_model(), recognizer_->config().diar.resolved_geometry());
+        // Words are tagged when their final is built, which for a stream
+        // without endpointing is at finish(). Keep their frames until then.
+        diar_->hold_frames_from(0);
     }
     if (coordinate_ingress_)
         recognizer_->register_streaming_ingress();
@@ -534,6 +537,15 @@ RecognitionStream::next() {
         flush_diar_deficit_(u);
         last_processed_sec_ = std::max(last_processed_sec_, u.audio_processed_sec);
         auto result = build_result_(u, /*is_final=*/true);
+        if (diar_ && !u.words.empty()) {
+            // This final's words are tagged and later words start at or after
+            // its last one. Not audio_processed_sec: the cache runner sets it
+            // to all pushed audio, which can hold utterances not decoded yet.
+            const double last_start_sec =
+                u.words.back().start_frame * recognizer_->ms_per_enc_frame() / 1000.0;
+            diar_->hold_frames_from(
+                static_cast<int64_t>(last_start_sec / diar_->seconds_per_frame()));
+        }
         result.late_punctuation = std::exchange(late_punctuation_, {});
         return result;
     }