Skip to content

parakeet: fix buffered-stream tail loss at end of stream #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cjpais wants to merge 1 commit into
base: main
Choose a base branch
from
Open

parakeet: fix buffered-stream tail loss at end of stream #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/environment-variables.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ its var is unset. Convention: `TRANSCRIBE_<FAMILY>_GGUF`.
| Variable | Test(s) |
| --- | --- |
| `TRANSCRIBE_PARAKEET_GGUF` | `parakeet_real_smoke`, `decoder_smoke` |
| `TRANSCRIBE_PARAKEET_UNIFIED_GGUF` | `parakeet_buffered_stream_eos_smoke` |
| `TRANSCRIBE_COHERE_GGUF` | `cohere_real_smoke`, `cohere_e2e_smoke` |
| `TRANSCRIBE_WHISPER_GGUF` | `whisper_e2e_smoke`, `whisper_tokenize_parity` |
| `TRANSCRIBE_QWEN3_ASR_GGUF` (+ `_0_6B_GGUF` / `_1_7B_GGUF`) | qwen3_asr smokes / parity |
Expand Down
28 changes: 18 additions & 10 deletions docs/models/parakeet-unified-en-0.6b.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,17 @@ variable-stride algorithm — step 0 consumes `chunk + right` audio,
steady-state consumes `chunk`, and the final step folds the trailing
right slot plus the ragged tail into one `is_last` emit.

### End of stream

Finalize flushes retained right context even when the input ends exactly
on a chunk boundary. It also appends `R` frames of silence so final speech
frames keep their trained lookahead and delayed RNN-T tokens are emitted.
Synthetic samples are not included in stream audio accounting, and token
timestamps are clamped to the real input duration. This deliberately
differs from NeMo on the final `R` real frames only. Decoding the masked
conv-overhang frame instead was rejected because it hallucinated trailing
words on naturally ending clips.

### Supported configurations

