diff --git a/THIRD-PARTY-LICENSES.md b/THIRD-PARTY-LICENSES.md index aaa95fb6..47a75013 100644 --- a/THIRD-PARTY-LICENSES.md +++ b/THIRD-PARTY-LICENSES.md @@ -15,6 +15,12 @@ Prebuilt artifacts (the Swift xcframework, the native Python wheels, and the npm platform packages) carry these same texts alongside the binaries — see each binding's packaging for the bundled `LICENSE.ggml` / `LICENSE.miniz` / `LICENSE.silero_vad` files. +The Whisper alignment-head presets come from +[whisper.cpp v1.8.3](https://github.com/ggml-org/whisper.cpp/blob/v1.8.3/src/whisper.cpp), +copyright (c) 2023–2024 The ggml authors, under the +[MIT license](https://github.com/ggml-org/whisper.cpp/blob/v1.8.3/LICENSE) +(the terms are reproduced under ggml below). + Conversion-time tooling also includes `scripts/lib/phonon.py`, adapted from Fermion Research's Apache-2.0 container reader at `FermionResearch/Phonon-2@7e153e4054c0a10db6d47d6fff5c062d4c77c154`. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 9fa946e6..ba5c1028 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -94,6 +94,7 @@ its var is unset. Convention: `TRANSCRIBE__GGUF`. | `TRANSCRIBE_COHERE_GGUF` | `cohere_real_smoke`, `cohere_e2e_smoke` | | `TRANSCRIBE_GRANITE5_CTC_GGUF` | `granite5_ctc_real_smoke`, `granite5_ctc_e2e_smoke` | | `TRANSCRIBE_WHISPER_GGUF` | `whisper_e2e_smoke`, `whisper_tokenize_parity` | +| `TRANSCRIBE_WHISPER_WORD_GGUF` | `whisper_word_timestamps_smoke` (stock OpenAI Whisper GGUF, e.g. Base) | | `TRANSCRIBE_QWEN3_ASR_GGUF` (+ `_0_6B_GGUF` / `_1_7B_GGUF`) | qwen3_asr smokes / parity | | `TRANSCRIBE_MOONSHINE_STREAMING_TINY_GGUF` | moonshine_streaming smokes, `stream_offline_interleave_smoke` | | `TRANSCRIBE_VOXTRAL_REALTIME_GGUF` | `voxtral_realtime_real_smoke`, `stream_offline_interleave_smoke` | diff --git a/docs/models/whisper.md b/docs/models/whisper.md index ac292a36..bea57c06 100644 --- a/docs/models/whisper.md +++ b/docs/models/whisper.md @@ -97,9 +97,7 @@ All Whisper variants support: - **Transcription** of 16 kHz mono WAV input. - **Long-form audio** via 30-second chunked decoding with the prev-context window assembly described in the family doc. -- **Segment timestamps** — the finest granularity the library emits for - Whisper (`max_timestamp_kind = segment`). Word-level timestamps are not - currently exposed. +- **Segment timestamps** — the default under `AUTO`, when no transcript prefix is set. - **Translation** (any supported language → English) on multilingual checkpoints — `.en` variants are transcribe-only. @@ -107,3 +105,18 @@ What's not supported (consistent across the family): real-time streaming (whisper is not streaming-first; chunked 30-second windows only), VAD, speaker diarization. See the family doc for the full runtime contract. + +## Word timestamps + +The OpenAI Whisper GGUF variants advertise `TRANSCRIBE_TIMESTAMPS_WORD` +as their maximum timestamp granularity. Request it with +`transcribe_run_params::timestamps` or `--timestamps word` in the CLI. +The decoder aligns the accepted text to the audio with a second cross-attention +pass and dynamic time warping (DTW). Word times are in milliseconds relative to +the input, including across 30-second windows. Languages such as Japanese +use complete UTF-8 token pieces as timing units. + +This mode adds decoder work and uses the serial path for offline batches. +It requires a model-specific alignment-head preset, so legacy `.bin` models +and other Whisper variants advertise segment granularity. A transcript prefix +requires `AUTO` or `NONE`. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 68606dcc..39e05df6 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -97,6 +97,7 @@ add_library(transcribe arch/whisper/weights.cpp arch/whisper/encoder.cpp arch/whisper/decoder.cpp + arch/whisper/alignment.cpp arch/whisper/bin_load.cpp arch/whisper/public.cpp arch/moonshine/model.cpp diff --git a/src/arch/whisper/alignment.cpp b/src/arch/whisper/alignment.cpp new file mode 100644 index 00000000..121962a3 --- /dev/null +++ b/src/arch/whisper/alignment.cpp @@ -0,0 +1,363 @@ +#include "alignment.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace transcribe::whisper { + +namespace { + +struct alignment_preset { + std::string_view variant; + int n_layers; + int n_heads; + const alignment_head * heads; + size_t count; +}; + +// OpenAI's alignment-head coordinates, as tabulated in ggml-org/whisper.cpp +// (MIT, Copyright (c) 2023-2024 The ggml authors). +// https://github.com/ggml-org/whisper.cpp/blob/v1.8.3/src/whisper.cpp#L360-L371 +constexpr alignment_head k_tiny_en[] = { + { 1, 0 }, + { 2, 0 }, + { 2, 5 }, + { 3, 0 }, + { 3, 1 }, + { 3, 2 }, + { 3, 3 }, + { 3, 4 }, +}; +constexpr alignment_head k_tiny[] = { + { 2, 2 }, + { 3, 0 }, + { 3, 2 }, + { 3, 3 }, + { 3, 4 }, + { 3, 5 }, +}; +constexpr alignment_head k_base_en[] = { + { 3, 3 }, + { 4, 7 }, + { 5, 1 }, + { 5, 5 }, + { 5, 7 }, +}; +constexpr alignment_head k_base[] = { + { 3, 1 }, + { 4, 2 }, + { 4, 3 }, + { 4, 7 }, + { 5, 1 }, + { 5, 2 }, + { 5, 4 }, + { 5, 6 }, +}; +constexpr alignment_head k_small_en[] = { + { 6, 6 }, + { 7, 0 }, + { 7, 3 }, + { 7, 8 }, + { 8, 2 }, + { 8, 5 }, + { 8, 7 }, + { 9, 0 }, + { 9, 4 }, + { 9, 8 }, + { 9, 10 }, + { 10, 0 }, + { 10, 1 }, + { 10, 2 }, + { 10, 3 }, + { 10, 6 }, + { 10, 11 }, + { 11, 2 }, + { 11, 4 }, +}; +constexpr alignment_head k_small[] = { + { 5, 3 }, + { 5, 9 }, + { 8, 0 }, + { 8, 4 }, + { 8, 7 }, + { 8, 8 }, + { 9, 0 }, + { 9, 7 }, + { 9, 9 }, + { 10, 5 }, +}; +constexpr alignment_head k_medium_en[] = { + { 11, 4 }, + { 14, 1 }, + { 14, 12 }, + { 14, 14 }, + { 15, 4 }, + { 16, 0 }, + { 16, 4 }, + { 16, 9 }, + { 17, 12 }, + { 17, 14 }, + { 18, 7 }, + { 18, 10 }, + { 18, 15 }, + { 20, 0 }, + { 20, 3 }, + { 20, 9 }, + { 20, 14 }, + { 21, 12 }, +}; +constexpr alignment_head k_medium[] = { + { 13, 15 }, + { 15, 4 }, + { 15, 15 }, + { 16, 1 }, + { 20, 0 }, + { 23, 4 }, +}; +constexpr alignment_head k_large_v1[] = { + { 9, 19 }, + { 11, 2 }, + { 11, 4 }, + { 11, 17 }, + { 22, 7 }, + { 22, 11 }, + { 22, 17 }, + { 23, 2 }, + { 23, 15 }, +}; +constexpr alignment_head k_large_v2[] = { + { 10, 12 }, + { 13, 17 }, + { 16, 11 }, + { 16, 12 }, + { 16, 13 }, + { 17, 15 }, + { 17, 16 }, + { 18, 4 }, + { 18, 11 }, + { 18, 19 }, + { 19, 11 }, + { 21, 2 }, + { 21, 3 }, + { 22, 3 }, + { 22, 9 }, + { 22, 12 }, + { 23, 5 }, + { 23, 7 }, + { 23, 13 }, + { 25, 5 }, + { 26, 1 }, + { 26, 12 }, + { 27, 15 }, +}; +constexpr alignment_head k_large_v3[] = { + { 7, 0 }, + { 10, 17 }, + { 12, 18 }, + { 13, 12 }, + { 16, 1 }, + { 17, 14 }, + { 19, 11 }, + { 21, 4 }, + { 24, 1 }, + { 25, 6 }, +}; +constexpr alignment_head k_large_v3_turbo[] = { + { 2, 4 }, + { 2, 11 }, + { 3, 3 }, + { 3, 6 }, + { 3, 11 }, + { 3, 14 }, +}; + +constexpr alignment_preset k_presets[] = { + { "whisper-tiny.en", 4, 6, k_tiny_en, std::size(k_tiny_en) }, + { "whisper-tiny", 4, 6, k_tiny, std::size(k_tiny) }, + { "whisper-base.en", 6, 8, k_base_en, std::size(k_base_en) }, + { "whisper-base", 6, 8, k_base, std::size(k_base) }, + { "whisper-small.en", 12, 12, k_small_en, std::size(k_small_en) }, + { "whisper-small", 12, 12, k_small, std::size(k_small) }, + { "whisper-medium.en", 24, 16, k_medium_en, std::size(k_medium_en) }, + { "whisper-medium", 24, 16, k_medium, std::size(k_medium) }, + { "whisper-large-v1", 32, 20, k_large_v1, std::size(k_large_v1) }, + { "whisper-large-v2", 32, 20, k_large_v2, std::size(k_large_v2) }, + { "whisper-large-v3", 32, 20, k_large_v3, std::size(k_large_v3) }, + { "whisper-large-v3-turbo", 4, 20, k_large_v3_turbo, std::size(k_large_v3_turbo) }, +}; + +// Whisper smooths each normalized token row along the encoder frame axis. +constexpr int k_median_radius = 3; + +float median_at(const std::vector & row, int frame) { + const int64_t last_frame = static_cast(row.size()) - 1; + if (last_frame < k_median_radius) { + return row[frame]; + } + + std::array window; + for (int offset = -k_median_radius; offset <= k_median_radius; ++offset) { + int64_t index = static_cast(frame) + offset; + if (index < 0) { + index = -index; + } else if (index > last_frame) { + index = 2 * last_frame - index; + } + window[offset + k_median_radius] = row[static_cast(index)]; + } + std::nth_element(window.begin(), window.begin() + k_median_radius, window.end()); + return window[k_median_radius]; +} + +enum class step : uint8_t { FRAME, TOKEN, BOTH }; + +void dtw_boundaries(const std::vector & scores, + int n_rows, + int n_frames, + int n_heads, + std::vector & token_frames) { + // Store one direction per cell and keep only the two live cost rows. + const size_t width = static_cast(n_frames) + 1; + const double inf = std::numeric_limits::infinity(); + std::vector previous(width, inf); + std::vector current(width, inf); + std::vector trace(scores.size()); + previous[0] = 0.0; + + for (int row = 0; row < n_rows; ++row) { + current[0] = inf; + for (int frame = 0; frame < n_frames; ++frame) { + const size_t index = static_cast(row) * n_frames + frame; + const size_t col = static_cast(frame) + 1; + + // Among minimum-cost steps, prefer a frame advance, then a diagonal. + double cost = current[col - 1]; + step direction = step::FRAME; + if (previous[col - 1] < cost) { + cost = previous[col - 1]; + direction = step::BOTH; + } + if (previous[col] < cost) { + cost = previous[col]; + direction = step::TOKEN; + } + current[col] = cost - static_cast(scores[index]) / n_heads; + trace[index] = direction; + } + std::swap(previous, current); + } + + token_frames.resize(n_rows); + int row = n_rows - 1; + int frame = n_frames - 1; + while (true) { + token_frames[row] = frame; + if (row == 0 && frame == 0) { + break; + } + const step direction = trace[static_cast(row) * n_frames + frame]; + if (direction != step::FRAME) { + --row; + } + if (direction != step::TOKEN) { + --frame; + } + } +} + +} // namespace + +std::vector alignment_heads_for_variant(std::string_view variant, int n_layers, int n_heads) { + // The converter's original large GGUF uses the v1 checkpoint. + if (variant == "whisper-large") { + variant = "whisper-large-v1"; + } + for (const auto & preset : k_presets) { + if (variant == preset.variant && n_layers == preset.n_layers && n_heads == preset.n_heads) { + return { preset.heads, preset.heads + preset.count }; + } + } + return {}; +} + +bool align_attention(const std::vector & attention, + int n_heads, + int n_tokens, + int n_frames, + int first_alignment_row, + int n_alignment_rows, + std::vector & token_frames) { + token_frames.clear(); + if (n_heads <= 0 || n_tokens <= 0 || n_frames <= 0 || first_alignment_row < 0 || n_alignment_rows <= 0 || + first_alignment_row > n_tokens || n_alignment_rows > n_tokens - first_alignment_row) { + return false; + } + if (static_cast(n_tokens) > attention.size() / n_frames) { + return false; + } + const size_t head_size = static_cast(n_tokens) * n_frames; + if (static_cast(n_heads) > attention.size() / head_size || head_size * n_heads != attention.size()) { + return false; + } + + // Attention processing follows OpenAI Whisper (MIT, Copyright (c) 2022 OpenAI). + // https://github.com/openai/whisper/blob/v20250625/whisper/timing.py#L192-L212 + // Accumulate smoothed rows one head at a time to avoid duplicating the input. + std::vector scores(static_cast(n_alignment_rows) * n_frames, 0.0f); + std::vector mean(n_frames); + std::vector inv_std(n_frames); + std::vector normalized(n_frames); + + for (int head = 0; head < n_heads; ++head) { + const float * values = attention.data() + head_size * head; + std::fill(mean.begin(), mean.end(), 0.0); + for (int token = 0; token < n_tokens; ++token) { + const float * row = values + static_cast(token) * n_frames; + for (int frame = 0; frame < n_frames; ++frame) { + const float value = row[frame]; + if (!std::isfinite(value) || value < 0.0f || value > 1.0f) { + return false; + } + mean[frame] += value; + } + } + for (double & value : mean) { + value /= n_tokens; + } + + std::fill(inv_std.begin(), inv_std.end(), 0.0); + for (int token = 0; token < n_tokens; ++token) { + const float * row = values + static_cast(token) * n_frames; + for (int frame = 0; frame < n_frames; ++frame) { + const double delta = row[frame] - mean[frame]; + inv_std[frame] += delta * delta; + } + } + for (double & value : inv_std) { + // Constant columns carry no evidence for any token at this frame. + value = value > 0.0 ? 1.0 / std::sqrt(value / n_tokens) : 0.0; + } + + for (int token = 0; token < n_alignment_rows; ++token) { + const float * row = values + static_cast(first_alignment_row + token) * n_frames; + float * score = scores.data() + static_cast(token) * n_frames; + for (int frame = 0; frame < n_frames; ++frame) { + normalized[frame] = static_cast((row[frame] - mean[frame]) * inv_std[frame]); + } + for (int frame = 0; frame < n_frames; ++frame) { + score[frame] += median_at(normalized, frame); + } + } + } + + dtw_boundaries(scores, n_alignment_rows, n_frames, n_heads, token_frames); + return true; +} + +} // namespace transcribe::whisper diff --git a/src/arch/whisper/alignment.h b/src/arch/whisper/alignment.h new file mode 100644 index 00000000..07494ec8 --- /dev/null +++ b/src/arch/whisper/alignment.h @@ -0,0 +1,33 @@ +#ifndef TRANSCRIBE_WHISPER_ALIGNMENT_H +#define TRANSCRIBE_WHISPER_ALIGNMENT_H + +#include +#include + +namespace transcribe::whisper { + +struct alignment_head { + int layer; + int head; +}; + +// Empty for variants without a known preset or with mismatched decoder geometry. +std::vector alignment_heads_for_variant(std::string_view variant, int n_layers, int n_heads); + +// Attention is contiguous [head][decoder token][encoder frame], with each row +// softmaxed over valid encoder frames. All decoder rows contribute to normalization. +// DTW uses the given subrange, normally no_timestamps followed by the text tokens. +// On success, token_frames holds the first frame reached in each DTW row. Text +// token i spans [token_frames[i], token_frames[i + 1]). Invalid input returns false +// and clears token_frames. +bool align_attention(const std::vector & attention, + int n_heads, + int n_tokens, + int n_frames, + int first_alignment_row, + int n_alignment_rows, + std::vector & token_frames); + +} // namespace transcribe::whisper + +#endif diff --git a/src/arch/whisper/capabilities.cpp b/src/arch/whisper/capabilities.cpp index 68c4c582..4ff1ee2e 100644 --- a/src/arch/whisper/capabilities.cpp +++ b/src/arch/whisper/capabilities.cpp @@ -15,9 +15,8 @@ void apply_family_invariants(transcribe_model & model) { caps.supports_translate = true; // supports_streaming left at its zero-init default (false). - // Segment timestamps via the timestamp-token stream (ids 50364+). - // Word-level DTW timing is not yet implemented; cap at SEGMENT so the - // dispatcher rejects WORD-grain requests rather than falling back. + // Segment timestamps are available for all variants. The GGUF loader + // advertises word timestamps when it can identify an alignment preset. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_SEGMENT; // Run knobs are reached through transcribe_whisper_run_ext. No PNC/ITN diff --git a/src/arch/whisper/decoder.cpp b/src/arch/whisper/decoder.cpp index 4808dbbf..abbe1bb1 100644 --- a/src/arch/whisper/decoder.cpp +++ b/src/arch/whisper/decoder.cpp @@ -16,6 +16,7 @@ #include "decoder.h" +#include "alignment.h" #include "conformer/conformer.h" #include "encoder.h" // for transcribe::conformer namespace alias context #include "ggml.h" @@ -368,21 +369,27 @@ ggml_tensor * mha_self_step(ggml_context * ctx, return o; } +struct CrossAttention { + ggml_tensor * queries = nullptr; + ggml_tensor * keys = nullptr; +}; + // Cross-attention reading the pre-populated cross KV cache. // x: [d_model, n_tokens] query input -> returns [d_model, n_tokens]. -ggml_tensor * mha_cross_cached(ggml_context * ctx, - ggml_tensor * x, - WhisperKvCache & kv_cache, - ggml_tensor * mask, - ggml_tensor * q_w, - ggml_tensor * q_b, - ggml_tensor * out_w, - ggml_tensor * out_b, - int n_heads, - int d_model, - int il, - int T_enc, - bool use_flash) { +ggml_tensor * mha_cross_cached(ggml_context * ctx, + ggml_tensor * x, + const WhisperKvCache & kv_cache, + ggml_tensor * mask, + ggml_tensor * q_w, + ggml_tensor * q_b, + ggml_tensor * out_w, + ggml_tensor * out_b, + int n_heads, + int d_model, + int il, + int T_enc, + bool use_flash, + CrossAttention * attention = nullptr) { const int head_dim = d_model / n_heads; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); const int64_t n_tokens = x->ne[1]; @@ -391,7 +398,6 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, // written). FA reads the padded shape — kernel-friendly seq dim at the cost // of small dilution unless the caller masks the padded slots. const int T_enc_pad = kv_cache.T_enc_pad > 0 ? kv_cache.T_enc_pad : T_enc; - (void) T_enc; // Q from current input; K / V read from the cross cache. ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x); @@ -407,6 +413,11 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, ggml_view_3d(ctx, kv_cache.cross_k, head_dim, T_enc_pad, n_heads, k_elem * d_model, k_elem * head_dim, k_elem * static_cast(static_cast(il) * T_enc_pad * d_model)); + if (attention != nullptr) { + attention->queries = Q; + attention->keys = K; + } + const size_t v_elem = ggml_element_size(kv_cache.cross_v); ggml_tensor * V = ggml_view_3d(ctx, kv_cache.cross_v, head_dim, T_enc_pad, n_heads, v_elem * d_model, v_elem * head_dim, @@ -435,6 +446,108 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, } // namespace +AlignmentBuild build_decoder_alignment_graph(ggml_context * ctx, + const WhisperWeights & w, + const WhisperHParams & hp, + const WhisperKvCache & kv_cache, + int seq_len, + int n_audio_frames, + const std::vector & heads, + bool use_flash) { + AlignmentBuild ab{}; + const int T_enc = kv_cache.T_enc; + const int d_model = hp.dec_d_model; + const int n_heads = hp.dec_n_heads; + if (ctx == nullptr || seq_len <= 0 || seq_len > hp.dec_max_target_positions || n_audio_frames <= 0 || + n_audio_frames > T_enc || kv_cache.n_batch != 1 || kv_cache.cross_k == nullptr || kv_cache.cross_v == nullptr || + heads.empty()) { + return ab; + } + + int last_layer = 0; + for (const auto & head : heads) { + if (head.layer < 0 || head.layer >= static_cast(w.dec_blocks.size()) || head.head < 0 || + head.head >= n_heads) { + return ab; + } + last_layer = std::max(last_layer, head.layer); + } + + ab.graph = ggml_new_graph_custom(ctx, 8192, false); + if (ab.graph == nullptr) { + return ab; + } + ab.attention_out.resize(heads.size()); + ab.token_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, seq_len); + named(ab.token_ids_in, "alignment.token_ids"); + ggml_set_input(ab.token_ids_in); + + ggml_tensor * causal_mask = nullptr; + if (seq_len > 1) { + ab.causal_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, seq_len, seq_len); + named(ab.causal_mask_in, "alignment.causal_mask"); + ggml_set_input(ab.causal_mask_in); + causal_mask = ggml_cast(ctx, ab.causal_mask_in, GGML_TYPE_F16); + } + + ggml_tensor * cross_mask = nullptr; + if (last_layer > 0 && kv_cache.T_enc_pad > T_enc) { + ab.cross_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kv_cache.T_enc_pad, seq_len); + named(ab.cross_mask_in, "alignment.cross_mask"); + ggml_set_input(ab.cross_mask_in); + cross_mask = ggml_cast(ctx, ab.cross_mask_in, GGML_TYPE_F16); + } + + ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w, ab.token_ids_in); + ggml_tensor * pos_emb = ggml_view_2d(ctx, w.dec_top.pos_emb_w, d_model, seq_len, w.dec_top.pos_emb_w->nb[1], 0); + ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb); + + const int head_dim = d_model / n_heads; + const float scale = 1.0f / std::sqrt(static_cast(head_dim)); + for (int i = 0; i <= last_layer; ++i) { + const auto & b = w.dec_blocks[i]; + + ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, b.norm_self_b); + y = mha_decoder(ctx, y, /*context=*/nullptr, causal_mask, b.self_q_w, b.self_q_b, b.self_k_w, b.self_v_w, + b.self_v_b, b.self_out_w, b.self_out_b, n_heads, d_model, use_flash); + x = ggml_add(ctx, x, y); + + CrossAttention attention; + y = layer_norm(ctx, x, b.norm_cross_w, b.norm_cross_b); + y = mha_cross_cached(ctx, y, kv_cache, cross_mask, b.cross_q_w, b.cross_q_b, b.cross_out_w, b.cross_out_b, + n_heads, d_model, i, T_enc, use_flash, &attention); + + // Use only the selected heads and real audio frames for the alignment + // softmax. The full cross-attention is still needed by subsequent layers. + for (size_t j = 0; j < heads.size(); ++j) { + if (heads[j].layer != i) { + continue; + } + const size_t head = static_cast(heads[j].head); + ggml_tensor * q = ggml_view_2d(ctx, attention.queries, head_dim, seq_len, attention.queries->nb[1], + head * attention.queries->nb[2]); + ggml_tensor * k = ggml_view_2d(ctx, attention.keys, head_dim, n_audio_frames, attention.keys->nb[1], + head * attention.keys->nb[2]); + ggml_tensor * scores = ggml_mul_mat(ctx, k, q); + ggml_tensor * probs = ggml_soft_max_ext(ctx, scores, nullptr, scale, 0.0f); + ggml_set_output(probs); + ggml_build_forward_expand(ab.graph, probs); + ab.attention_out[j] = probs; + } + + // No subsequent decoder layer needs this residual or the logits. + if (i == last_layer) { + break; + } + x = ggml_add(ctx, x, y); + y = layer_norm(ctx, x, b.norm_ffn_w, b.norm_ffn_b); + y = ffn(ctx, y, b.ffn_fc1_w, b.ffn_fc1_b, b.ffn_fc2_w, b.ffn_fc2_b); + x = ggml_add(ctx, x, y); + } + + return ab; +} + DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, const WhisperWeights & w, const WhisperHParams & hp, diff --git a/src/arch/whisper/decoder.h b/src/arch/whisper/decoder.h index 42b83326..90a9a446 100644 --- a/src/arch/whisper/decoder.h +++ b/src/arch/whisper/decoder.h @@ -30,6 +30,7 @@ namespace transcribe::whisper { struct WhisperHParams; struct WhisperWeights; struct WhisperKvCache; +struct alignment_head; struct DecoderDumps { ggml_tensor * token_emb = nullptr; @@ -87,6 +88,29 @@ DecoderBuild build_cross_kv_graph(ggml_context * compute_ctx, ggml_tensor * encoder_out, int T_enc); +struct AlignmentBuild { + ggml_tensor * token_ids_in = nullptr; // i32 [seq_len] + ggml_tensor * causal_mask_in = nullptr; // f32 [seq_len, seq_len], optional + ggml_tensor * cross_mask_in = nullptr; // f32 [T_enc_pad, seq_len], optional + std::vector attention_out; // one f32 [n_audio_frames, seq_len] per selected head + ggml_cgraph * graph = nullptr; +}; + +// Forced decoder pass over the alignment prompt and accepted text. Reads the +// populated cross cache without writing the self cache. The residual stream +// attends to all T_enc encoded frames. Output probabilities are normalized +// over the first n_audio_frames only, in the order of the supplied heads. +// Caller fills masks with 0 for visible positions and -inf for future tokens +// (causal_mask_in) or cache padding beyond T_enc (cross_mask_in). +AlignmentBuild build_decoder_alignment_graph(ggml_context * compute_ctx, + const WhisperWeights & weights, + const WhisperHParams & hp, + const WhisperKvCache & kv_cache, + int seq_len, + int n_audio_frames, + const std::vector & heads, + bool use_flash = true); + // Build a KV-cached decoder graph. Works for both prompt pass // (n_tokens > 1, n_past = 0) and step pass (n_tokens = 1, n_past = // prev total). Self-K/V are written into the self cache at position diff --git a/src/arch/whisper/model.cpp b/src/arch/whisper/model.cpp index 75cb2660..77aa95a3 100644 --- a/src/arch/whisper/model.cpp +++ b/src/arch/whisper/model.cpp @@ -4,6 +4,7 @@ // path below; TRANSCRIBE_MEL_FROM_REF can inject a reference mel tensor // to isolate encoder/decoder drift during numerical validation. +#include "alignment.h" #include "decoder.h" #include "encoder.h" #include "ggml-alloc.h" @@ -21,6 +22,7 @@ #include "transcribe-log.h" #include "transcribe-meta.h" #include "transcribe-prompting.h" +#include "transcribe-unicode.h" #include "weights.h" #include "whisper.h" @@ -40,6 +42,7 @@ #include #include #include +#include #include #include #include @@ -494,6 +497,9 @@ transcribe_status whisper_load(Loader & loader, if (const transcribe_status st = read_capability_kv(loader.gguf(), m->caps); st != TRANSCRIBE_OK) { return st; } + if (!alignment_heads_for_variant(m->variant, m->hparams.dec_n_layers, m->hparams.dec_n_heads).empty()) { + m->caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_WORD; + } if (const transcribe_status st = read_languages_kv(loader.gguf(), *m); st != TRANSCRIBE_OK) { return st; } @@ -1107,6 +1113,8 @@ void apply_whisper_timestamp_rules(std::vector & logits, struct WhisperSegmentResult { std::vector segments; std::vector> prev_chunk_segments; + std::vector word_token_ids; + std::vector word_token_ends; // exclusive end of each segment in word_token_ids int segment_offset_frames = 0; }; @@ -1120,7 +1128,8 @@ WhisperSegmentResult whisper_retrieve_segment(const std::vector & gene int seek_num_frames, bool want_segment_timestamps, int timestamp_begin, - int64_t vocab_size) { + int64_t vocab_size, + bool collect_word_tokens = false) { WhisperSegmentResult out; out.segment_offset_frames = seek_num_frames; @@ -1153,7 +1162,12 @@ WhisperSegmentResult whisper_retrieve_segment(const std::vector & gene if (ids_int.empty()) { return std::string(); } - return trim_ascii(tok.decode(ids_int.data(), static_cast(ids_int.size()))); + std::string text = trim_ascii(tok.decode(ids_int.data(), static_cast(ids_int.size()))); + if (collect_word_tokens && !text.empty()) { + out.word_token_ids.insert(out.word_token_ids.end(), ids_int.begin(), ids_int.end()); + out.word_token_ends.push_back(out.word_token_ids.size()); + } + return text; }; const int gn = static_cast(generated_ids.size()); @@ -1248,6 +1262,154 @@ WhisperSegmentResult whisper_retrieve_segment(const std::vector & gene return out; } +transcribe_status whisper_align_tokens(WhisperSession * cc, + const std::vector & heads, + const std::vector & ids, + int first_row, + int n_audio_frames, + std::vector & boundaries) { + const auto * cm = static_cast(cc->model); + const int n_tokens = static_cast(ids.size()); + // Align against the real audio within the encoded window. + const int n_frames = std::min(cc->kv_cache.T_enc, n_audio_frames); + if (!ensure_compute_ctx(cc, 16 * 1024 * 1024)) { + return TRANSCRIBE_ERR_OOM; + } + const AlignmentBuild build = build_decoder_alignment_graph(cc->compute_ctx, cm->weights, cm->hparams, cc->kv_cache, + n_tokens, n_frames, heads, cc->decoder_use_flash); + if (build.graph == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!TRANSCRIBE_ALLOC_SCHED_GRAPH(cc->sched, build.graph, "whisper")) { + return TRANSCRIBE_ERR_OOM; + } + + ggml_backend_tensor_set(build.token_ids_in, ids.data(), 0, ids.size() * sizeof(ids[0])); + if (build.causal_mask_in != nullptr) { + std::vector mask(ids.size() * ids.size(), -INFINITY); + for (size_t q = 0; q < ids.size(); ++q) { + std::fill_n(mask.data() + q * ids.size(), q + 1, 0.0f); + } + ggml_backend_tensor_set(build.causal_mask_in, mask.data(), 0, mask.size() * sizeof(float)); + } + if (build.cross_mask_in != nullptr) { + const int stride = cc->kv_cache.T_enc_pad; + std::vector mask(ids.size() * stride, -INFINITY); + for (size_t q = 0; q < ids.size(); ++q) { + std::fill_n(mask.data() + q * stride, cc->kv_cache.T_enc, 0.0f); + } + ggml_backend_tensor_set(build.cross_mask_in, mask.data(), 0, mask.size() * sizeof(float)); + } + if (const ggml_status st = ggml_backend_sched_graph_compute(cc->sched, build.graph); st != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper run: alignment graph compute failed (%d)", static_cast(st)); + return TRANSCRIBE_ERR_BACKEND; + } + + const size_t head_size = ids.size() * n_frames; + std::vector attention(heads.size() * head_size); + for (size_t head = 0; head < heads.size(); ++head) { + ggml_backend_tensor_get(build.attention_out[head], attention.data() + head * head_size, 0, + head_size * sizeof(float)); + } + // The no_timestamps row predicts the first text token. The last text row + // predicts EOT and provides the final boundary. EOT itself contributes to + // normalization only, like the preceding SOT/language/task rows. + if (!align_attention(attention, static_cast(heads.size()), n_tokens, n_frames, first_row, + n_tokens - first_row - 1, boundaries)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper run: could not align decoded text"); + return TRANSCRIBE_ERR_BACKEND; + } + return TRANSCRIBE_OK; +} + +// Whisper BPE tokens can contain partial UTF-8 codepoints. Hold their bytes +// until the full piece is available before assigning it to a word. +bool whisper_complete_piece(const std::string & piece, size_t & begin, size_t & end, bool & punctuation) { + begin = piece.size(); + end = 0; + punctuation = true; + for (size_t pos = 0; pos < piece.size();) { + const size_t start = pos; + if (piece.size() - start < unicode::len_utf8(piece[start])) { + return false; + } + uint32_t cp = 0; + try { + cp = unicode::cpt_from_utf8(piece, pos); + } catch (const std::invalid_argument &) { + // Treat a malformed byte as text. A truncated, otherwise valid + // sequence is handled above and may complete in the next token. + pos = start + 1; + } + const auto flags = unicode::flags_from_cpt(cp); + if (!flags.is_whitespace()) { + begin = std::min(begin, start); + end = pos; + punctuation = punctuation && (flags.bits & unicode::CptFlags::PUNCTUATION) != 0; + } + } + return true; +} + +void whisper_append_words(WhisperSession & cc, + WhisperSegmentResult & chunk, + const std::vector & frames, + int64_t start_ms, + int64_t end_ms, + std::string_view language) { + const auto & tok = static_cast(cc.model)->tok; + // Same language split as OpenAI Whisper's tokenizer. For languages + // without spaces between words, use complete UTF-8 token pieces. + const bool split_on_unicode = language == "zh" || language == "ja" || language == "th" || language == "lo" || + language == "my" || language == "yue"; + auto time_at = [&](size_t token) { + return std::min(start_ms + int64_t{ 20 } * frames[token], end_ms); + }; + size_t token = 0; + for (size_t s = 0; s < chunk.segments.size(); ++s) { + auto & segment = chunk.segments[s]; + const size_t segment_end = chunk.word_token_ends[s]; + const int seg_index = static_cast(cc.segments.size() + s); + segment.first_word = static_cast(cc.words.size()); + std::string piece; + size_t piece_start = token; + for (; token < segment_end; ++token) { + piece += tok.decode(&chunk.word_token_ids[token], 1); + size_t begin = 0; + size_t end = 0; + bool punctuation = false; + if (!whisper_complete_piece(piece, begin, end, punctuation)) { + continue; + } + if (begin < end) { + const bool has_word = static_cast(cc.words.size()) > segment.first_word; + if (!has_word || (!punctuation && (split_on_unicode || begin > 0))) { + transcribe_session::WordEntry word; + word.text = piece.substr(begin, end - begin); + word.t0_ms = time_at(piece_start); + word.t1_ms = time_at(token + 1); + word.seg_index = seg_index; + cc.words.push_back(std::move(word)); + } else { + auto & word = cc.words.back(); + word.text.append(piece, 0, end); + if (!punctuation) { + word.t1_ms = time_at(token + 1); + } + } + } + piece.clear(); + piece_start = token + 1; + } + segment.n_words = static_cast(cc.words.size()) - segment.first_word; + if (segment.n_words > 0) { + segment.t0_ms = cc.words[segment.first_word].t0_ms; + segment.t1_ms = cc.words.back().t1_ms; + } + } +} + // Render an accepted raw token stream for transcribe_raw_text. Timestamp // tokens are model-emitted but absent from the GGUF tokens array (see the // tokenizer-load note), so Tokenizer::decode silently drops them; synthesize @@ -1403,15 +1565,21 @@ transcribe_status whisper_run(transcribe_session * session, const transcribe_timestamp_kind requested_timestamps = params != nullptr ? params->timestamps : TRANSCRIBE_TIMESTAMPS_NONE; - if (requested_timestamps == TRANSCRIBE_TIMESTAMPS_WORD) { + const bool want_word_timestamps = requested_timestamps == TRANSCRIBE_TIMESTAMPS_WORD; + const auto alignment_heads = + want_word_timestamps ? + alignment_heads_for_variant(cm->variant, cm->hparams.dec_n_layers, cm->hparams.dec_n_heads) : + std::vector{}; + if (want_word_timestamps && alignment_heads.empty()) { return TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS; } // AUTO resolves to NONE under a transcript prefix: the timestamp rules // restart after the prefix and force an initial timestamp near 0 s while // the prefix's speech is still playing, which derails the continuation. - // An explicit SEGMENT request with a prefix is rejected in run_validate. + // Explicit SEGMENT and WORD requests with a prefix are rejected in run_validate. const bool has_prefix = params != nullptr && transcribe::prompting::has_text(params->prefix); - const bool want_segment_timestamps = (requested_timestamps == TRANSCRIBE_TIMESTAMPS_AUTO && !has_prefix) || + const bool want_segment_timestamps = want_word_timestamps || + (requested_timestamps == TRANSCRIBE_TIMESTAMPS_AUTO && !has_prefix) || requested_timestamps == TRANSCRIBE_TIMESTAMPS_SEGMENT; // Multilingual variants emit <|lang|> + <|task|> in the decoder prefix; @@ -1527,7 +1695,9 @@ transcribe_status whisper_run(transcribe_session * session, } cc->raw_text = whisper_decode_raw(cm->tok, all_raw_ids, timestamp_begin); cc->full_text = std::move(text); - cc->result_kind = want_segment_timestamps ? TRANSCRIBE_TIMESTAMPS_SEGMENT : TRANSCRIBE_TIMESTAMPS_NONE; + cc->result_kind = want_word_timestamps ? TRANSCRIBE_TIMESTAMPS_WORD : + want_segment_timestamps ? TRANSCRIBE_TIMESTAMPS_SEGMENT : + TRANSCRIBE_TIMESTAMPS_NONE; cc->has_result = true; if (whisper_perf_enabled()) { @@ -2537,8 +2707,48 @@ transcribe_status whisper_run(transcribe_session * session, // - no pairs: emit one full-chunk segment, advance seek_num_frames. // TIMESTAMPS_NONE includes <|notimestamps|>, so no ts tokens -> "no // pairs" branch, full-chunk advance, no per-chunk segment emission. - WhisperSegmentResult seg_res = whisper_retrieve_segment(generated_ids, cm->tok, time_offset_ms, seek_num_frames, - want_segment_timestamps, timestamp_begin, vocab_size); + WhisperSegmentResult seg_res = + whisper_retrieve_segment(generated_ids, cm->tok, time_offset_ms, seek_num_frames, want_segment_timestamps, + timestamp_begin, vocab_size, want_word_timestamps); + if (want_word_timestamps && !seg_res.word_token_ids.empty()) { + if (cc->poll_abort()) { + commit_result(); + return TRANSCRIBE_ERR_ABORTED; + } + const int consumed_frames = seg_res.segment_offset_frames > 0 ? + std::min(seg_res.segment_offset_frames, seek_num_frames) : + seek_num_frames; + const int64_t audio_end_ms = static_cast(n_samples) * 1000 / cm->caps.native_sample_rate; + const int64_t chunk_end_ms = std::min(audio_end_ms, time_offset_ms + int64_t{ 10 } * consumed_frames); + // Align the selected text with the standard unconditioned Whisper + // prompt. History and the caller's prompt are not part of this audio. + std::vector alignment_ids(prompt_ids.begin() + sot_index, prompt_ids.end()); + const int first_row = static_cast(alignment_ids.size()); + alignment_ids.push_back(cm->hparams.no_timestamps_token_id); + alignment_ids.insert(alignment_ids.end(), seg_res.word_token_ids.begin(), seg_res.word_token_ids.end()); + alignment_ids.push_back(eos_id); + std::string_view language = "en"; + if (is_multilingual && params->task != TRANSCRIBE_TASK_TRANSLATE) { + language = params->language != nullptr && params->language[0] != '\0' ? + std::string_view(params->language) : + std::string_view(cc->detected_language); + } + // Use full 20 ms frames within the PCM bounds, as in Whisper's + // reference alignment. The mel frontend can include a final + // centered frame that extends past the end of the recording. + const int n_audio_frames = std::max(1, static_cast((chunk_end_ms - time_offset_ms) / 20)); + std::vector boundaries; + const transcribe_status st = + whisper_align_tokens(cc, alignment_heads, alignment_ids, first_row, n_audio_frames, boundaries); + if (st != TRANSCRIBE_OK) { + return st; + } + if (cc->poll_abort()) { + commit_result(); + return TRANSCRIBE_ERR_ABORTED; + } + whisper_append_words(*cc, seg_res, boundaries, time_offset_ms, chunk_end_ms, language); + } for (auto & seg : seg_res.segments) { cc->segments.push_back(std::move(seg)); } @@ -2654,7 +2864,7 @@ transcribe_status whisper_run_batch(transcribe_session * session, // ---- Global gates: only what the batched graph genuinely can't do. ---- // Segment timestamps, temperature>0 fallback, and initial_prompt are now // all handled in-batch; what still peels the WHOLE call to serial is the - // device/topology constraints + word timestamps (unsupported) + the + // device/topology constraints + the per-utterance word alignment pass + the // ALL_SEGMENTS-without-condition error (let serial surface it). const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && cm->plan.primary_kind != transcribe::BackendKind::Accel && @@ -3518,15 +3728,16 @@ static transcribe_status whisper_run_validate(const transcribe_session * ctx, co st != TRANSCRIBE_OK) { return st; } - // Transcript prefix: with explicit SEGMENT timestamps the timestamp rules + // Transcript prefix: with explicit SEGMENT or WORD the timestamp rules // restart after the prefix and force an initial timestamp while the // prefix's speech is still playing, which in practice ends the decode // (AUTO resolves to NONE instead; see whisper_run). And like openai, the // prefix may take at most half the decoder window. if (params != nullptr && transcribe::prompting::has_text(params->prefix)) { - if (params->timestamps == TRANSCRIBE_TIMESTAMPS_SEGMENT) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "whisper run: a transcript prefix does not compose with segment timestamps; use NONE or AUTO"); + if (params->timestamps == TRANSCRIBE_TIMESTAMPS_SEGMENT || params->timestamps == TRANSCRIBE_TIMESTAMPS_WORD) { + log_msg( + TRANSCRIBE_LOG_LEVEL_ERROR, + "whisper run: a transcript prefix does not compose with segment or word timestamps; use NONE or AUTO"); return TRANSCRIBE_ERR_INVALID_ARG; } const auto * cm = static_cast(ctx->model); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 110645b0..313dbd7f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1252,9 +1252,26 @@ endif() # Loads a real Whisper GGUF and verifies the public capabilities match the # runtime behavior: language detection runs when no language hint is supplied, # default params use no-timestamp decode, explicit SEGMENT returns segment -# timestamps, and WORD timestamps are rejected until the DTW alignment path lands. +# timestamps, and TOKEN timestamps are rejected. WORD timestamps (DTW) are +# covered by whisper_word_timestamps_smoke. if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) + add_executable(transcribe_whisper_word_timestamps_smoke + whisper_word_timestamps_smoke.cpp) + + target_link_libraries(transcribe_whisper_word_timestamps_smoke + PRIVATE transcribe transcribe-common-example) + + target_compile_definitions(transcribe_whisper_word_timestamps_smoke PRIVATE + "TRANSCRIBE_TEST_SAMPLES_DIR=\"${CMAKE_SOURCE_DIR}/samples\"") + + transcribe_apply_warnings(transcribe_whisper_word_timestamps_smoke) + + add_test(NAME transcribe_whisper_word_timestamps_smoke + COMMAND transcribe_whisper_word_timestamps_smoke) + set_tests_properties(transcribe_whisper_word_timestamps_smoke PROPERTIES + SKIP_RETURN_CODE 77) + add_executable(transcribe_whisper_e2e_smoke whisper_e2e_smoke.cpp) @@ -1596,6 +1613,19 @@ transcribe_apply_warnings(transcribe_whisper_bin_suppress_unit) add_test(NAME transcribe_whisper_bin_suppress_unit COMMAND transcribe_whisper_bin_suppress_unit) +add_executable(transcribe_whisper_alignment_unit + whisper_alignment_unit.cpp) + +target_link_libraries(transcribe_whisper_alignment_unit PRIVATE transcribe) + +target_include_directories(transcribe_whisper_alignment_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_whisper_alignment_unit) + +add_test(NAME transcribe_whisper_alignment_unit + COMMAND transcribe_whisper_alignment_unit) + # ----------------------------------------------------------------------------- # Whisper .bin parser unit test # ----------------------------------------------------------------------------- diff --git a/tests/whisper_alignment_unit.cpp b/tests/whisper_alignment_unit.cpp new file mode 100644 index 00000000..a1d86554 --- /dev/null +++ b/tests/whisper_alignment_unit.cpp @@ -0,0 +1,208 @@ +#include "arch/whisper/alignment.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +using transcribe::whisper::align_attention; +using transcribe::whisper::alignment_head; +using transcribe::whisper::alignment_heads_for_variant; + +int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \ + ++g_failures; \ + } \ + } while (0) + +void test_presets() { + const struct { + std::string_view variant; + int n_layers; + int n_heads; + size_t count; + } presets[] = { + { "whisper-tiny", 4, 6, 6 }, + { "whisper-tiny.en", 4, 6, 8 }, + { "whisper-base", 6, 8, 8 }, + { "whisper-base.en", 6, 8, 5 }, + { "whisper-small", 12, 12, 10 }, + { "whisper-small.en", 12, 12, 19 }, + { "whisper-medium", 24, 16, 6 }, + { "whisper-medium.en", 24, 16, 18 }, + { "whisper-large", 32, 20, 9 }, + { "whisper-large-v1", 32, 20, 9 }, + { "whisper-large-v2", 32, 20, 23 }, + { "whisper-large-v3", 32, 20, 10 }, + { "whisper-large-v3-turbo", 4, 20, 6 }, + }; + + for (const auto & preset : presets) { + const auto heads = alignment_heads_for_variant(preset.variant, preset.n_layers, preset.n_heads); + CHECK(heads.size() == preset.count); + int prev_index = -1; + for (const alignment_head & head : heads) { + CHECK(head.layer >= 0 && head.layer < preset.n_layers); + CHECK(head.head >= 0 && head.head < preset.n_heads); + const int index = head.layer * preset.n_heads + head.head; + CHECK(index > prev_index); + prev_index = index; + } + CHECK(alignment_heads_for_variant(preset.variant, preset.n_layers + 1, preset.n_heads).empty()); + CHECK(alignment_heads_for_variant(preset.variant, preset.n_layers, preset.n_heads + 1).empty()); + } + + const auto base_en = alignment_heads_for_variant("whisper-base.en", 6, 8); + const alignment_head expected[] = { + { 3, 3 }, + { 4, 7 }, + { 5, 1 }, + { 5, 5 }, + { 5, 7 } + }; + if (base_en.size() == 5) { + for (size_t i = 0; i < base_en.size(); ++i) { + CHECK(base_en[i].layer == expected[i].layer && base_en[i].head == expected[i].head); + } + } + CHECK(alignment_heads_for_variant("whisper", 6, 8).empty()); + CHECK(alignment_heads_for_variant("breeze-asr-25", 32, 20).empty()); +} + +void test_reference_attention() { + // OpenAI Whisper v20250625, Base.en, first 12,800 samples of samples/jfk.wav + // at 16 kHz. Decoder input is SOT, no_timestamps, " And", " so", EOT. + // Two preset heads are softmaxed over the 40 valid frames and stored as + // round(probability * 65535). OpenAI's normalization, median filter and DTW + // give the same boundaries before and after quantizing this fixture. + constexpr uint16_t attention_packed[2][5][40] = { + // Decoder (layer, head) (3, 3). + { + { + 5707, 4695, 6863, 3603, 2143, 1492, 1443, 2335, 763, 894, 627, 989, 1541, 2542, + 2676, 3656, 3033, 2988, 2528, 1948, 1160, 845, 529, 395, 479, 504, 467, 674, + 725, 625, 619, 649, 723, 781, 870, 677, 665, 550, 521, 608, + }, { + 2119, 1663, 4720, 3429, 2347, 2468, 2075, 3843, 1363, 1464, 1211, 1184, 1736, 2994, + 3242, 4185, 4251, 3723, 3910, 3050, 2124, 1535, 874, 511, 443, 522, 325, 344, + 289, 284, 451, 572, 447, 328, 366, 413, 302, 174, 125, 129, + }, { + 125, 101, 184, 171, 78, 44, 50, 63, 29, 33, 61, 72, 101, 114, + 98, 44, 28, 25, 25, 24, 28, 43, 41, 27, 16, 35, 441, 3625, + 4219, 4113, 4135, 4742, 8874, 10027, 11708, 6401, 3304, 1210, 784, 292, + }, { + 577, 195, 481, 227, 122, 71, 67, 144, 58, 68, 64, 59, 69, 96, 170, 368, 632, 901, 667, 657, + 891, 1167, 1253, 1261, 965, 805, 293, 80, 34, 12, 12, 17, 38, 107, 181, 850, 3534, 8257, 9838, 30245, + }, { + 7090, 2075, 4603, 2117, 1140, 916, 951, 1161, 430, 416, 436, 595, 712, 1168, + 1928, 2201, 1453, 1172, 1029, 1054, 1145, 985, 830, 833, 813, 1258, 1164, 719, + 350, 202, 190, 138, 124, 189, 367, 457, 693, 1665, 2413, 18356, + }, }, + // Decoder (layer, head) (5, 1). + { + { + 1622, 2711, 12745, 1913, 472, 602, 769, 1112, 613, 1032, 716, 1554, 3134, 6150, + 7885, 7791, 5882, 2737, 1511, 709, 357, 250, 234, 194, 268, 416, 374, 231, + 202, 114, 178, 199, 154, 79, 84, 97, 78, 88, 94, 187, + }, { + 5094, 2289, 1595, 1949, 487, 245, 240, 526, 401, 589, 836, 1013, 1004, 864, + 823, 1590, 4021, 10059, 6158, 3502, 1704, 864, 503, 428, 482, 631, 773, 585, + 644, 712, 718, 1033, 1306, 1604, 2965, 1996, 1228, 777, 996, 2298, + }, { + 228, 126, 139, 212, 59, 30, 27, 81, 35, 44, 34, 33, 32, 21, + 16, 32, 89, 338, 508, 706, 936, 1263, 1576, 1944, 3221, 5347, 6303, 9402, + 8239, 7149, 2712, 1879, 3017, 3358, 3212, 1411, 749, 294, 314, 417, + }, { + 788, 349, 325, 346, 95, 48, 46, 150, 50, 78, 55, 84, 87, 98, 68, 87, 154, 332, 335, 349, + 254, 229, 174, 123, 99, 66, 34, 29, 34, 52, 59, 78, 149, 267, 623, 1100, 2022, 4657, 12960, 38603, + }, { + 4373, 2307, 1982, 3439, 833, 453, 407, 760, 423, 543, 483, 595, 524, 525, + 496, 776, 1415, 2114, 1597, 1241, 776, 549, 376, 325, 335, 314, 251, 326, + 490, 662, 573, 584, 696, 980, 1920, 2641, 3397, 3580, 5771, 15703, + }, }, + }; + + std::vector attention; + for (const auto & head : attention_packed) { + for (const auto & row : head) { + for (uint16_t weight : row) { + attention.push_back(weight / 65535.0f); + } + } + } + std::vector frames; + CHECK(align_attention(attention, 2, 5, 40, 1, 3, frames)); + CHECK((frames == std::vector{ 0, 22, 36 })); +} + +void test_short_attention() { + std::vector frames; + // Two identical rows still occupy distinct DTW rows and may share a frame. + const std::vector repeated = { 0.875f, 0.125f, 0.125f, 0.875f, 0.125f, 0.875f }; + CHECK(align_attention(repeated, 1, 3, 2, 0, 3, frames)); + CHECK((frames == std::vector{ 0, 1, 1 })); + + const std::vector one_frame = { 1.0f, 1.0f, 1.0f }; + CHECK(align_attention(one_frame, 1, 3, 1, 0, 3, frames)); + CHECK((frames == std::vector{ 0, 0, 0 })); +} + +void test_flat_attention() { + // A zero-variance head contributes zero to all cells. Frame advances win + // DTW cost ties, so the path visits every token in the first frame. + const std::vector flat(2 * 4 * 8, 0.125f); + std::vector frames; + CHECK(align_attention(flat, 2, 4, 8, 1, 3, frames)); + CHECK((frames == std::vector{ 0, 0, 0 })); +} + +void test_tied_predecessors() { + // At the center of the second row, the diagonal and token predecessors + // have equal minimum cost. Choosing the frame predecessor would lose it. + const std::vector attention = { 0.625f, 0.25f, 0.125f, 0.125f, 0.25f, 0.625f }; + std::vector frames; + CHECK(align_attention(attention, 1, 2, 3, 0, 2, frames)); + CHECK((frames == std::vector{ 0, 1 })); +} + +void test_invalid_input() { + std::vector frames = { 123 }; + CHECK(!align_attention({}, 1, 1, 1, 0, 1, frames)); + CHECK(frames.empty()); + CHECK(!align_attention({ 1.0f, 1.0f }, 1, 1, 1, 0, 1, frames)); + CHECK(!align_attention({ 1.0f }, 0, 1, 1, 0, 1, frames)); + CHECK(!align_attention({ 1.0f }, 1, 1, 0, 0, 1, frames)); + CHECK(!align_attention({ 1.0f }, 1, 1, 1, -1, 1, frames)); + CHECK(!align_attention({ 1.0f }, 1, 1, 1, 1, 1, frames)); + CHECK(!align_attention({ 1.0f }, 1, 1, 1, 0, 0, frames)); + CHECK(!align_attention({ 1.0f }, 1, std::numeric_limits::max(), 2, 0, 1, frames)); + CHECK(!align_attention({ std::numeric_limits::quiet_NaN(), 1.0f }, 1, 2, 1, 1, 1, frames)); + CHECK(!align_attention({ -0.01f }, 1, 1, 1, 0, 1, frames)); + CHECK(!align_attention({ 1.01f }, 1, 1, 1, 0, 1, frames)); +} + +} // namespace + +int main() { + test_presets(); + test_reference_attention(); + test_short_attention(); + test_flat_attention(); + test_tied_predecessors(); + test_invalid_input(); + + if (g_failures != 0) { + std::fprintf(stderr, "FAILED: %d check(s)\n", g_failures); + return 1; + } + std::fprintf(stderr, "OK\n"); + return 0; +} diff --git a/tests/whisper_e2e_smoke.cpp b/tests/whisper_e2e_smoke.cpp index 6e92524f..bd39ffbe 100644 --- a/tests/whisper_e2e_smoke.cpp +++ b/tests/whisper_e2e_smoke.cpp @@ -94,7 +94,8 @@ int main() { CHECK_EQ_INT(caps->native_sample_rate, 16000); CHECK(caps->supports_language_detect); CHECK(caps->supports_translate); - CHECK_EQ_INT(caps->max_timestamp_kind, TRANSCRIBE_TIMESTAMPS_SEGMENT); + CHECK(caps->max_timestamp_kind == TRANSCRIBE_TIMESTAMPS_SEGMENT || + caps->max_timestamp_kind == TRANSCRIBE_TIMESTAMPS_WORD); CHECK(caps->n_languages > 0); CHECK(caps->languages != nullptr); @@ -177,7 +178,7 @@ int main() { { transcribe_run_params rp; transcribe_run_params_init(&rp); - rp.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + rp.timestamps = TRANSCRIBE_TIMESTAMPS_TOKEN; st = transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), &rp); CHECK(st == TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); } diff --git a/tests/whisper_word_timestamps_smoke.cpp b/tests/whisper_word_timestamps_smoke.cpp new file mode 100644 index 00000000..316aac47 --- /dev/null +++ b/tests/whisper_word_timestamps_smoke.cpp @@ -0,0 +1,171 @@ +// Real-model public ABI check for Whisper's opt-in word timestamps. + +#include "transcribe.h" +#include "wav.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \ + ++failures; \ + } \ + } while (0) + +// The bundled JFK clip contains 22 words and ends near 10.4 s. Also check +// linkage, coverage, and forward motion so flat but in-range DTW fails. +void check_words(const transcribe_session * ctx, int64_t audio_ms) { + CHECK(transcribe_returned_timestamp_kind(ctx) == TRANSCRIBE_TIMESTAMPS_WORD); + CHECK(transcribe_n_words(ctx) >= 20); + CHECK(transcribe_n_tokens(ctx) == 0); + int seen = 0; + int advancing_words = 0; + int64_t previous_end = 0; + for (int i = 0; i < transcribe_n_segments(ctx); ++i) { + transcribe_segment segment; + transcribe_segment_init(&segment); + CHECK(transcribe_get_segment(ctx, i, &segment) == TRANSCRIBE_OK); + CHECK(segment.first_word == seen); + CHECK(segment.n_words > 0); + for (int j = 0; j < segment.n_words; ++j) { + transcribe_word word; + transcribe_word_init(&word); + CHECK(transcribe_get_word(ctx, seen++, &word) == TRANSCRIBE_OK); + CHECK(word.seg_index == i); + CHECK(word.n_tokens == 0); + CHECK(word.text != nullptr && word.text[0] != '\0'); + CHECK(word.t0_ms >= previous_end && word.t1_ms >= word.t0_ms); + CHECK(word.t0_ms >= segment.t0_ms && word.t1_ms <= segment.t1_ms); + CHECK(word.t1_ms <= audio_ms); + advancing_words += word.t1_ms > word.t0_ms; + previous_end = word.t1_ms; + } + } + CHECK(seen == transcribe_n_words(ctx)); + CHECK(advancing_words + 2 >= seen); + CHECK(previous_end > 9000 && previous_end < 10900); +} + +} // namespace + +int main() { + const char * path = std::getenv("TRANSCRIBE_WHISPER_WORD_GGUF"); + if (path == nullptr || path[0] == '\0') { + std::fprintf(stderr, "whisper_word_timestamps_smoke: TRANSCRIBE_WHISPER_WORD_GGUF not set; skipping.\n"); + return 77; + } + + transcribe_model_load_params mp; + transcribe_model_load_params_init(&mp); + mp.backend = TRANSCRIBE_BACKEND_CPU; + transcribe_model * raw = nullptr; + const auto load = transcribe_model_load_file(path, &mp, &raw); + if (load != TRANSCRIBE_OK || raw == nullptr) { + std::fprintf(stderr, "FAIL load: %s\n", transcribe_status_string(load)); + return EXIT_FAILURE; + } + const std::unique_ptr model(raw, transcribe_model_free); + + transcribe_capabilities caps; + transcribe_capabilities_init(&caps); + if (transcribe_model_get_capabilities(model.get(), &caps) != TRANSCRIBE_OK || + std::strcmp(transcribe_model_arch_string(model.get()), "whisper") != 0) { + return EXIT_FAILURE; + } + // This fixture is intended for stock Whisper GGUFs with known heads. + if (caps.max_timestamp_kind != TRANSCRIBE_TIMESTAMPS_WORD) { + std::fprintf(stderr, "whisper_word_timestamps_smoke: model does not advertise word timestamps.\n"); + return EXIT_FAILURE; + } + + transcribe_session_params sp; + transcribe_session_params_init(&sp); + sp.n_threads = 4; + sp.kv_type = TRANSCRIBE_KV_TYPE_F32; + transcribe_session * raw_session = nullptr; + if (transcribe_session_init(model.get(), &sp, &raw_session) != TRANSCRIBE_OK || raw_session == nullptr) { + return EXIT_FAILURE; + } + const std::unique_ptr session(raw_session, + transcribe_session_free); + auto * ctx = session.get(); + + const std::string sample = std::string(TRANSCRIBE_TEST_SAMPLES_DIR) + "/jfk.wav"; + std::vector pcm; + std::string error; + if (!transcribe_cli::load_wav_mono_16k(sample, pcm, error)) { + std::fprintf(stderr, "FAIL wav load: %s\n", error.c_str()); + return EXIT_FAILURE; + } + + transcribe_run_params rp; + transcribe_run_params_init(&rp); + rp.language = "en"; + const int n_samples = static_cast(pcm.size()); + CHECK(transcribe_run(ctx, pcm.data(), n_samples, &rp) == TRANSCRIBE_OK); + CHECK(transcribe_returned_timestamp_kind(ctx) == TRANSCRIBE_TIMESTAMPS_SEGMENT); + CHECK(transcribe_n_words(ctx) == 0); + const std::string text = transcribe_full_text(ctx); + + rp.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK(transcribe_run(ctx, pcm.data(), n_samples, &rp) == TRANSCRIBE_OK); + CHECK(text == transcribe_full_text(ctx)); + check_words(ctx, static_cast(n_samples) * 1000 / 16000); + + rp.prefix = "And so"; + const int word_count = transcribe_n_words(ctx); + CHECK(transcribe_run(ctx, pcm.data(), n_samples, &rp) == TRANSCRIBE_ERR_INVALID_ARG); + CHECK(transcribe_n_words(ctx) == word_count); + rp.prefix = nullptr; + + // The second utterance crosses the 30 s internal window. Its second + // copy of the speech must have absolute word times in the later window. + std::vector delayed = pcm; + delayed.resize(32 * 16000); + delayed.insert(delayed.end(), pcm.begin(), pcm.end()); + const float * audio[] = { pcm.data(), delayed.data() }; + const int lengths[] = { n_samples, static_cast(delayed.size()) }; + CHECK(transcribe_run_batch(ctx, audio, lengths, 2, &rp) == TRANSCRIBE_OK); + CHECK(transcribe_batch_n_results(ctx) == 2); + CHECK(transcribe_batch_n_words(ctx, 0) == word_count); + CHECK(transcribe_batch_returned_timestamp_kind(ctx, 1) == TRANSCRIBE_TIMESTAMPS_WORD); + const int delayed_words = transcribe_batch_n_words(ctx, 1); + CHECK(delayed_words > word_count); + if (delayed_words > 0) { + transcribe_word last; + transcribe_word_init(&last); + CHECK(transcribe_batch_get_word(ctx, 1, delayed_words - 1, &last) == TRANSCRIBE_OK); + CHECK(last.t1_ms > 40000 && last.t1_ms <= static_cast(lengths[1]) * 1000 / 16000); + } + + rp.timestamps = TRANSCRIBE_TIMESTAMPS_SEGMENT; + CHECK(transcribe_run(ctx, pcm.data(), n_samples, &rp) == TRANSCRIBE_OK); + CHECK(transcribe_returned_timestamp_kind(ctx) == TRANSCRIBE_TIMESTAMPS_SEGMENT); + CHECK(transcribe_n_words(ctx) == 0); + + if (caps.supports_language_detect) { + const std::string japanese = std::string(TRANSCRIBE_TEST_SAMPLES_DIR) + "/fleurs-ja.wav"; + if (!transcribe_cli::load_wav_mono_16k(japanese, pcm, error)) { + std::fprintf(stderr, "FAIL Japanese wav load: %s\n", error.c_str()); + return EXIT_FAILURE; + } + rp.language = "ja"; + rp.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK(transcribe_run(ctx, pcm.data(), static_cast(pcm.size()), &rp) == TRANSCRIBE_OK); + CHECK(transcribe_returned_timestamp_kind(ctx) == TRANSCRIBE_TIMESTAMPS_WORD); + // A no-space language should expose multiple complete UTF-8 pieces. + CHECK(transcribe_n_words(ctx) > 5); + } + + return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE; +}