From c35746f1b06aa0a0d0dfa740347a3ada9e4933f4 Mon Sep 17 00:00:00 2001 From: Rafael Wille Date: Fri, 31 Jul 2026 17:50:28 -0300 Subject: [PATCH] cohere: add long-form chunking Cohere Transcribe trains on clips of at most 35 seconds, but the runtime gated only against the encoder positional table (~400 s). Past the trained span the decoder's cross-attention loses alignment, emits EOS early, and drops the rest of the audio, returning TRANSCRIBE_OK with no WARN and no transcribe_was_truncated(). Window longer audio to the trained span with 1 s of overlap, preferring low-energy boundaries, and stitch the windows in token space by matching the tail of the accumulated hypothesis against the head of the next one. Decode once at the end so the tokenizer's own spacing is preserved. The seam search is ported from canary_token_seam(), including its asymmetric search widths, so both families stitch the same way. Report input capacity as unbounded (max_audio_ms = 0) and move the family to the chunked bucket in docs/input-limits.md. n_ctx still bounds the per-window output budget. Batches carrying long inputs fall back to the serial path, since the batched graph assumes one encoder pass per utterance. Two assertions in cohere_e2e_smoke.cpp expected a rejected transcribe_run to clear the previous result; src/transcribe.cpp documents preserving it as intentional, so they now check the result survives. --- docs/input-limits.md | 44 ++- docs/models/cohere-transcribe-03-2026.md | 22 +- examples/cli/main.cpp | 12 +- include/transcribe.h | 2 +- src/arch/cohere/cohere.h | 28 ++ src/arch/cohere/model.cpp | 475 +++++++++++++++++++++-- tests/cohere_e2e_smoke.cpp | 104 ++++- tests/cohere_smoke.cpp | 132 +++++++ 8 files changed, 759 insertions(+), 60 deletions(-) diff --git a/docs/input-limits.md b/docs/input-limits.md index fb7a6ae0..65385b31 100644 --- a/docs/input-limits.md +++ b/docs/input-limits.md @@ -57,6 +57,7 @@ value at the default context; for the effective limit under a lowered | Families | Limit | Behavior | | --- | --- | --- | | whisper, parakeet (all variants), voxtral_realtime | none (`max_audio_ms = 0`) | Long audio is windowed internally and stitched (whisper), processed by an unbounded/streaming encoder (parakeet, voxtral_realtime), or padded if short. No practical limit; all three **ignore `n_ctx`** (it cannot lower a limit they don't have). whisper and parakeet never error on length. voxtral_realtime has one exception — its absolute `dec_max_position` cap (~2.9 h, see below): a clip past it returns `INPUT_TOO_LONG` (one-shot/batch), and a stream that reaches it flags `was_truncated`. | +| cohere | none (`max_audio_ms = 0`) | Windowed to the trained clip span (35 s) before the encoder, then the per-window transcripts are concatenated. Unlike the three above, cohere **does honor `n_ctx`**: it still bounds the decoder self-KV and therefore the per-window output budget, but no longer bounds input length. | Whisper slices audio into 30 s windows with prev-context stitching; parakeet's conformer is effectively unbounded (the encoder positional table is recomputed @@ -67,11 +68,27 @@ constant-memory caches (`cache_last_channel` / `cache_last_time` + the decoder LSTM state) rather than a growing KV, so it is unbounded. These families do not need and do not have a length gate. +Cohere is windowed for a different reason. Its encoder positional table spans +~400 s, so a much longer clip than the model was trained on still encodes +cleanly and an input gate keyed on that table lets it through. The decoder is +what breaks: its cross-attention has no monotonicity constraint (unlike an +RNN-T, which walks encoder frames in order), so past the trained clip span +(`max_audio_clip_s`, 35 s) it loses alignment, emits EOS early, and drops the +rest of the audio — returning `TRANSCRIBE_OK` with a silently incomplete +transcript, which is exactly what this document promises never happens. So the +family windows to the trained span and stitches the windows back together. +Windows overlap by one second and are joined in token space by matching the tail +of the accumulated hypothesis against the head of the next window, so a boundary +landing mid-word costs nothing — the word survives whole in the next window and +the duplicate is trimmed at the join. Boundaries still prefer a detected pause, +which makes that match easier to find. When no join is identified the window is +appended whole, repeating a few words rather than dropping any. + ### 2. Hard context cap — reject up front | Families | Limit source | Behavior | | --- | --- | --- | -| qwen3_asr, canary_qwen, funasr_nano, granite, granite_nar, voxtral, cohere, canary | decoder context window (`dec_max_position_embeddings` / `dec_max_seq`), or the encoder positional table (`enc_pos_emb_max_len`, for cohere/canary) — all from GGUF | KV cache grows to fit, clamped to the model's true max. Over-length input is **rejected before the decode** (or before the encoder, where the encoder table is the binding limit) with `TRANSCRIBE_ERR_INPUT_TOO_LONG`. | +| qwen3_asr, canary_qwen, funasr_nano, granite, granite_nar, voxtral, canary | decoder context window (`dec_max_position_embeddings` / `dec_max_seq`), or the encoder positional table (`enc_pos_emb_max_len`, for canary) — all from GGUF | KV cache grows to fit, clamped to the model's true max. Over-length input is **rejected before the decode** (or before the encoder, where the encoder table is the binding limit) with `TRANSCRIBE_ERR_INPUT_TOO_LONG`. | These families wrap an LLM-style decoder whose context window (`audio_tokens + prompt + generation`) is the binding constraint. The number of @@ -138,16 +155,21 @@ reports the model's default-context ceiling (`n_ctx == 0`); it is not re-derived for a session that narrows `n_ctx`. A session that lowers `n_ctx` may therefore reject audio shorter than the advertised `max_audio_ms`. -Encoder-bound families are different. For cohere and canary, the input-audio -limit is the encoder positional table, while `n_ctx` only bounds the decoder -self-KV / output budget. In those families `transcribe_session_get_limits()` -reports a smaller `effective_n_ctx` and `max_kv_bytes` when `n_ctx` is lowered, -but `effective_max_audio_ms` stays pinned to the encoder input bound. - -**Chunked / unbounded families (bucket 1) ignore `n_ctx` entirely.** whisper, -parakeet, and voxtral_realtime have no lowerable context ceiling, so a non-zero -`n_ctx` is a documented no-op and `transcribe_session_get_limits()` keeps -reporting them unbounded. voxtral_realtime is the subtle case: it *does* have an +Encoder-bound families are different. For canary, the input-audio limit is the +encoder positional table, while `n_ctx` only bounds the decoder self-KV / output +budget. There `transcribe_session_get_limits()` reports a smaller +`effective_n_ctx` and `max_kv_bytes` when `n_ctx` is lowered, but +`effective_max_audio_ms` stays pinned to the encoder input bound. + +Cohere is the hybrid: it windows its input (bucket 1, `max_audio_ms = 0`) but +its decoder self-KV is still `n_ctx`-bounded, so a lowered `n_ctx` shrinks the +per-window output budget without changing how much audio it accepts. + +**Most chunked / unbounded families (bucket 1) ignore `n_ctx` entirely** — +cohere, described just above, is the exception. whisper, parakeet, and +voxtral_realtime have no lowerable context ceiling, so a non-zero `n_ctx` is a +documented no-op and `transcribe_session_get_limits()` keeps reporting them +unbounded. voxtral_realtime is the subtle case: it *does* have an absolute decoder position cap (`dec_max_position`, ~2.9 h of audio), but that is the model's true RoPE wall — not a memory ceiling a caller can lower — so `n_ctx` does not narrow it (its decoder KV is a constant-memory sliding ring; diff --git a/docs/models/cohere-transcribe-03-2026.md b/docs/models/cohere-transcribe-03-2026.md index f605605b..9a945545 100644 --- a/docs/models/cohere-transcribe-03-2026.md +++ b/docs/models/cohere-transcribe-03-2026.md @@ -21,10 +21,24 @@ pinned 2026-04-16. ## Input limits -Accepts up to about **6.7 minutes (400 s)** of 16 kHz mono audio per call — the -encoder's positional table is the binding limit. Longer audio is rejected up -front with `TRANSCRIBE_ERR_INPUT_TOO_LONG` rather than silently truncated; split -it into shorter segments. See the [input-length contract](../input-limits.md). +**No practical limit** (`max_audio_ms = 0`): audio longer than the model's +trained clip span (`max_audio_clip_s`, 35 s) is windowed internally and the +per-window transcripts are concatenated, so a caller never has to segment input +by hand. + +Windowing is required, not an optimization. The encoder's positional table +spans ~400 s, so a much longer clip than the model ever saw in training still +encodes cleanly; the decoder is what fails. Its cross-attention has no +monotonicity constraint, so past the trained span it loses alignment, emits EOS +early, and drops the rest of the audio — returning `TRANSCRIBE_OK` with a +silently incomplete transcript. Windows overlap by one second and are joined in +token space, so a boundary that lands mid-word is harmless — the word survives +whole in the next window and the duplicate is trimmed at the join. Boundaries +still prefer a detected pause, which makes that match easier to find. + +`n_ctx` still applies: it bounds the decoder self-KV, and therefore the output +budget *per window*, but no longer bounds input length. See the +[input-length contract](../input-limits.md). ## Download diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index a4d0ed13..8b4373c3 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1166,6 +1166,16 @@ int main(int argc, char ** argv) { { struct transcribe_session_limits lim; transcribe_session_limits_init(&lim); + + // A model that advertises no input ceiling of its own is unbounded + // however small its decoder context is -- the family windows the + // audio internally. Without this, a family that is both windowed + // and context-capped (cohere) reads as "context too small". + struct transcribe_capabilities lim_caps; + transcribe_capabilities_init(&lim_caps); + const bool model_is_unbounded = + transcribe_model_get_capabilities(model, &lim_caps) == TRANSCRIBE_OK && lim_caps.max_audio_ms == 0; + if (transcribe_session_get_limits(ctx, &lim) == TRANSCRIBE_OK) { if (lim.effective_max_audio_ms > 0) { std::printf(" max audio: %.1f s", (double) lim.effective_max_audio_ms / 1000.0); @@ -1174,7 +1184,7 @@ int main(int argc, char ** argv) { (long long) (lim.max_kv_bytes >> 20)); } std::printf("\n"); - } else if (lim.effective_n_ctx > 0) { + } else if (lim.effective_n_ctx > 0 && !model_is_unbounded) { // Capped family whose context is too small to fit any audio // plus a prompt (e.g. an aggressively low --n-ctx). std::printf( diff --git a/include/transcribe.h b/include/transcribe.h index e789bbd7..db4350c8 100644 --- a/include/transcribe.h +++ b/include/transcribe.h @@ -243,7 +243,7 @@ typedef enum { * Returned by transcribe_run / transcribe_run_batch when the input * audio is longer than the loaded model can process in a single * decode. Hard-context-cap families (LLM-style decoders: qwen3_asr, - * canary_qwen, funasr_nano, granite, granite_nar, voxtral, cohere, + * canary_qwen, funasr_nano, granite, granite_nar, voxtral, * canary) reject an over-length clip UP FRONT — before the decode * (and, where the binding limit is the encoder's positional table, * before the encoder) — by comparing the audio's prefill token count diff --git a/src/arch/cohere/cohere.h b/src/arch/cohere/cohere.h index 319e6151..962730a9 100644 --- a/src/arch/cohere/cohere.h +++ b/src/arch/cohere/cohere.h @@ -33,6 +33,30 @@ namespace transcribe::cohere { void apply_family_invariants(transcribe_model & model); +// Long-form window boundary picker: searches back from `target` over `search` +// samples for a genuine pause to cut on, and returns -1 when the region is +// continuous speech and no candidate qualifies. Declared here so the boundary +// logic can be unit-tested without a model; see model.cpp for the full +// contract and its limitations. +int pick_split(const float * pcm, int n_samples, int target, int search); + +// Where two overlapping windows should be joined, in token space. +struct TokenSeam { + int previous_keep; // tokens to keep from the accumulated hypothesis + int current_skip; // tokens to drop from the front of the new window + bool matched; // false when no overlap was identified +}; + +// Find the join point between consecutive windows by matching the tail of +// `previous` against the head of `current`. Timestamp-free, so it works for a +// family that emits no alignment data. Ported from the canary long-form path +// (see PR #112) so both families stitch the same way. Declared here for unit +// testing; see model.cpp for the acceptance rules. +TokenSeam token_seam(const std::vector & previous, + const std::vector & current, + int previous_search, + int current_search); + // KV cache for the autoregressive decoder. Flat 1D K/V tensors (whisper.cpp // pattern), per-layer slices via views. Self cache grows per step; // cross cache is computed once from encoder output, then reused. @@ -172,6 +196,10 @@ struct CohereSession final : public transcribe_session { bool encoder_use_flash = true; bool decoder_use_flash = true; + // Token ids the last decoded window produced, kept so the long-form path + // can stitch windows in token space instead of re-parsing decoded text. + std::vector window_ids; + CohereSession() = default; ~CohereSession() override; }; diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index 368b3f86..3bfd20fc 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -220,19 +221,187 @@ CohereModel::~CohereModel() { plan.primary_kind = transcribe::BackendKind::Unknown; } +// Long-form windowing (see run() below for the rationale). +// +// Cohere Transcribe was trained on clips of at most 35 s -- `max_audio_clip_s` +// in the HF feature-extractor config, recorded in the porting intake. That +// value is not written to the GGUF today, so the trained window lives here as +// a family constant; emitting it as `stt.cohere.max_audio_clip_s` and reading +// it in read_cohere_hparams() (with this as the fallback for already-published +// GGUFs) would be the tidier follow-up. +constexpr double kTrainedClipSeconds = 35.0; + +// Fractions of a window searched backwards for a pause to cut on. The wider +// pass is a fallback when the narrow one finds only continuous speech. +constexpr double kSplitSearchFraction = 0.25; +constexpr double kSplitSearchFractionWide = 0.50; + +// A frame counts as a pause only if it is this much quieter than the mean of +// the searched region. A minimum always exists; a pause does not. +constexpr float kPauseRatio = 0.30f; + +// Search back from `target` for a genuine pause to cut on: the quietest 20 ms +// frame within the trailing `search` samples, accepted only when it is clearly +// quieter than that region's own mean amplitude. +// +// Returns the cut offset, or -1 when the region is continuous speech and no +// pause qualifies -- the caller decides what to do rather than getting a +// silently arbitrary cut. +// +// A pause is preferred but not required: windows overlap by kOverlapMs, so a +// cut that lands mid-word still leaves that word whole in the next window, and +// token_seam() removes the duplicate when the two are joined. Cutting on a +// pause simply makes the seam easier to find. +int pick_split(const float * pcm, int n_samples, int target, int search) { + constexpr int frame = 320; // 20 ms @ 16 kHz + if (target >= n_samples || search < frame * 2 || target < frame * 2) { + return -1; + } + + const int lo = std::max(frame, target - search); + const int step = frame / 2; + + double sum_amp = 0.0; + int n_frames = 0; + int best_at = -1; + float best_amp = std::numeric_limits::max(); + + for (int at = lo; at + frame <= target; at += step) { + float sum = 0.0f; + for (int i = 0; i < frame; ++i) { + sum += std::fabs(pcm[at + i]); + } + const float amp = sum / static_cast(frame); + + sum_amp += amp; + ++n_frames; + + if (amp < best_amp) { + best_amp = amp; + best_at = at + frame / 2; // cut in the middle of the pause + } + } + + if (n_frames == 0 || best_at < 0) { + return -1; + } + + // Reject a "quietest" frame that is really just ordinary speech. + const float mean_amp = static_cast(sum_amp / n_frames); + if (best_amp > mean_amp * kPauseRatio) { + return -1; + } + return best_at; +} + +// Overlap carried between consecutive windows. A boundary that lands in +// continuous speech splits a word; repeating a second of audio on both sides +// means the word survives whole in at least one window, and token_seam() drops +// the duplicate when the windows are joined. +constexpr int kOverlapMs = 1000; + +// Seam search widths, in encoder frames' worth of tokens, derived from the +// overlap rather than hardcoded. +// +// The two sides are deliberately ASYMMETRIC, and the small one is the load- +// bearing part. Only the first ~1 s of a new window can legitimately duplicate +// the previous one, so `current_search` must not reach past the overlap: a +// wider window lets a genuinely repeated sentence match the accumulated tail +// and be deleted as if it were overlap. `previous_search` can be looser because +// the at-previous-end rule already constrains where a match may finish. +// +// Mirrors the canary long-form path (PR #112), which computes the same two +// widths from its own overlap. +int seam_search_previous(const CohereHParams & hp) { + const int frame_ms = + hp.fe_sample_rate > 0 ? hp.fe_hop_length * hp.enc_subsampling_factor * 1000 / hp.fe_sample_rate : 0; + const int delay_frames = std::max(1, frame_ms > 0 ? kOverlapMs / frame_ms : 12); + return delay_frames * 2; +} + +int seam_search_current(const CohereHParams & hp) { + const int frame_ms = + hp.fe_sample_rate > 0 ? hp.fe_hop_length * hp.enc_subsampling_factor * 1000 / hp.fe_sample_rate : 0; + const int delay_frames = std::max(1, frame_ms > 0 ? kOverlapMs / frame_ms : 12); + return std::max(1, delay_frames * 3 / 5); +} + +// Find where to join two consecutive windows by matching the tail of the +// accumulated hypothesis against the head of the new window. +// +// The match is only accepted when it runs to the *end* of the previous +// hypothesis. A match that stops earlier means the repeated text sits in the +// interior, which is ambiguous -- the speaker may genuinely have repeated the +// phrase, and trimming it would delete real speech. A single-token match is +// accepted only at the very start of the new window, where it cannot be +// coincidence. +// +// Ported from canary_token_seam() in the canary long-form path (PR #112) so the +// two families stitch identically. Timestamp-free by construction, which is +// what makes it usable here: cohere advertises max_timestamp_kind == NONE and +// has no alignment data to fall back on. +TokenSeam token_seam(const std::vector & previous, + const std::vector & current, + int previous_search, + int current_search) { + TokenSeam seam{ static_cast(previous.size()), 0, false }; + if (previous.empty() || current.empty() || previous_search <= 0 || current_search <= 0) { + return seam; + } + + const int previous_begin = std::max(0, static_cast(previous.size()) - previous_search); + const int current_end = std::min(static_cast(current.size()), current_search); + + int best_length = 0; + int best_curr_end = 0; + + for (int i = previous_begin; i < static_cast(previous.size()); ++i) { + for (int j = 0; j < current_end; ++j) { + int length = 0; + while (i + length < static_cast(previous.size()) && j + length < current_end && + previous[i + length] == current[j + length]) { + ++length; + } + const bool at_previous_end = (i + length) == static_cast(previous.size()); + if (!at_previous_end || !(length >= 2 || (length == 1 && j == 0))) { + continue; + } + if (length > best_length) { + best_length = length; + best_curr_end = j + length; + } + } + } + + // A seam that would swallow the entire new window is never overlap -- the + // window covers ~34 s of fresh audio and the overlap is 1 s. Treat it as a + // failed match and append whole; repeating a phrase is recoverable, dropping + // a window is the exact defect this path exists to prevent. + if (best_length > 0 && best_curr_end < static_cast(current.size())) { + seam.current_skip = best_curr_end; + seam.matched = true; + } + return seam; +} + namespace { constexpr float kBnEps = 1e-5f; // Input-length contract (canonical reference; see docs/input-limits.md). -// Cohere ASR has TWO distinct limits: -// (a) INPUT — the encoder's relative positional-encoding table -// (enc_pos_emb_max_len). T_enc must stay within the trained pos-emb -// span or the runtime table aliases past the trained range. Binding -// input bound, gated up front; drives caps.max_audio_ms. -// (b) DECODER self-KV — dec_max_seq and a 512 max-new-tokens cap. Bound -// the *output*: an over-budget transcript is kept as a partial result -// and flagged via transcribe_was_truncated(), not rejected. +// Cohere ASR has THREE bounds, only one of which a caller ever sees: +// (a) TRAINED CLIP SPAN — kTrainedClipSeconds, the window the model was +// trained on. run() windows longer audio to it, so this is what actually +// shapes behaviour. Because run() handles it, caps.max_audio_ms is 0 +// (no practical limit) and the family sits in bucket 1. +// (b) ENCODER pos-emb table — enc_pos_emb_max_len. T_enc must stay within +// the trained span or the runtime table aliases. No longer the public +// limit: each window is ~35 s, far under it, so the up-front gate in +// run_window() is now a per-window assertion rather than a caller-facing +// bound. +// (c) DECODER self-KV — dec_max_seq and a 512 max-new-tokens cap. Bounds the +// *output* per window: an over-budget transcript is kept as a partial +// result and flagged via transcribe_was_truncated(), not rejected. // Predicted encoder frame count T_enc for a given mel frame count. The // FastConformer pre-encode downsamples time via three stride-2, kernel-3, @@ -461,21 +630,28 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } - // Publish the input-length ceiling now that the encoder pos-emb span - // and frontend rate are known (the encoder is the binding INPUT limit). - m->caps.max_audio_ms = cohere_max_audio_ms(m->hparams); + // No practical input-length limit: run() windows long audio to the trained + // clip span (kTrainedClipSeconds) before it ever reaches the encoder, so the + // encoder pos-emb table is never the binding bound for a caller. This puts + // cohere in the "chunked / unbounded" bucket of docs/input-limits.md, + // alongside whisper and parakeet. cohere_max_audio_ms() is still the + // per-window ceiling and is asserted inside run_window(). + m->caps.max_audio_ms = 0; // Basis for the session-level limits query (transcribe_session_get_limits). // The LimitsBasis is a single context cap, filled from the DECODER side - // (dec_max_seq) so effective_n_ctx and max_kv_bytes are exact. Its derived - // effective_max_audio_ms therefore keys off the decoder ceiling and will - // NOT match caps.max_audio_ms (the encoder input bound) — for cohere the - // decoder context bounds the OUTPUT transcript, so that figure is advisory. + // (dec_max_seq) so effective_n_ctx and max_kv_bytes are exact. It bounds the + // OUTPUT transcript per window, which is why lowering n_ctx shrinks the + // per-window generation budget without changing how much audio the family + // accepts. if (m->hparams.dec_max_seq > 0 && m->hparams.enc_subsampling_factor > 0 && m->hparams.fe_hop_length > 0 && m->hparams.fe_sample_rate > 0) { m->limits.has_context_cap = true; - // Audio bound is the encoder table (caps.max_audio_ms), not the - // decoder context, so effective_max_audio_ms must not shrink with n_ctx. + // Take the audio bound straight from caps (0 = unbounded) rather than + // deriving it from the decoder ceiling: run() windows the input, so the + // amount of audio accepted does not shrink with n_ctx. Reporting + // effective_max_audio_ms == 0 keeps the session query consistent with + // caps.max_audio_ms. m->limits.audio_from_caps = true; m->limits.model_max_ctx = m->hparams.dec_max_seq; // Fixed control-token preamble (see run()'s prompt_pieces). Audio is @@ -638,10 +814,13 @@ transcribe_status init_context(transcribe_model * model, return TRANSCRIBE_OK; } -transcribe_status run(transcribe_session * session, - const float * pcm, - int n_samples, - const transcribe_run_params * params) { +// Transcribe one window of PCM. This is the single-pass path: mel -> encoder +// -> autoregressive decode -> committed result. run() below drives it once for +// short-form audio and repeatedly for long-form. +transcribe_status run_window(transcribe_session * session, + const float * pcm, + int n_samples, + const transcribe_run_params * params) { if (session == nullptr || pcm == nullptr || n_samples <= 0) { return TRANSCRIBE_ERR_INVALID_ARG; } @@ -652,7 +831,7 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_INVALID_ARG; } - // Pre-run abort check (cohere is single-chunk today). + // Pre-run abort check. if (cc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } @@ -684,7 +863,7 @@ transcribe_status run(transcribe_session * session, if (cm->hparams.enc_pos_emb_max_len > 0) { const int t_enc_pred = cohere_predict_t_enc(mel_n_frames, cm->hparams.enc_subsampling_factor); if (t_enc_pred > cm->hparams.enc_pos_emb_max_len) { - const double max_s = static_cast(cm->caps.max_audio_ms) / 1000.0; + const double max_s = static_cast(cohere_max_audio_ms(cm->hparams)) / 1000.0; transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere run: input too long — %d encoder frames exceed the %d " "the model supports (~%.0f s max). See " @@ -1082,6 +1261,11 @@ transcribe_status run(transcribe_session * session, auto commit_result = [&]() { cc->t_decode_us = ggml_time_us() - t_dec_start; + // Publish the raw ids so the long-form path in run() can stitch + // windows in token space. Set before the empty check so a window + // that produced nothing clears the previous window's ids. + cc->window_ids = generated_ids; + if (generated_ids.empty()) { return; } @@ -1315,6 +1499,189 @@ transcribe_status run(transcribe_session * session, return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } +// Long-form entry point. +// +// The encoder's positional table spans ~400 s, so a clip well past the trained +// window still encodes cleanly and the up-front gate lets it through. The +// decoder is the problem: its cross-attention has no monotonicity constraint +// (unlike an RNN-T, which walks the encoder frames in order), so past the +// trained span it drifts, emits EOS early, and drops the tail. The run then +// returns TRANSCRIBE_OK with a silently incomplete transcript -- exactly what +// docs/input-limits.md promises never happens. +// +// So window the audio to the trained span. Windows overlap by kOverlapMs and +// are stitched in token space by token_seam(), so a boundary that lands +// mid-word costs nothing: the word survives whole in the next window and the +// duplicate is trimmed at the join. Boundaries still prefer a pause where one +// is detectable, which makes the seam easier to find. Short-form audio takes +// the single-pass path unchanged. +transcribe_status run(transcribe_session * session, + const float * pcm, + int n_samples, + const transcribe_run_params * params) { + if (session == nullptr || pcm == nullptr || n_samples <= 0) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + auto * cc = static_cast(session); + auto * cm = static_cast(cc->model); + if (cm == nullptr) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + const int sample_rate = cm->hparams.fe_sample_rate > 0 ? cm->hparams.fe_sample_rate : 16000; + const int window_samples = static_cast(kTrainedClipSeconds * sample_rate); + + // Short-form: one window, identical to the pre-windowing path. + if (window_samples <= 0 || n_samples <= window_samples) { + return run_window(session, pcm, n_samples, params); + } + + const int search = static_cast(window_samples * kSplitSearchFraction); + const int search_wide = static_cast(window_samples * kSplitSearchFractionWide); + + const int overlap_samples = std::min(kOverlapMs * sample_rate / 1000, window_samples / 2); + + // Windows are stitched in TOKEN space, then decoded once at the end. Two + // reasons: the seam search needs ids, and a single decode reproduces the + // tokenizer's own spacing exactly. Gluing decoded strings would either + // inject separators the model never emitted (fatal for unspaced scripts + // like Chinese and Japanese, both advertised by this family) or require + // re-deriving word boundaries from text. + std::vector merged_ids; + int64_t mel_us = 0; + int64_t encode_us = 0; + int64_t decode_us = 0; + int offset = 0; + + // Fold the accumulated ids into the session result. Windowing is an + // implementation detail: the caller still sees a single text-only segment, + // matching max_timestamp_kind == NONE. + auto commit_windows = [&]() { + cc->clear_result(); + cc->t_mel_us = mel_us; + cc->t_encode_us = encode_us; + cc->t_decode_us = decode_us; + + if (merged_ids.empty()) { + return; + } + + // raw_text keeps the untrimmed decode (the transcribe_raw_text + // contract); full_text applies the same single leading-space trim + // run_window() applies to a one-window run. + std::string raw_joined = cm->tok.decode(merged_ids.data(), static_cast(merged_ids.size())); + std::string full = raw_joined; + if (!full.empty() && full.front() == ' ') { + full.erase(full.begin()); + } + + transcribe_session::SegmentEntry seg; + seg.t0_ms = 0; + seg.t1_ms = 0; + seg.first_token = 0; + seg.n_tokens = 0; + seg.first_word = 0; + seg.n_words = 0; + seg.text = full; + + cc->raw_text = std::move(raw_joined); + cc->segments.push_back(std::move(seg)); + cc->full_text = std::move(full); + cc->result_kind = TRANSCRIBE_TIMESTAMPS_NONE; + cc->has_result = true; + }; + + int n_windows = 0; + while (offset < n_samples) { + const int remaining = n_samples - offset; + int take = remaining; + + if (remaining > window_samples) { + // Prefer a pause; widen the search once before settling for a + // mid-speech cut. + int split = pick_split(pcm + offset, remaining, window_samples, search); + if (split < 0) { + split = pick_split(pcm + offset, remaining, window_samples, search_wide); + } + if (split > 0) { + take = split; + } else { + take = window_samples; + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, + "cohere run: no pause near the window boundary at sample %d — " + "cutting mid-speech; a word may be split across windows.", + offset + window_samples); + } + } + if (take <= 0) { + take = std::min(remaining, window_samples); + } + + // Drop the previous window's state *before* running. run_window() can + // return early (abort, mel failure, over-length gate, OOM) before it + // reaches its own clear_result(), and stale text would otherwise be + // captured a second time as if it belonged to this window. clear_result() + // does not touch the timing fields, so zero those explicitly or an + // early-returning window re-adds the previous window's numbers. + cc->clear_result(); + cc->window_ids.clear(); + cc->t_mel_us = 0; + cc->t_encode_us = 0; + cc->t_decode_us = 0; + + const transcribe_status st = run_window(session, pcm + offset, take, params); + ++n_windows; + + mel_us += cc->t_mel_us; + encode_us += cc->t_encode_us; + decode_us += cc->t_decode_us; + + if (!cc->window_ids.empty()) { + if (merged_ids.empty()) { + merged_ids = cc->window_ids; + } else { + // The windows share `overlap_samples` of audio, so the new one + // re-transcribes the tail of the previous. Drop that duplicate + // prefix; when no seam is found the text is appended whole, + // which repeats a little rather than losing anything. + const TokenSeam seam = token_seam(merged_ids, cc->window_ids, seam_search_previous(cm->hparams), + seam_search_current(cm->hparams)); + if (!seam.matched) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, + "cohere run: no token seam at window %d — appending whole " + "(overlap may repeat a few words).", + n_windows); + } + merged_ids.insert(merged_ids.end(), cc->window_ids.begin() + seam.current_skip, cc->window_ids.end()); + } + } + + // OUTPUT_TRUNCATED keeps its partial text and is surfaced at the end + // through was_truncated; any other non-OK status stops the run, but the + // windows already decoded stay readable (the aborted-run contract). + if (st != TRANSCRIBE_OK && st != TRANSCRIBE_ERR_OUTPUT_TRUNCATED) { + commit_windows(); + return st; + } + + // Step back by the overlap so the next window re-reads the tail of this + // one. Never on the final window (nothing follows), and never far + // enough to stall: `take` always exceeds the overlap because + // overlap_samples is capped at half a window. + const int advance = (offset + take >= n_samples) ? take : take - overlap_samples; + offset += advance > 0 ? advance : take; + } + + commit_windows(); + + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, + "cohere run: long-form windowed into %d windows (<= %.0f s each, %d ms overlap)", n_windows, + kTrainedClipSeconds, kOverlapMs); + + return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; +} + // =========================================================================== // Offline batched decode (transcribe_run_batch) // =========================================================================== @@ -1434,20 +1801,50 @@ transcribe_status run_batch_serial(CohereSession * cc, const int * n_samples, int n, const transcribe_run_params * params) { + // Session-level truncation is the OR across utterances; the exact + // per-utterance status stays in batch_results. Same shape as arch/moss. + bool any_truncated = false; + for (int i = 0; i < n; ++i) { if (cc->poll_abort()) { + cc->was_truncated = any_truncated; return TRANSCRIBE_ERR_ABORTED; } + + // Per-utterance state. transcribe_run_batch() resets was_truncated once + // for the whole call, so without this a single truncated utterance would + // leave every later one flagged truncated as well. clear_result() covers + // the case where run() is skipped entirely (invalid args) and the + // previous utterance's text would otherwise be captured again. + cc->clear_result(); + cc->was_truncated = false; + cc->t_mel_us = 0; + cc->t_encode_us = 0; + cc->t_decode_us = 0; + const transcribe_status st = (pcm[i] == nullptr || n_samples[i] <= 0) ? TRANSCRIBE_ERR_INVALID_ARG : run(cc, pcm[i], n_samples[i], params); - if (st == TRANSCRIBE_OK) { - cc->batch_results.push_back(cc->capture_result(st)); - } else { - transcribe_session::ResultSet rs; - rs.status = st; - cc->batch_results.push_back(std::move(rs)); + any_truncated = any_truncated || st == TRANSCRIBE_ERR_OUTPUT_TRUNCATED; + + // Capture unconditionally, matching run_batched_encdec_*: a non-OK + // status that still produced text (OUTPUT_TRUNCATED, ABORTED) keeps that + // partial transcript readable. Discarding it here would contradict + // docs/input-limits.md, which guarantees the partial output survives. + cc->batch_results.push_back(cc->capture_result(st)); + + // A cancel that lands *inside* an utterance ends the batch. The + // poll_abort() at the top of the loop is not enough on its own: an + // edge-triggered callback may report true only once, and run() has + // already consumed it, so the remaining utterances would transcribe + // normally. Returning ABORTED lets the dispatcher fill the remaining + // slots via pad_batch_results_aborted(). + if (st == TRANSCRIBE_ERR_ABORTED) { + cc->was_truncated = any_truncated; + return TRANSCRIBE_ERR_ABORTED; } } + + cc->was_truncated = any_truncated; return TRANSCRIBE_OK; } @@ -1470,7 +1867,23 @@ transcribe_status run_batch(transcribe_session * session, const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && cm->plan.primary_kind != transcribe::BackendKind::Accel && cm->plan.primary_kind != transcribe::BackendKind::Unknown; - if (n == 1 || !cc->decoder_use_flash || !primary_is_gpu || transcribe::debug::enabled()) { + + // Long-form windowing lives in run(), which the serial path calls per + // utterance. The batched decoder assumes exactly one encode per utterance + // and has nowhere to fan one input out into several windows, so a batch + // carrying any long clip goes serial. Without this, the same audio would + // be silently truncated or not depending on backend and batch size. + const int sample_rate = cm->hparams.fe_sample_rate > 0 ? cm->hparams.fe_sample_rate : 16000; + const int window_samples = static_cast(kTrainedClipSeconds * sample_rate); + bool any_long_form = false; + for (int i = 0; i < n; ++i) { + if (n_samples[i] > window_samples) { + any_long_form = true; + break; + } + } + + if (n == 1 || !cc->decoder_use_flash || !primary_is_gpu || transcribe::debug::enabled() || any_long_form) { return run_batch_serial(cc, pcm, n_samples, n, params); } @@ -1539,7 +1952,7 @@ transcribe_status run_batch(transcribe_session * session, if (cm->hparams.enc_pos_emb_max_len > 0) { const int t_enc_pred = cohere_predict_t_enc(mel_nf[b], cm->hparams.enc_subsampling_factor); if (t_enc_pred > cm->hparams.enc_pos_emb_max_len) { - const double max_s = static_cast(cm->caps.max_audio_ms) / 1000.0; + const double max_s = static_cast(cohere_max_audio_ms(cm->hparams)) / 1000.0; transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere run_batch: utterance %d too long — %d encoder frames " "exceed the %d the model supports (~%.0f s max). See " diff --git a/tests/cohere_e2e_smoke.cpp b/tests/cohere_e2e_smoke.cpp index 66fa5ce7..b9ac14df 100644 --- a/tests/cohere_e2e_smoke.cpp +++ b/tests/cohere_e2e_smoke.cpp @@ -130,6 +130,11 @@ const char * const k_jfk_reference_text = constexpr int k_max_edit_distance = 3; +// How many times the sample is repeated to build the long-form clip. Four +// copies of jfk.wav is ~44 s, comfortably past the 35 s trained window, and +// repeated speech is the adversarial case for overlap trimming. +constexpr int k_long_form_repeats = 4; + } // namespace int main() { @@ -321,6 +326,80 @@ int main() { CHECK_EQ_INT(seg.n_tokens, 0); } + // ---- Long-form windowing ----------------------------------------- + // + // Repeat the clip past the trained clip span (35 s) so run() has to split + // it into several windows. The assertions are model-independent invariants + // rather than an exact transcript: windowing must not lose text, must not + // invent separators the model never emitted, and must keep raw_text as the + // untrimmed decode. + { + std::vector long_pcm; + for (int r = 0; r < k_long_form_repeats; ++r) { // 4 x 11 s = 44 s -> at least two windows + long_pcm.insert(long_pcm.end(), pcm.begin(), pcm.end()); + } + + const transcribe_status st = transcribe_run(ctx, long_pcm.data(), static_cast(long_pcm.size()), &rp); + CHECK(st == TRANSCRIBE_OK); + + const char * full_c = transcribe_full_text(ctx); + const char * raw_c = transcribe_raw_text(ctx); + const std::string lf_full = full_c ? full_c : ""; + const std::string lf_raw = raw_c ? raw_c : ""; + + // The clip is the same sentence four times over, so the transcript must + // contain it four times. Counting a distinctive phrase is the assertion + // that matters: a length-only check passes with three repetitions, which + // is precisely the silent-loss defect this path exists to prevent. The + // seam search must trim the ~1 s overlap without mistaking a genuinely + // repeated sentence for it. + CHECK(!lf_full.empty()); + CHECK(lf_full.size() > actual.size()); + + const auto count_occurrences = [](const std::string & haystack, const std::string & needle) { + int n = 0; + size_t pos = haystack.find(needle); + while (pos != std::string::npos) { + ++n; + pos = haystack.find(needle, pos + needle.size()); + } + return n; + }; + CHECK_EQ_INT(count_occurrences(lf_full, "fellow Americans"), k_long_form_repeats); + CHECK_EQ_INT(count_occurrences(lf_full, "ask not"), k_long_form_repeats); + + // Bound the other direction too. Counting phrases catches dropped + // speech but not duplicated speech: when the seam search finds no exact + // token match -- a capitalised word at a sentence start tokenises + // differently from its lower-case form mid-sentence -- the window is + // appended whole and the overlap leaves a short repeated fragment. That + // trade is deliberate (repeating is recoverable, dropping is not), but + // it must stay small. One overlap's worth per seam is ~1 s of speech; + // allow a generous margin and fail if duplication ever runs away. + const size_t expected = actual.size() * static_cast(k_long_form_repeats); + CHECK(lf_full.size() >= expected * 9 / 10); + CHECK(lf_full.size() <= expected * 6 / 5); + + // raw_text must stay the untrimmed decode, exactly as a single-window + // run leaves it: this clip is English, so the tokenizer emits a leading + // word-start space that full_text trims and raw_text keeps. Asserting + // the space is present is what distinguishes accumulating raw decodes + // from accumulating already-trimmed text and re-joining it, which would + // leave raw_text == full_text and quietly break the + // transcribe_raw_text contract. + CHECK(!lf_raw.empty()); + CHECK(lf_raw[0] == ' '); + CHECK(lf_raw.size() == lf_full.size() + 1); + CHECK(lf_raw.compare(1, std::string::npos, lf_full) == 0); + + // Windows are concatenated as raw decodes so the tokenizer's own + // spacing survives untouched. A doubled space at a seam would mean a + // separator was injected on top of it — the bug that a hardcoded ' ' + // join introduces (and which would corrupt unspaced scripts like + // Chinese and Japanese outright). + CHECK(lf_full.find(" ") == std::string::npos); + } + // ---- Verify capability-aware dispatcher rejections --------------- // // A second run with a request that exceeds NONE must be rejected @@ -329,22 +408,23 @@ int main() { // alignment returns ERR_UNSUPPORTED_TIMESTAMPS, not a synthetic // fake entry. // - // Critically, a failed transcribe_run must replace the prior - // result with the empty-sentinel state — not silently leave - // the last successful run's text hanging on the context. This - // is the "transcribe_run replaces the previous result" rule - // from include/transcribe.h applied to the failure path. + // The rejection happens in caller-param validation, which runs BEFORE the + // dispatcher's clear_result(), so the previous snapshot survives: see the + // "Caller-param rejections ... preserve the previous snapshot" contract in + // src/transcribe.cpp. A typo'd timestamp request must not destroy the + // utterance the caller already has. + const std::string before_rejection = transcribe_full_text(ctx) ? transcribe_full_text(ctx) : ""; + CHECK(!before_rejection.empty()); { transcribe_run_params rp2; transcribe_run_params_init(&rp2); rp2.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; const transcribe_status st = transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), &rp2); CHECK(st == TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); - // The prior successful result must be gone: no stale text, - // no stale segments/words/tokens, returned_timestamp_kind - // back to the pre-run sentinel. - CHECK(std::strcmp(transcribe_full_text(ctx), "") == 0); - CHECK_EQ_INT(transcribe_n_segments(ctx), 0); + CHECK_STR_EQ(transcribe_full_text(ctx), before_rejection); + CHECK_EQ_INT(transcribe_n_segments(ctx), 1); + // Cohere advertises NONE, so the preserved result carries no word or + // token substructure either way. CHECK_EQ_INT(transcribe_n_words(ctx), 0); CHECK_EQ_INT(transcribe_n_tokens(ctx), 0); CHECK(transcribe_returned_timestamp_kind(ctx) == TRANSCRIBE_TIMESTAMPS_NONE); @@ -355,8 +435,8 @@ int main() { rp2.timestamps = TRANSCRIBE_TIMESTAMPS_SEGMENT; const transcribe_status st = transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), &rp2); CHECK(st == TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); - CHECK(std::strcmp(transcribe_full_text(ctx), "") == 0); - CHECK_EQ_INT(transcribe_n_segments(ctx), 0); + CHECK_STR_EQ(transcribe_full_text(ctx), before_rejection); + CHECK_EQ_INT(transcribe_n_segments(ctx), 1); } // ---- Teardown ---------------------------------------------------- diff --git a/tests/cohere_smoke.cpp b/tests/cohere_smoke.cpp index c62c0b2d..04c59105 100644 --- a/tests/cohere_smoke.cpp +++ b/tests/cohere_smoke.cpp @@ -51,6 +51,7 @@ #include #include #include +#include #ifndef TRANSCRIBE_TEST_FIXTURES_DIR # error "TRANSCRIBE_TEST_FIXTURES_DIR must be defined by the build system" @@ -384,6 +385,137 @@ int main() { ++g_failures; } + // ----- Long-form window boundary picker ----- + // pick_split() must find a real pause when one exists and must refuse to + // invent one in continuous speech: a minimum always exists, so returning + // it unconditionally would cut a word in half at the window boundary. + { + const int n = 32000; // 2 s @ 16 kHz + const int tgt = 16000; // nominal boundary at 1 s + const int srch = 8000; // search the preceding 0.5 s + std::vector pcm(static_cast(n), 0.5f); + + // Uniform speech: no frame is meaningfully quieter -> refusal. + CHECK_EQ_INT(transcribe::cohere::pick_split(pcm.data(), n, tgt, srch), -1); + + // Carve a silence inside the search span; the cut must land in it. + const int gap_lo = 12000; + const int gap_hi = 15200; + for (int i = gap_lo; i < gap_hi; ++i) { + pcm[static_cast(i)] = 0.0f; + } + const int split = transcribe::cohere::pick_split(pcm.data(), n, tgt, srch); + CHECK(split >= gap_lo && split <= gap_hi); + + // Degenerate inputs are refusals, not arbitrary offsets. + CHECK_EQ_INT(transcribe::cohere::pick_split(pcm.data(), n, n, srch), -1); + CHECK_EQ_INT(transcribe::cohere::pick_split(pcm.data(), n, tgt, 100), -1); + } + + // ----- Long-form window stitching ----- + // token_seam() joins overlapping windows by matching the tail of the + // accumulated hypothesis against the head of the next window. It must trim + // a genuine overlap and must refuse an ambiguous one: text repeated in the + // interior may be speech the speaker actually repeated, and trimming it + // would delete real words. + { + using transcribe::cohere::token_seam; + + // Clean overlap: previous ends 4,5 and current starts 4,5. + { + const auto seam = token_seam({ 1, 2, 3, 4, 5 }, { 4, 5, 6, 7 }, 5, 4); + CHECK(seam.matched); + CHECK_EQ_INT(seam.current_skip, 2); + } + + // Match sits in the previous hypothesis's interior (2,3 is followed by + // 9), so it is not a seam -- refuse rather than delete the tail. + { + const auto seam = token_seam({ 1, 2, 3, 9 }, { 2, 3, 4 }, 4, 3); + CHECK(!seam.matched); + CHECK_EQ_INT(seam.current_skip, 0); + } + + // No overlap at all: append the whole window. + { + const auto seam = token_seam({ 1, 2, 3 }, { 7, 8, 9 }, 3, 3); + CHECK(!seam.matched); + CHECK_EQ_INT(seam.current_skip, 0); + } + + // A single shared token counts only at the very head of the new window. + { + const auto seam = token_seam({ 1, 2, 3 }, { 3, 8, 9 }, 3, 3); + CHECK(seam.matched); + CHECK_EQ_INT(seam.current_skip, 1); + } + + // Degenerate inputs never claim a match. + { + CHECK(!token_seam({}, { 1, 2 }, 4, 4).matched); + CHECK(!token_seam({ 1, 2 }, {}, 4, 4).matched); + CHECK(!token_seam({ 1, 2 }, { 1, 2 }, 0, 4).matched); + } + } + + // ----- Batch cancellation contract ----- + // A cancel that lands *inside* an utterance must end the whole batch, not + // just that utterance. The abort callback here is edge-triggered: it fires + // once and then reports false forever, which is exactly the case a + // poll_abort()-at-loop-top check alone would miss — run() consumes the + // single true, and the remaining utterances would transcribe normally. + // + // Per include/transcribe.h, an aborted batch still exposes one slot per + // input utterance (the dispatcher pads the ones that never ran). + { + transcribe_session_params cp; + transcribe_session_params_init(&cp); + + struct transcribe_session * ctx = nullptr; + if (transcribe_session_init(model, &cp, &ctx) == TRANSCRIBE_OK && ctx != nullptr) { + struct AbortOnce { + int polls = 0; + bool fired = false; + int fire_at = 2; // let the loop-top poll pass, trip inside run() + } gate; + + transcribe_set_abort_callback( + ctx, + [](void * ud) -> bool { + auto * g = static_cast(ud); + if (g->fired) { + return false; // edge-triggered: never true twice + } + if (++g->polls >= g->fire_at) { + g->fired = true; + return true; + } + return false; + }, + &gate); + + std::vector utt(4000, 0.1f); // toy PCM; never fully decoded + const float * pcm3[3] = { utt.data(), utt.data(), utt.data() }; + const int lens3[3] = { static_cast(utt.size()), static_cast(utt.size()), + static_cast(utt.size()) }; + transcribe_run_params rp; + transcribe_run_params_init(&rp); + + const transcribe_status bst = transcribe_run_batch(ctx, pcm3, lens3, 3, &rp); + + CHECK(bst == TRANSCRIBE_ERR_ABORTED); + CHECK_EQ_INT(transcribe_batch_n_results(ctx), 3); + // The trailing utterances never ran, so they must report ABORTED + // rather than a stale OK from a previous utterance. + CHECK(transcribe_batch_status(ctx, 2) == TRANSCRIBE_ERR_ABORTED); + + transcribe_session_free(ctx); + } else { + std::fprintf(stderr, "FAIL: could not init session for batch-abort check\n"); + ++g_failures; + } + } + // ----- Lifecycle ----- transcribe_model_free(model);