The runtime `(L, C, R)` is set via `--stream-buf-left-ms /
Expand All @@ -159,20 +170,17 @@ F32 cpp vs the same NeMo buffered-streaming reference loop. The

| `(L, C, R)` | Latency | REF WER% (CI) | cpp F32 WER% (CI) | Δpp |
| ----------- | ------: | --------------------- | --------------------- | ------: |
| `(70, 1, 0)` | 80ms | 5.57% [4.86, 6.31] | 5.76% [5.02, 6.58] | **+0.19** |
| `(70, 1, 0)` | 80ms | 5.57% [4.86, 6.31] | 5.76% [5.03, 6.59] | **+0.19** |
| `(70, 1, 1)` | 160ms | 1.90% [1.61, 2.26] | 1.90% [1.61, 2.26] | +0.00 |
| `(70, 2, 2)` | 320ms | 1.64% [1.33, 1.95] | 1.64% [1.33, 1.95] | -0.00 |
| `(70, 2, 4)` | 480ms | 1.54% [1.26, 1.88] | 1.57% [1.28, 1.91] | +0.03 |
| `(70, 2, 2)` | 320ms | 1.64% [1.33, 1.95] | 1.64% [1.33, 1.95] | +0.00 |
| `(70, 2, 4)` | 480ms | 1.54% [1.26, 1.88] | 1.55% [1.27, 1.89] | +0.01 |
| `(70, 7, 7)` | 1.12s | 1.42% [1.15, 1.74] | 1.40% [1.13, 1.72] | -0.02 |
| `(70, 13, 13)` | 2.08s | 1.44% [1.16, 1.79] | 1.44% [1.16, 1.78] | +0.00 |

Five of six configurations land within 0.03pp of the NeMo reference
(well inside the parakeet family's 0.5% gate); `(70, 1, 1)` and
`(70, 2, 2)` are bit-identical at the Sub/Del/Ins level. The
`(70, 1, 0)` zero-lookahead row is an outlier in both REF (5.57%)
and cpp (5.76%): the model itself doesn't generalize well to this
configuration, and the 4× WER jump vs. `(70, 1, 1)` makes it a poor
choice in practice — use `(70, 1, 1)` if you want the lowest
The cpp column includes the end-of-stream handling described above.
Five of six configurations land within 0.02pp of the NeMo reference,
well inside the parakeet family's 0.5% gate. The zero-lookahead row is
an outlier in both implementations; use `(70, 1, 1)` for the lowest
practical latency.

### Streaming parity reproduction
Expand Down
59 changes: 50 additions & 9 deletions scripts/validate_buffered_streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,8 +156,9 @@ def load_f32(path: Path, expected_shape: list[int] | None = None) -> np.ndarray:
return raw


def compare_pair(name: str, ref_path: Path, cpp_path: Path) -> dict:
"""Per-tensor diff. Returns a row with max_abs, mean_abs, p99_abs, rel."""
def compare_pair(name: str, ref_path: Path, cpp_path: Path,
eos_pad_frames: int | None = None) -> dict:
"""Compare common tensor rows, excluding the expected EOS divergence."""
if not ref_path.exists() or not cpp_path.exists():
return {
"name": name,
Expand All @@ -167,6 +168,24 @@ def compare_pair(name: str, ref_path: Path, cpp_path: Path) -> dict:
}
ref = load_f32(ref_path).astype(np.float64)
cpp = load_f32(cpp_path).astype(np.float64)
extra = {}
if eos_pad_frames is not None and ref.shape != cpp.shape:
if ref.ndim == 1 and cpp.shape[0] >= ref.shape[0]:
# Real audio must match and the synthetic tail must be zero.
pad = cpp[ref.shape[0]:]
extra["pad_samples"] = int(pad.size)
extra["pad_max_abs"] = float(np.abs(pad).max()) if pad.size else 0.0
cpp = cpp[:ref.shape[0]]
elif ref.ndim == 2 and cpp.shape[0] >= ref.shape[0] and cpp.shape[1:] == ref.shape[1:]:
n_real = ref.shape[0] - 1 # drop NeMo's conv-overhang row
extra["pad_rows"] = int(cpp.shape[0] - ref.shape[0])
tail = max(0, min(eos_pad_frames, n_real))
if tail > 0:
extra["eos_tail_rows"] = tail
extra["eos_tail_max_abs"] = float(
np.abs(ref[n_real - tail:n_real] - cpp[n_real - tail:n_real]).max())
ref = ref[:n_real - tail]
cpp = cpp[:n_real - tail]
if ref.shape != cpp.shape:
# Some chunks may have differently-shaped tensors (last-chunk
# divergence). Report and skip.
Expand All @@ -176,6 +195,9 @@ def compare_pair(name: str, ref_path: Path, cpp_path: Path) -> dict:
"ref_shape": list(ref.shape),
"cpp_shape": list(cpp.shape),
}
if ref.size == 0:
return {"name": name, "status": "OK", "max_abs": 0.0, "mean_abs": 0.0,
"p99_abs": 0.0, "rel_max": 0.0, "n_elem": 0, **extra}
diff = np.abs(ref - cpp)
max_abs = float(diff.max())
mean_abs = float(diff.mean())
Expand All @@ -189,6 +211,7 @@ def compare_pair(name: str, ref_path: Path, cpp_path: Path) -> dict:
"p99_abs": p99_abs,
"rel_max": rel,
"n_elem": int(ref.size),
**extra,
}


Expand Down Expand Up @@ -263,13 +286,17 @@ def main() -> int:

rows: list[dict] = []
fail_chunks = 0
last_step = common[-1] if common else None
eos_pad_frames = int(round(args.right_secs * 1000)) // 80
for step in common:
for kind in ("audio_in", "enc_out"):
name = f"stream.chunk.{step}.{kind}"
# The final R rows intentionally differ because cpp adds lookahead.
r = compare_pair(
name,
ref_dir / f"{name}.f32",
cpp_dir / f"{name}.f32",
eos_pad_frames=(eos_pad_frames if step == last_step else None),
)
rows.append({"step": step, **r})
if r.get("status") == "OK":
Expand All @@ -279,8 +306,17 @@ def main() -> int:
if over_max or over_mean:
tag = "FAIL"
fail_chunks += 1
if "pad_max_abs" in r and r["pad_max_abs"] != 0.0:
tag = "FAIL"
fail_chunks += 1
extra = ""
if "pad_samples" in r:
extra = f" eos_pad_samples={r['pad_samples']} (zeros: {r['pad_max_abs'] == 0.0})"
if "eos_tail_max_abs" in r:
extra = (f" eos_tail_rows={r['eos_tail_rows']} eos_tail_max_abs={r['eos_tail_max_abs']:.3e}"
f" pad_rows={r.get('pad_rows', 0)} (expected divergence, not gated)")
print(f" step {step:>2} {kind:9s}: max_abs={r['max_abs']:.3e} "
f"mean_abs={r['mean_abs']:.3e} rel_max={r['rel_max']:.3e} [{tag}]")
f"mean_abs={r['mean_abs']:.3e} rel_max={r['rel_max']:.3e} [{tag}]{extra}")
else:
# Variable-stride algorithm produces identical chunk
# geometry to ref, so SHAPE_DIFF or MISSING is now a
Expand Down Expand Up @@ -310,12 +346,17 @@ def main() -> int:
print(f"FAIL: {fail_chunks} chunks exceed tolerance or have wrong shape")
return 1
if not transcript_match:
# Greedy RNN-T can tip a single-token decision on fp32 noise
# even when per-chunk encoder outputs match within tolerance.
# Report informationally — per-chunk parity is the algorithmic
# gate; WER on test-clean is the corpus-level gate.
print("WARN: transcript byte-match differs (per-chunk parity gate passes — "
"likely fp32-noise tipping a greedy decision)")
if cpp_text_norm.startswith(ref_text_norm):
tail = cpp_text_norm[len(ref_text_norm):]
print(f"INFO: cpp transcript extends ref by {tail!r} "
"(EOS silence lookahead; expected)")
else:
# Greedy RNN-T can tip a single-token decision on fp32 noise
# even when per-chunk encoder outputs match within tolerance.
# Report informationally — per-chunk parity is the algorithmic
# gate; WER on test-clean is the corpus-level gate.
print("WARN: transcript byte-match differs (per-chunk parity gate passes — "
"likely fp32-noise tipping a greedy decision)")
print("OK")
return 0

Expand Down
50 changes: 33 additions & 17 deletions src/arch/parakeet/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2292,6 +2292,13 @@ void rebuild_streaming_result_text(ParakeetSession * pc, const ParakeetModel * p
}
pc->raw_text = tok.decode(raw_ids.data(), static_cast<int>(raw_ids.size()));
}
if (pc->buf_active) {
const int64_t max_ms = us_to_ms(pc->stream_audio_input_us);
for (auto & token : pc->tokens) {
token.t1_ms = std::min(token.t1_ms, max_ms);
token.t0_ms = std::min(token.t0_ms, token.t1_ms);
}
}
pc->has_result = true;
pc->result_kind = TRANSCRIBE_TIMESTAMPS_TOKEN;
}
Expand All @@ -2300,12 +2307,8 @@ void rebuild_streaming_result_text(ParakeetSession * pc, const ParakeetModel * p
//
// Mirrors NeMo's speech_to_text_streaming_infer_rnnt.py. Variable-stride
// per step: step 0 num_new = samples_chunk + samples_right; steady state
// num_new = samples_chunk; the final step (finalize) consumes the rest
// and folds the right slot into chunk. Each step updates the buffer's
// ContextSize (buf_ctx_*), slices the encoder window, computes mel,
// builds the graph with a BufferedStreamMaskOverride, then slices off the
// ctx_left frames and decodes ctx_chunk (or all remaining on last) with a
// carried RNN-T LstmState.
// num_new = samples_chunk. Finalize folds the retained right slot and any
// ragged tail into the decoded chunk.
static void buf_ctx_add_frames(ParakeetSession * pc, int64_t num_new, bool is_last) {
pc->buf_ctx_left += pc->buf_ctx_chunk;
pc->buf_ctx_chunk = 0;
Expand All @@ -2324,10 +2327,13 @@ static void buf_ctx_add_frames(ParakeetSession * pc, int64_t num_new, bool is_la
pc->buf_ctx_left -= extra;
}

// The final chunk may append silence for right-context lookahead. Padding
// extends the decoded window but never advances the real-audio cursor.
transcribe_status emit_buffered_chunk(ParakeetSession * pc,
ParakeetModel * pm,
int64_t num_new_samples,
bool is_last_chunk) {
bool is_last_chunk,
int64_t eos_pad_samples = 0) {
if (pc->poll_abort()) {
return TRANSCRIBE_ERR_ABORTED;
}
Expand All @@ -2345,9 +2351,10 @@ transcribe_status emit_buffered_chunk(ParakeetSession * pc,

// ----- Update buffer ContextSize (mirrors NeMo's add_frames_get_removed_) -----
buf_ctx_add_frames(pc, num_new_samples, is_last_chunk);
pc->buf_ctx_chunk += eos_pad_samples;

// ----- Build the [left | chunk | right] PCM window from absolute coords -----
const int64_t end_abs = pc->buf_next_audio_read + num_new_samples;
const int64_t end_abs = pc->buf_next_audio_read + num_new_samples + eos_pad_samples;
const int64_t total_now = pc->buf_ctx_left + pc->buf_ctx_chunk + pc->buf_ctx_right;
const int effective_T = static_cast<int>(total_now / samples_per_frame);
const int64_t start_abs = end_abs - total_now;
Expand Down Expand Up @@ -2831,6 +2838,8 @@ transcribe_status stream_begin(transcribe_session * session,
return TRANSCRIBE_ERR_NOT_IMPLEMENTED;
}

pc->stream_audio_input_samples = 0;

// -------- Buffered streaming path (parakeet-unified-en-0.6b) --------
//
// chunked_limited_with_rc with a 3-tuple training menu. Re-runs the
Expand Down Expand Up @@ -2961,7 +2970,8 @@ transcribe_status stream_feed(transcribe_session * session,
}

pc->stream_pcm_buffer.insert(pc->stream_pcm_buffer.end(), pcm, pcm + n_samples);
pc->stream_audio_input_us += samples_to_us(n_samples);
pc->stream_audio_input_samples += n_samples;
pc->stream_audio_input_us = samples_to_us(pc->stream_audio_input_samples);

const int prev_n_tokens = static_cast<int>(pc->raw_tokens.size());

Expand All @@ -2979,6 +2989,7 @@ transcribe_status stream_feed(transcribe_session * session,
// job). Step 0 needs samples_chunk + samples_right; steady-state
// needs samples_chunk.
if (pc->buf_active) {
bool emitted = false;
while (true) {
if (pc->poll_abort()) {
return TRANSCRIBE_ERR_ABORTED;
Expand All @@ -2994,6 +3005,7 @@ transcribe_status stream_feed(transcribe_session * session,
st != TRANSCRIBE_OK) {
return st;
}
emitted = true;
}
const bool tokens_changed = static_cast<int>(pc->raw_tokens.size()) != prev_n_tokens;
if (tokens_changed) {
Expand All @@ -3002,8 +3014,10 @@ transcribe_status stream_feed(transcribe_session * session,
pc->n_committed_words = 0;
pc->n_committed_segments = 0;
pc->stream_revision += 1;
pc->stream_audio_committed_us =
pc->buf_next_audio_read * 1000000LL / std::max<int64_t>(pm->hparams.fe_sample_rate, 1);
}
if (emitted) {
// The retained right slot has been read but not decoded.
pc->stream_audio_committed_us = samples_to_us(pc->buf_next_audio_read - pc->buf_ctx_right);
}
if (update != nullptr) {
update->result_changed = tokens_changed;
Expand Down Expand Up @@ -3183,18 +3197,19 @@ transcribe_status stream_finalize(transcribe_session * session, transcribe_strea
const int prev_n_tokens = static_cast<int>(pc->raw_tokens.size());

// -------- Buffered streaming finalize --------
//
// One final emit consuming all remaining audio with is_last_chunk=true;
// add_frames_get_removed_ folds the right slot + this num_new into the
// chunk slot so the decoder gets every frame past ctx_left (no zero-pad).
// Flush retained right context even when the read cursor is already at
// EOS, then append R frames of silence so the final speech frames keep
// their trained lookahead. Synthetic samples are not counted as input.
if (pc->buf_active) {
const int64_t total = static_cast<int64_t>(pc->stream_pcm_buffer.size());
if (pc->buf_next_audio_read < total) {
// An exact C+R+k*C input leaves right context retained at EOS.
if (pc->buf_next_audio_read < total || pc->buf_ctx_right > 0) {
if (pc->poll_abort()) {
return TRANSCRIBE_ERR_ABORTED;
}
const int64_t num_new = total - pc->buf_next_audio_read;
if (const transcribe_status st = emit_buffered_chunk(pc, pm, num_new, /*is_last_chunk=*/true);
const int64_t eos_pad = static_cast<int64_t>(pc->buf_samples_right);
if (const transcribe_status st = emit_buffered_chunk(pc, pm, num_new, /*is_last_chunk=*/true, eos_pad);
st != TRANSCRIBE_OK) {
return st;
}
Expand Down Expand Up @@ -3305,6 +3320,7 @@ transcribe_status stream_finalize(transcribe_session * session, transcribe_strea
void stream_reset(transcribe_session * session) {
auto * pc = static_cast<ParakeetSession *>(session);
pc->stream_pcm_buffer.clear(); // keep the allocation
pc->stream_audio_input_samples = 0;
}

// Kind+slot probe. No run-slot extensions (always false on _RUN). On
Expand Down
13 changes: 6 additions & 7 deletions src/arch/parakeet/parakeet.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,8 @@ struct ParakeetSession final : public transcribe_session {
// clear_result (the family owns its per-utterance audio scratch).
std::vector<float> stream_pcm_buffer;
transcribe_run_params stream_run_params{};
// Convert from cumulative samples so odd feed sizes do not lose fractions.
int64_t stream_audio_input_samples = 0;

// ---- incremental streaming state (cache-aware) ----
//
Expand All @@ -285,13 +287,10 @@ struct ParakeetSession final : public transcribe_session {
//
// Mirrors NeMo's StreamingBatchedAudioBuffer + reference inference
// loop. Variable-stride: step 0 consumes samples_chunk + samples_right
// of new audio; subsequent feeds consume samples_chunk; finalize's
// last step consumes the rest (chunk slot absorbs the trailing right
// context, no zero-pad). The buf_* geometry reflects the active
// (L, C, R) tuple; the ctx_* fields track the buffer's internal
// ContextSize, updated per-chunk via buf_ctx_add_frames (NeMo's
// add_frames_get_removed_). The RNN-T state rides on stream_dec_state
// (same predictor/joint path, no per-layer encoder cache).
// of new audio; subsequent feeds consume samples_chunk. Finalize folds
// retained right context into the decoded chunk and appends right-context
// silence without advancing the real-audio cursor. The RNN-T state rides
// on stream_dec_state (same predictor/joint path, no per-layer encoder cache).
int32_t buf_left_frames = 0; // L (expected)
int32_t buf_chunk_frames = 0; // C
int32_t buf_right_frames = 0; // R
Expand Down
17 changes: 17 additions & 0 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -703,6 +703,23 @@ if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS)
COMMAND transcribe_parakeet_multitalker_e2e_smoke)
set_tests_properties(transcribe_parakeet_multitalker_e2e_smoke PROPERTIES
SKIP_RETURN_CODE 77)

# Buffered-stream tail regression; skips when the unified model is absent.
add_executable(transcribe_parakeet_buffered_stream_eos_smoke
parakeet_buffered_stream_eos_smoke.cpp)

target_link_libraries(transcribe_parakeet_buffered_stream_eos_smoke
PRIVATE transcribe transcribe-common-example)

target_compile_definitions(transcribe_parakeet_buffered_stream_eos_smoke PRIVATE
"TRANSCRIBE_TEST_SAMPLES_DIR=\"${CMAKE_SOURCE_DIR}/samples\"")

transcribe_apply_warnings(transcribe_parakeet_buffered_stream_eos_smoke)

add_test(NAME transcribe_parakeet_buffered_stream_eos_smoke
COMMAND transcribe_parakeet_buffered_stream_eos_smoke)
set_tests_properties(transcribe_parakeet_buffered_stream_eos_smoke PROPERTIES
SKIP_RETURN_CODE 77)
endif()

# -----------------------------------------------------------------------------
Expand Down
Loading
Loading