From c37a1e27997370016051067373e72962bbf498b9 Mon Sep 17 00:00:00 2001 From: Braulio Oliveira Date: Sat, 25 Jul 2026 09:11:58 -0300 Subject: [PATCH 1/2] canary: add v2 word timestamps --- docs/models/canary-1b-v2.md | 31 +- docs/porting/families/canary.md | 9 +- examples/cli/main.cpp | 152 +++++- scripts/convert-canary.py | 218 ++++++++- scripts/dump_reference_canary_nemo.py | 20 +- scripts/lib/ref_dump.py | 3 + scripts/validate.py | 252 ++++++++++ src/arch/canary/canary.h | 17 + src/arch/canary/capabilities.cpp | 4 +- src/arch/canary/encoder.cpp | 69 +++ src/arch/canary/encoder.h | 15 + src/arch/canary/model.cpp | 461 +++++++++++++++++- src/arch/canary/weights.cpp | 148 ++++++ src/arch/canary/weights.h | 30 ++ src/transcribe-tokenizer.cpp | 188 +++++++ src/transcribe-tokenizer.h | 14 + tests/CMakeLists.txt | 39 +- tests/canary_timestamp_real_smoke.cpp | 338 +++++++++++++ tests/canary_timestamp_unit.cpp | 117 +++++ .../golden/canary/canary-1b-v2.manifest.json | 11 +- tests/tokenizer_decode_only_unit.cpp | 75 +++ 21 files changed, 2163 insertions(+), 48 deletions(-) create mode 100644 tests/canary_timestamp_real_smoke.cpp create mode 100644 tests/canary_timestamp_unit.cpp diff --git a/docs/models/canary-1b-v2.md b/docs/models/canary-1b-v2.md index 6358f791..a88a434e 100644 --- a/docs/models/canary-1b-v2.md +++ b/docs/models/canary-1b-v2.md @@ -23,8 +23,9 @@ The model takes a 16 kHz mono WAV and produces a transcript. Supports: This is the broadest-coverage canary variant. The other multilingual variants (180m-flash, 1b-flash) cover only English/German/Spanish/French. -Not a streaming model. Word and segment timestamps from the upstream -model are not exposed in the v1 port. +Not a streaming model. Word and segment timestamps are available when the +GGUF includes the upstream auxiliary CTC aligner. Older GGUFs without the +aligner remain text-only and reject explicit timestamp requests. See NVIDIA's [model card](https://huggingface.co/nvidia/canary-1b-v2) for training data, intended use, and upstream evaluation methodology. @@ -44,6 +45,9 @@ pinned 2026-05-08. | Q5_K_M | [canary-1b-v2-Q5_K_M.gguf](https://huggingface.co/handy-computer/canary-1b-v2-gguf/resolve/main/canary-1b-v2-Q5_K_M.gguf) | 798 MB | 1.93% | | Q4_K_M | [canary-1b-v2-Q4_K_M.gguf](https://huggingface.co/handy-computer/canary-1b-v2-gguf/resolve/main/canary-1b-v2-Q4_K_M.gguf) | 701 MB | 1.91% | +These published artifacts predate embedded-aligner support and remain text-only. +Use the [conversion recipe](#convert) below to create a timestamp-enabled GGUF. + WER is measured on the full LibriSpeech test-clean split (2620 utterances) with greedy decoding and no external LM. F32 reference baseline: 1.92%. NVIDIA's self-reported number on the upstream model card is 2.18%; our @@ -58,7 +62,7 @@ cmake --build build # ASR (any supported language) build/bin/transcribe-cli \ - -m models/canary-1b-v2/canary-1b-v2-Q8_0.gguf \ + -m models/canary-1b-v2/canary-1b-v2-timestamps-Q8_0.gguf \ -l en \ samples/jfk.wav @@ -74,6 +78,13 @@ build/bin/transcribe-cli \ --task translate \ -l en --target-language de \ samples/jfk.wav + +# Word timestamps (requires a GGUF converted with the auxiliary aligner) +build/bin/transcribe-cli \ + -m models/canary-1b-v2/canary-1b-v2-Q8_0.gguf \ + -l en --timestamps word \ + --output-json transcript.json \ + samples/jfk.wav ``` If your audio is not already 16 kHz mono WAV, convert it first: @@ -90,7 +101,13 @@ CLI flags specific to canary: - `-l ` — source language code (one of the 25 supported BCP-47 codes). - `--task translate` + `--target-language ` — switch to translation - mode. + mode. Translation returns text only; explicit timestamp requests are rejected + because the auxiliary CTC model aligns source-language speech. +- `--timestamps word|segment` — run the embedded CTC aligner. Timestamp-enabled + v2 conversions advertise word support; GGUFs converted from stripped sources + without `timestamps_asr_model` advertise none. +- `--output-json ` — write structured segment and nested word timings in + integer milliseconds for subtitle integrations. ## Performance @@ -158,13 +175,15 @@ For the full porting writeup, see ```bash uv run --project scripts/envs/canary \ - scripts/convert-canary.py nvidia/canary-1b-v2 --repo-id nvidia/canary-1b-v2 + scripts/convert-canary.py nvidia/canary-1b-v2 \ + models/canary-1b-v2/canary-1b-v2-timestamps-F32.gguf ``` ### Quantize ```bash -uv run scripts/quantize-all.py models/canary-1b-v2/canary-1b-v2-F32.gguf +uv run scripts/quantize-all.py \ + models/canary-1b-v2/canary-1b-v2-timestamps-F32.gguf ``` ### Validate diff --git a/docs/porting/families/canary.md b/docs/porting/families/canary.md index 5351a93c..789d2bdd 100644 --- a/docs/porting/families/canary.md +++ b/docs/porting/families/canary.md @@ -178,7 +178,7 @@ See per-variant `reports/porting/canary//intake.json::known_risks`. Fam 5. **Cross-attention from transformer decoder to encoder output.** Padding mask on the cross-attention keys must propagate from encoder input lengths after the factor-8 subsampling. 6. **Decoding default differs by variant.** canary-1b: beam=5, length_penalty=1.0. flash variants: beam=1 greedy. Mismatch with upstream's measured-config beam will move WER. 7. **Audio length contract.** Native ≤40 s; <1 s is zero-padded to 1 s; long-form needs an external chunker. Streaming is not native — out of scope for v1. -8. **Timestamps are not native to the AED.** canary-1b-v2 / 1b-flash / 180m-flash advertise word + segment timestamps but produce them via a side CTC aligner inside the .nemo archive, not from the AED itself. Out of scope for v1. +8. **Timestamps are not native to the v2 AED.** canary-1b-v2 aligns the AED transcript with a side CTC Conformer stored in the `.nemo`. The aligner tokenizer IDs must remain identical to the AED tokenizer IDs. Flash timestamp-token decoding remains out of scope. 9. **License divergence.** canary-1b is CC-BY-NC-4.0; the rest are CC-BY-4.0. Converter must surface license correctly in `general.license` so non-commercial restrictions on canary-1b ride along with the GGUF. 10. **.nemo archive distribution.** No HF `config.json` / `tokenizer_config.json` / `generation_config.json` (except canary-1b-flash's encoder-only HF Transformers shim). Converter must mirror the parakeet pattern of unpacking the .nemo before extraction. @@ -209,8 +209,9 @@ flagged in the row notes. | Translate | EN→de | `build/bin/transcribe-cli -m models//-F32.gguf --language en --translate --target-language de samples/jfk.wav` | non-empty plausible German translation of JFK | PASS — verified on all four variants 2026-05-07 (`--target-language` wired through `params->target_language`; canary-1 uses explicit `<\|translate\|>` task token, canary2 infers from src!=tgt) | | Translate | EN↔es / EN↔fr | same as above with `--target-language es\|fr` | non-empty plausible translation in the requested language | ACCEPTED GAP — no es/fr reference samples shipped; runtime path identical to the de case | | Translate | X→EN for 24 langs (canary-1b-v2 only) | `build/bin/transcribe-cli -m models/canary-1b-v2/canary-1b-v2-F32.gguf --language --translate --target-language en samples/.wav` | non-empty plausible English translation | ACCEPTED GAP — multilingual sample set not shipped; runtime path identical to the EN→de case | -| Segment timestamps | only canary-1b-v2 / 1b-flash / 180m-flash | n/a — not exposed by runtime in v1 | n/a | SKIP — not exposed by runtime (timestamps decoder path / `_timestamps_asr_model` aligner not ported in v1; intake known_risks notes timestamps are explicitly experimental upstream) | -| Word timestamps | only canary-1b-v2 / 1b-flash / 180m-flash | same | n/a | SKIP — not exposed by runtime (same reason) | +| Segment timestamps | canary-1b-v2 | `build/bin/transcribe-cli -m models/canary-1b-v2/canary-1b-v2-F32.gguf --timestamps segment samples/jfk.wav` | non-zero `segments:` rows matching NeMo within one 80 ms encoder frame | PASS — embedded CTC aligner matched NeMo on `samples/jfk.wav` 2026-07-23 | +| Word timestamps | canary-1b-v2 | `build/bin/transcribe-cli -m models/canary-1b-v2/canary-1b-v2-F32.gguf --timestamps word samples/jfk.wav` | non-zero `words:` rows matching NeMo within one 80 ms encoder frame | PASS — 22 words matched NeMo within one frame; covered by the golden manifest and gated real-model API smoke 2026-07-23 | +| Word/segment timestamps | canary-1b-flash / canary-180m-flash | n/a | n/a | SKIP — upstream AED timestamp-token path is not exposed by the runtime | | Streaming | n/a | not advertised by upstream | n/a | SKIP — not advertised | | Voice activity detection | n/a | not advertised by upstream | n/a | SKIP — not advertised | | Speaker diarization | n/a | not advertised by upstream | n/a | SKIP — not advertised | @@ -222,7 +223,7 @@ flagged in the row notes. - Canary's FastConformer encoder shares the rel_pos relative-shift implementation that Parakeet already needs. The C++ encoder code can probably be re-used with parameter-only changes (n_layers, ffn_dim, etc.). - Use `samples/jfk.wav` for the English smoke test. German/Spanish/French sample clips and the wider 25-language multilingual set need to be sourced for canary-1b-v2 — pull from FLEURS or Common Voice during Stage 2 dumper setup. - canary-1b is CC-BY-NC-4.0. Do not bake "Apache-2.0" or default license metadata into its converter output; surface the non-commercial restriction in `general.license`. -- canary-1b-v2 ships timestamps via a separate CTC aligner inside the .nemo archive (`_timestamps_asr_model`). The model card says these can be deleted to save memory; Stage 2 should decide whether to keep them in the converted GGUF or strip (default: strip — timestamps out of scope for v1). +- canary-1b-v2 ships timestamps via a separate CTC aligner inside the `.nemo` archive (`timestamps_asr_model`, deliberately excluded from the AED `state_dict`). The converter embeds it in the same GGUF under `aligner.*`; older v2 GGUFs without those tensors remain text-only. The F32 aligner adds about 2.33 GiB before quantization. - **Stage 3 tensor-name decisions** (`scripts/convert-canary.py`): - Encoder block names mirror `src/arch/parakeet/weights.cpp` verbatim (FastConformer module class is shared); the converter is a pass-through with no transposes. - Decoder layer names use `dec.layer.{i}.{norm1, self_attn.{q,k,v,o}, norm2, cross_attn.{q,k,v,o}, norm3, ffn.{up,down}}` rather than NeMo's `first_sub_layer`/`second_sub_layer`/`third_sub_layer`. The latter is fine for source comprehension but unfriendly to a reader of the future Stage 4 `src/arch/canary/` code. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index df820403..9ec36e2f 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -17,35 +17,72 @@ #include #include #include +#include #include +#include #include #include namespace { -// Minimal JSON string escape: covers the characters MUST be escaped by -// the JSON spec (quote, backslash, control chars). Transcribed text is -// short UTF-8 in practice; we don't need unicode escaping. +// JSON string escape with UTF-8 validation. A truncated decoder can stop in +// the middle of a byte sequence; replace malformed bytes instead of emitting +// syntactically invalid JSON. std::string json_escape(const char * s) { std::string out; - for (const char * p = s ? s : ""; *p; ++p) { + for (const char * p = s ? s : ""; *p;) { const unsigned char c = static_cast(*p); if (c == '"') { out += "\\\""; + ++p; } else if (c == '\\') { out += "\\\\"; + ++p; } else if (c == '\n') { out += "\\n"; + ++p; } else if (c == '\r') { out += "\\r"; + ++p; } else if (c == '\t') { out += "\\t"; + ++p; } else if (c < 0x20) { char buf[8]; std::snprintf(buf, sizeof(buf), "\\u%04x", c); out += buf; - } else { + ++p; + } else if (c < 0x80) { out += static_cast(c); + ++p; + } else { + int n = 0; + if (c >= 0xC2 && c <= 0xDF) { + n = 2; + } else if (c >= 0xE0 && c <= 0xEF) { + n = 3; + } else if (c >= 0xF0 && c <= 0xF4) { + n = 4; + } + bool valid = n > 0; + for (int i = 1; valid && i < n; ++i) { + const unsigned char continuation = static_cast(p[i]); + valid = continuation >= 0x80 && continuation <= 0xBF; + } + if (valid && n == 3) { + const unsigned char second = static_cast(p[1]); + valid = (c != 0xE0 || second >= 0xA0) && (c != 0xED || second <= 0x9F); + } else if (valid && n == 4) { + const unsigned char second = static_cast(p[1]); + valid = (c != 0xF0 || second >= 0x90) && (c != 0xF4 || second <= 0x8F); + } + if (valid) { + out.append(p, static_cast(n)); + p += n; + } else { + out += "\xEF\xBF\xBD"; + ++p; + } } } return out; @@ -123,6 +160,85 @@ std::string batch_segments_json(const transcribe_session * ctx, int i) { return out; } +const char * timestamp_kind_json_name(transcribe_timestamp_kind kind) { + switch (kind) { + case TRANSCRIBE_TIMESTAMPS_NONE: + return "none"; + case TRANSCRIBE_TIMESTAMPS_SEGMENT: + return "segment"; + case TRANSCRIBE_TIMESTAMPS_WORD: + return "word"; + case TRANSCRIBE_TIMESTAMPS_TOKEN: + return "token"; + case TRANSCRIBE_TIMESTAMPS_AUTO: + return "auto"; + } + return "unknown"; +} + +bool write_result_json(const std::string & path, + const transcribe_session * ctx, + const std::string & language, + transcribe_status status) { + const char * detected = transcribe_detected_language(ctx); + const char * lang = (detected != nullptr && detected[0] != '\0') ? detected : language.c_str(); + const auto kind = transcribe_returned_timestamp_kind(ctx); + const bool truncated = transcribe_was_truncated(ctx); + std::ostringstream json; + json << "{\"status\":\"" << json_escape(transcribe_status_string(status)) + << "\",\"truncated\":" << (truncated ? "true" : "false") + << ",\"aborted\":" << (status == TRANSCRIBE_ERR_ABORTED ? "true" : "false") << ",\"text\":\"" + << json_escape(transcribe_full_text(ctx)) << "\",\"language\":\"" << json_escape(lang) + << "\",\"timestamp_kind\":\"" << timestamp_kind_json_name(kind) << "\",\"segments\":["; + + const bool emit_segments = kind != TRANSCRIBE_TIMESTAMPS_NONE || transcribe_n_speaker_segments(ctx) > 0; + const int n_segments = emit_segments ? transcribe_n_segments(ctx) : 0; + for (int i = 0; i < n_segments; ++i) { + if (i > 0) { + json << ','; + } + transcribe_segment segment; + transcribe_segment_init(&segment); + if (transcribe_get_segment(ctx, i, &segment) != TRANSCRIBE_OK) { + std::fprintf(stderr, "error: cannot read segment %d for JSON output\n", i); + return false; + } + json << "{\"t0_ms\":" << segment.t0_ms << ",\"t1_ms\":" << segment.t1_ms; + if (segment.speaker_id > 0) { + json << ",\"speaker_id\":" << segment.speaker_id; + } + json << ",\"text\":\"" << json_escape(segment.text) << "\",\"words\":["; + for (int j = 0; j < segment.n_words; ++j) { + if (j > 0) { + json << ','; + } + transcribe_word word; + transcribe_word_init(&word); + if (transcribe_get_word(ctx, segment.first_word + j, &word) != TRANSCRIBE_OK) { + std::fprintf(stderr, "error: cannot read word %d for JSON output\n", segment.first_word + j); + return false; + } + json << "{\"t0_ms\":" << word.t0_ms << ",\"t1_ms\":" << word.t1_ms << ",\"text\":\"" + << json_escape(word.text) << "\"}"; + } + json << "]}"; + } + json << "]}\n"; + + std::ofstream out(std::filesystem::u8path(path), std::ios::binary | std::ios::trunc); + if (!out) { + std::fprintf(stderr, "error: cannot open JSON output %s\n", path.c_str()); + return false; + } + out << json.str(); + out.close(); + if (out.fail()) { + std::fprintf(stderr, "error: failed writing JSON output %s\n", path.c_str()); + return false; + } + return true; +} + // ",\"raw_text\":\"...\"" fragment: the model's pre-cleanup decode // (transcribe_raw_text). Emitted only when it differs from the clean text, // so JSONL stays compact for families with no post-processing. @@ -197,6 +313,7 @@ struct cli_args { std::string language; std::string target_language; // --target-language: target lang for translation std::string batch_file; // --batch: one wav path per line + std::string output_json_path; // --output-json: structured single-file result int batch_size = 0; // --batch-size: >1 groups utterances into // transcribe_run_batch calls (offline only). // 0/1 keeps the per-file serial loop. @@ -295,6 +412,8 @@ void print_usage(const char * argv0) { " --timestamps TYPE timestamps: auto, none, segment, word, token (default: auto)\n" " --batch FILE batch mode: FILE has one wav path per line\n" " --batch-jsonl output one JSON line per file (for batch)\n" + " --output-json PATH write a structured single-file result with nested\n" + " segment and word timings in milliseconds\n" " --batch-size N group N utterances into one transcribe_run_batch\n" " call (offline only; 0/1 = per-file serial loop)\n" " --initial-prompt TEXT (whisper) initial prompt text for context biasing\n" @@ -518,6 +637,12 @@ bool parse_args(int argc, char ** argv, cli_args & out) { out.batch_file = v; } else if (a == "--batch-jsonl") { out.batch_jsonl = true; + } else if (a == "--output-json") { + const char * v = take_value(a.c_str()); + if (!v) { + return false; + } + out.output_json_path = v; } else if (a == "--batch-size") { const char * v = take_value(a.c_str()); if (!v) { @@ -656,6 +781,14 @@ bool parse_args(int argc, char ** argv, cli_args & out) { std::fprintf(stderr, "error: cannot combine positional audio.wav with --batch\n"); return false; } + if (!out.output_json_path.empty() && !out.batch_file.empty()) { + std::fprintf(stderr, "error: --output-json is only available for single-file runs\n"); + return false; + } + if (!out.output_json_path.empty() && out.model_path.empty()) { + std::fprintf(stderr, "error: --output-json requires --model\n"); + return false; + } if (out.stream_chunk_ms > 0 && out.repeat > 1) { std::fprintf(stderr, "error: --stream-chunk-ms cannot be combined with --repeat\n"); return false; @@ -1236,7 +1369,8 @@ int main(int argc, char ** argv) { // CLI prints the live tentative text on each such feed. // Families that only commit at finalize keep result_changed // false until the finalize call. - transcribe_status run_st = TRANSCRIBE_OK; + transcribe_status run_st = TRANSCRIBE_OK; + bool json_write_ok = true; if (args.stream_chunk_ms > 0) { struct transcribe_capabilities caps; transcribe_capabilities_init(&caps); @@ -1400,6 +1534,10 @@ int main(int argc, char ** argv) { (tok.text != nullptr) ? tok.text : ""); } } + if (!args.output_json_path.empty() && + !write_result_json(args.output_json_path, ctx, args.language, run_st)) { + json_write_ok = false; + } } transcribe_print_timings(ctx); @@ -1418,7 +1556,7 @@ int main(int argc, char ** argv) { transcribe_session_free(ctx); transcribe_model_free(model); - if (run_st != TRANSCRIBE_OK) { + if (run_st != TRANSCRIBE_OK || !json_write_ok) { return EXIT_FAILURE; } } else { diff --git a/scripts/convert-canary.py b/scripts/convert-canary.py index 15a47f73..f07fd94d 100644 --- a/scripts/convert-canary.py +++ b/scripts/convert-canary.py @@ -209,6 +209,7 @@ def _extract_aggregate(tk) -> dict: return { "tokens": tokens, "scores": scores, "types": types, "offsets": offsets, "codes": codes, "sizes": sizes, + "model": "unigram", "specials": dict(tk.special_tokens), "unk_id": 0, "bos_id": int(tk.bos_id) if tk.bos_id is not None else None, @@ -257,6 +258,7 @@ def _extract_bpe(tk) -> dict: return { "tokens": tokens, "scores": scores, "types": types, "offsets": [0], "codes": ["all"], "sizes": [vocab_size], + "model": "bpe", "specials": specials, "unk_id": int(tk.unk_id) if getattr(tk, "unk_id", None) is not None else None, "bos_id": int(tk.bos_id) if tk.bos_id is not None else None, @@ -300,11 +302,10 @@ def read_hparams(config: dict) -> dict: f"converter only handles AudioToMelSpectrogramPreprocessor" ) - sample_rate = int(pre["sample_rate"]) - window_size = float(pre["window_size"]) - window_stride = float(pre["window_stride"]) - win_length = int(round(window_size * sample_rate)) - hop_length = int(round(window_stride * sample_rate)) + frontend = read_mel_hparams(pre) + sample_rate = frontend["sample_rate"] + win_length = frontend["win_length"] + hop_length = frontend["hop_length"] # transf_encoder is configured but disabled across all four shipping # canary variants. If a future variant turns it on we need new @@ -357,9 +358,10 @@ def read_hparams(config: dict) -> dict: "fe_normalize": str(pre["normalize"]), "fe_dither": float(pre["dither"]), "fe_log": bool(pre.get("log", True)), - "fe_pre_emphasis": 0.97, # NeMo FilterbankFeatures default - "fe_f_min": 0.0, - "fe_f_max": float(sample_rate) / 2.0, + "fe_pre_emphasis": frontend["preemphasis"], + "fe_f_min": frontend["f_min"], + "fe_f_max": frontend["f_max"], + "frontend": frontend, } @@ -473,6 +475,38 @@ def read_hparams(config: dict) -> dict: ("third_sub_layer.dense_out.bias", "ffn.down.bias"), ] +# Bias-free Conformer used by canary-1b-v2's auxiliary CTC aligner. +ALIGNER_BLOCK_TABLE: list[tuple[str, str]] = [ + ("norm_feed_forward1.weight", "norm_ff1.weight"), + ("norm_feed_forward1.bias", "norm_ff1.bias"), + ("feed_forward1.linear1.weight", "ff1.linear1.weight"), + ("feed_forward1.linear2.weight", "ff1.linear2.weight"), + ("norm_self_att.weight", "norm_attn.weight"), + ("norm_self_att.bias", "norm_attn.bias"), + ("self_attn.linear_q.weight", "attn.linear_q.weight"), + ("self_attn.linear_k.weight", "attn.linear_k.weight"), + ("self_attn.linear_v.weight", "attn.linear_v.weight"), + ("self_attn.linear_out.weight", "attn.linear_out.weight"), + ("self_attn.linear_pos.weight", "attn.linear_pos.weight"), + ("self_attn.pos_bias_u", "attn.pos_bias_u"), + ("self_attn.pos_bias_v", "attn.pos_bias_v"), + ("norm_conv.weight", "norm_conv.weight"), + ("norm_conv.bias", "norm_conv.bias"), + ("conv.pointwise_conv1.weight", "conv.pointwise1.weight"), + ("conv.depthwise_conv.weight", "conv.depthwise.weight"), + ("conv.pointwise_conv2.weight", "conv.pointwise2.weight"), + ("conv.batch_norm.weight", "conv.bn.weight"), + ("conv.batch_norm.bias", "conv.bn.bias"), + ("conv.batch_norm.running_mean", "conv.bn.running_mean"), + ("conv.batch_norm.running_var", "conv.bn.running_var"), + ("norm_feed_forward2.weight", "norm_ff2.weight"), + ("norm_feed_forward2.bias", "norm_ff2.bias"), + ("feed_forward2.linear1.weight", "ff2.linear1.weight"), + ("feed_forward2.linear2.weight", "ff2.linear2.weight"), + ("norm_out.weight", "norm_out.weight"), + ("norm_out.bias", "norm_out.bias"), +] + EXPECTED_UNUSED_PREFIXES = ( "preprocessor.", # mel filterbank + window — C++ recomputes @@ -507,6 +541,72 @@ def tensor_to_fp32_numpy(t) -> np.ndarray: return np.ascontiguousarray(arr) +def read_mel_hparams(pre: dict) -> dict: + sample_rate = int(pre["sample_rate"]) + highfreq = pre.get("highfreq") + frontend = { + "num_mels": int(pre["features"]), + "sample_rate": sample_rate, + "n_fft": int(pre["n_fft"]), + "win_length": int(round(float(pre["window_size"]) * sample_rate)), + "hop_length": int(round(float(pre["window_stride"]) * sample_rate)), + "window": str(pre["window"]), + "normalize": str(pre["normalize"]), + "dither": float(pre["dither"]), + "log": bool(pre.get("log", True)), + "preemphasis": float(pre.get("preemph", pre.get("pre_emphasis", 0.97))), + "f_min": float(pre.get("lowfreq") or 0.0), + "f_max": float(highfreq) if highfreq is not None else sample_rate / 2.0, + } + for key in ( + "exact_pad", + "frame_splicing", + "log_zero_guard_type", + "log_zero_guard_value", + "mag_power", + "mel_norm", + "mel_scale", + "pad_to", + "pad_value", + ): + if key in pre: + frontend[key] = pre[key] + return frontend + + +def read_aligner_hparams(config: dict, state_dict: dict) -> dict: + enc = config["encoder"] + pre = config["preprocessor"] + head_w = state_dict["decoder.decoder_layers.0.weight"] + if head_w.ndim != 3 or int(head_w.shape[2]) != 1: + raise ValueError(f"unexpected aligner CTC head shape: {tuple(head_w.shape)}") + + conv_norm_type = str(enc.get("conv_norm_type", "batch_norm")) + return { + "n_layers": int(enc["n_layers"]), + "d_model": int(enc["d_model"]), + "n_heads": int(enc["n_heads"]), + "d_ff": int(enc["d_model"]) * int(enc["ff_expansion_factor"]), + "conv_kernel": int(enc["conv_kernel_size"]), + "subsampling_factor": int(enc["subsampling_factor"]), + "subsampling_channels": int(enc["subsampling_conv_channels"]), + "pos_emb_max_len": int(enc["pos_emb_max_len"]), + "conv_norm_type": conv_norm_type, + "vocab_size": int(head_w.shape[0]), + "blank_id": int(head_w.shape[0]) - 1, + "frontend": read_mel_hparams(pre), + } + + +def tokenizer_piece_list(tokenizer, vocab_size: int) -> list[str]: + if hasattr(tokenizer, "ids_to_tokens"): + return [str(tokenizer.ids_to_tokens([i])[0]) for i in range(vocab_size)] + sp = getattr(tokenizer, "tokenizer", None) + if sp is not None and hasattr(sp, "IdToPiece"): + return [str(sp.IdToPiece(i)) for i in range(vocab_size)] + raise TypeError(f"cannot enumerate aligner tokenizer {type(tokenizer).__name__}") + + # --------------------------------------------------------------------------- # Variant resolution # --------------------------------------------------------------------------- @@ -569,6 +669,39 @@ def convert(model_spec: str, out_path: Path, variant: str, profile: dict, langua sd = model.state_dict() sd_keys = set(sd.keys()) + aligner = getattr(model, "timestamps_asr_model", None) + aligner_sd = None + aligner_sd_keys: set[str] = set() + aligner_hp = None + if aligner is not None: + aligner_sd = aligner.state_dict() + aligner_sd_keys = set(aligner_sd.keys()) + aligner_cfg = OmegaConf.to_container(aligner.cfg, resolve=True) + aligner_hp = read_aligner_hparams(aligner_cfg, aligner_sd) + if variant != "canary-1b-v2": + raise ValueError(f"unexpected auxiliary timestamps model on {variant}") + if aligner_hp["conv_norm_type"] != "batch_norm": + raise ValueError( + f"unsupported aligner conv_norm_type={aligner_hp['conv_norm_type']!r}" + ) + shared_frontend_fields = sorted( + set(aligner_hp["frontend"]) & set(hp["frontend"]) + ) + frontend_mismatches = [ + key + for key in shared_frontend_fields + if aligner_hp["frontend"][key] != hp["frontend"][key] + ] + if frontend_mismatches: + details = ", ".join( + f"{key}: aligner={aligner_hp['frontend'][key]!r}, " + f"AED={hp['frontend'][key]!r}" + for key in frontend_mismatches + ) + raise ValueError(f"aligner and AED mel frontends differ ({details})") + aligner_pieces = tokenizer_piece_list(aligner.tokenizer, aligner_hp["vocab_size"] - 1) + if aligner_pieces != tok["tokens"]: + raise ValueError("aligner tokenizer IDs differ from the AED tokenizer") has_proj = "encoder_decoder_proj.weight" in sd_keys if has_proj: proj_w = sd["encoder_decoder_proj.weight"] @@ -588,6 +721,7 @@ def convert(model_spec: str, out_path: Path, variant: str, profile: dict, langua f"max_pos={hp['dec_max_position']}") print(f"Vocab: {hp['head_num_classes']} ({len(tok['codes'])} sub-tokenizers)") print(f"encoder_decoder_proj: {'present' if has_proj else 'absent (dims match)'}") + print(f"timestamp aligner: {'present' if aligner_hp is not None else 'absent (text-only GGUF)'}") print(f"Writing GGUF to {out_path}") writer = gguf_writer(str(out_path), "canary") @@ -620,17 +754,29 @@ def convert(model_spec: str, out_path: Path, variant: str, profile: dict, langua xlat_langs = [lang for lang in languages if lang not in xlat_exclude] writer.add_array ("stt.translation.target_languages", xlat_langs) writer.add_array ("stt.translation.pairs", english_pivot_pairs(xlat_langs)) - # All shipping variants except canary-1b expose timestamp tokens. - has_timestamps = "<|timestamp|>" in tok["specials"] - writer.add_bool ("stt.capability.timestamps", has_timestamps) + writer.add_bool ("stt.capability.timestamps", aligner_hp is not None) + + writer.add_bool("stt.canary.aligner.present", aligner_hp is not None) + if aligner_hp is not None: + writer.add_uint32("stt.canary.aligner.encoder.n_layers", aligner_hp["n_layers"]) + writer.add_uint32("stt.canary.aligner.encoder.d_model", aligner_hp["d_model"]) + writer.add_uint32("stt.canary.aligner.encoder.n_heads", aligner_hp["n_heads"]) + writer.add_uint32("stt.canary.aligner.encoder.d_ff", aligner_hp["d_ff"]) + writer.add_uint32("stt.canary.aligner.encoder.conv_kernel", aligner_hp["conv_kernel"]) + writer.add_uint32("stt.canary.aligner.encoder.subsampling_factor", aligner_hp["subsampling_factor"]) + writer.add_uint32("stt.canary.aligner.encoder.subsampling_channels", aligner_hp["subsampling_channels"]) + writer.add_uint32("stt.canary.aligner.encoder.pos_emb_max_len", aligner_hp["pos_emb_max_len"]) + writer.add_string("stt.canary.aligner.encoder.conv_norm_type", aligner_hp["conv_norm_type"]) + writer.add_uint32("stt.canary.aligner.vocab_size", aligner_hp["vocab_size"]) + writer.add_uint32("stt.canary.aligner.blank_id", aligner_hp["blank_id"]) # ----- tokenizer.ggml.* ----- - # Use "unigram" so the C++ tokenizer's SentencePiece decode path - # picks up the pieces directly. The aggregate-over-multi-SP nature + # Both SentencePiece model tags share the same C++ decode path. The + # aggregate-over-multi-SP nature # of canary's tokenizer is preserved separately under # stt.canary.tokenizer.{lang_codes,lang_offsets,lang_sizes} for any # future encode() implementation to consume. - writer.add_string("tokenizer.ggml.model", "unigram") + writer.add_string("tokenizer.ggml.model", tok["model"]) writer.add_array ("tokenizer.ggml.tokens", tok["tokens"]) writer.add_array ("tokenizer.ggml.scores", tok["scores"]) writer.add_array ("tokenizer.ggml.token_type", tok["types"]) @@ -800,6 +946,41 @@ def add(nemo_name: str, gguf_name: str) -> None: raise RuntimeError( f"tensor count mismatch: added {n_added}, expected {expected}" ) + + aligner_consumed: set[str] = set() + if aligner_hp is not None and aligner_sd is not None: + aligner_added = 0 + + def add_aligner(nemo_name: str, gguf_name: str, squeeze_head: bool = False) -> None: + nonlocal n_added, bytes_out, aligner_added + if nemo_name not in aligner_sd_keys: + raise KeyError(f"aligner state_dict missing tensor: {nemo_name!r}") + arr = tensor_to_fp32_numpy(aligner_sd[nemo_name]) + if squeeze_head: + arr = np.ascontiguousarray(arr[:, :, 0]) + writer.add_tensor(gguf_name, arr) + aligner_consumed.add(nemo_name) + bytes_out += int(arr.nbytes) + n_added += 1 + aligner_added += 1 + + for nemo_name, gguf_name in PRE_ENCODE_TABLE: + add_aligner(nemo_name, gguf_name.replace("enc.", "aligner.enc.", 1)) + for i in range(aligner_hp["n_layers"]): + for suffix_nemo, suffix_gguf in ALIGNER_BLOCK_TABLE: + add_aligner( + f"encoder.layers.{i}.{suffix_nemo}", + f"aligner.enc.blocks.{i}.{suffix_gguf}", + ) + add_aligner("decoder.decoder_layers.0.weight", "aligner.head.weight", squeeze_head=True) + add_aligner("decoder.decoder_layers.0.bias", "aligner.head.bias") + + aligner_expected = len(PRE_ENCODE_TABLE) + aligner_hp["n_layers"] * len(ALIGNER_BLOCK_TABLE) + 2 + if aligner_added != aligner_expected: + raise RuntimeError( + f"aligner tensor count mismatch: added {aligner_added}, expected {aligner_expected}" + ) + print(f"Added {n_added} tensors ({bytes_out / (1024 * 1024):.1f} MB)") unused = sorted( @@ -815,6 +996,15 @@ def add(nemo_name: str, gguf_name: str) -> None: if len(unused) > 10: print(f" ... and {len(unused) - 10} more", file=sys.stderr) + if aligner_sd is not None: + aligner_unused = sorted( + k for k in (aligner_sd_keys - aligner_consumed) if not is_expected_unused(k) + ) + if aligner_unused: + raise RuntimeError( + "auxiliary aligner tensors were not consumed: " + ", ".join(aligner_unused[:10]) + ) + print("Writing header + KV + tensor info...") writer.write_header_to_file() writer.write_kv_data_to_file() diff --git a/scripts/dump_reference_canary_nemo.py b/scripts/dump_reference_canary_nemo.py index 501fe6c1..208d2c7f 100644 --- a/scripts/dump_reference_canary_nemo.py +++ b/scripts/dump_reference_canary_nemo.py @@ -492,6 +492,7 @@ def cmd_decode(args: argparse.Namespace) -> int: text: str = "" tokens: list[int] | None = None + timestamps: dict[str, Any] | None = None def dump(name: str, t, stage: str) -> None: a = to_np(t) @@ -552,9 +553,14 @@ def dump(name: str, t, stage: str) -> None: print(f" warning: could not set beam_size={args.beam}: {exc}") print("\n Running canonical transcribe() ...") + timestamps_arg = None + if args.toggle_timestamps in ("yes", "no"): + timestamps_arg = args.toggle_timestamps == "yes" output = model.transcribe( str(manifest_path), batch_size=1, + return_hypotheses=True, + timestamps=timestamps_arg, ) if isinstance(output, (list, tuple)) and output: first = output[0] @@ -562,6 +568,9 @@ def dump(name: str, t, stage: str) -> None: first = first[0] if hasattr(first, "text"): text = first.text + raw_timestamps = getattr(first, "timestamp", None) + if isinstance(raw_timestamps, dict): + timestamps = raw_timestamps raw_tokens = getattr(first, "y_sequence", None) if raw_tokens is not None: try: @@ -579,6 +588,15 @@ def dump(name: str, t, stage: str) -> None: if manifest_path is not None and manifest_path.exists(): manifest_path.unlink() + if not args.skip_transcript and args.toggle_timestamps == "yes": + if not isinstance(timestamps, dict) or any( + not isinstance(timestamps.get(level), list) or not timestamps[level] + for level in ("word", "segment") + ): + raise RuntimeError( + "Canary timestamp decoding returned no non-empty word and segment lists" + ) + # ----- Decoder intermediates (captured from the transcribe() pass) --- for i in dec_blocks: for sub in ("self_attn", "cross_attn", "ffn"): @@ -588,7 +606,7 @@ def dump(name: str, t, stage: str) -> None: # ----- Transcript artifact ------------------------------------------ if text: - write_transcript(out_dir, text, source=source, tokens=tokens) + write_transcript(out_dir, text, source=source, tokens=tokens, timestamps=timestamps) print(f" wrote {out_dir / 'transcript.json'}") for h in enc_hooks + dec_hooks: diff --git a/scripts/lib/ref_dump.py b/scripts/lib/ref_dump.py index c484ea76..2d425efa 100644 --- a/scripts/lib/ref_dump.py +++ b/scripts/lib/ref_dump.py @@ -102,6 +102,7 @@ def write_transcript( *, source: dict[str, Any], tokens: list[int] | None = None, + timestamps: dict[str, Any] | None = None, ) -> None: """Write transcript.json alongside a decode dump. @@ -115,4 +116,6 @@ def write_transcript( payload: dict[str, Any] = {"text": text, "source": source} if tokens is not None: payload["tokens"] = list(tokens) + if timestamps is not None: + payload["timestamps"] = timestamps (out_dir / "transcript.json").write_text(json.dumps(payload, indent=2) + "\n") diff --git a/scripts/validate.py b/scripts/validate.py index 14f21304..bc1e5ece 100644 --- a/scripts/validate.py +++ b/scripts/validate.py @@ -40,6 +40,7 @@ import datetime as dt import glob import json +import math import os import re import shutil @@ -309,6 +310,175 @@ def parse_cli_transcript(output: str) -> str | None: return None +def parse_cli_timestamps(output: str) -> dict[str, list[dict[str, Any]]] | None: + header_pattern = re.compile(r"^(segments|words|tokens):\s*(\d+)\s*$") + row_pattern = re.compile(r"^\s*\[\s*(\S+)\s*->\s*(\S+)\s*\]\s*(.*)$") + sections = {"segments": "segment", "words": "word", "tokens": "token"} + timestamps: dict[str, list[dict[str, Any]]] = {} + declared_counts: dict[str, int] = {} + section: str | None = None + + def finish_section() -> None: + nonlocal section + if section is not None and len(timestamps[section]) != declared_counts[section]: + raise ValueError( + f"CLI declared {declared_counts[section]} {section} timestamp rows, " + f"parsed {len(timestamps[section])}" + ) + section = None + + for line in output.splitlines(): + header = header_pattern.match(line) + if header is not None: + finish_section() + section = sections[header.group(1)] + if section in declared_counts: + raise ValueError(f"CLI emitted duplicate {section} timestamp section") + declared_counts[section] = int(header.group(2)) + timestamps[section] = [] + continue + if section is None: + continue + match = row_pattern.match(line) + if match is None: + finish_section() + continue + label = match.group(3) + if section == "token": + token_match = re.match(r"^p=\S+\s+(.*)$", label) + if token_match is None: + raise ValueError("malformed CLI token timestamp row") + label = token_match.group(1) + try: + start = float(match.group(1)) + end = float(match.group(2)) + except ValueError as exc: + raise ValueError(f"malformed CLI {section} timestamp edge") from exc + timestamps[section].append({ + section: label, + "start": start, + "end": end, + }) + finish_section() + return timestamps or None + + +def manifest_timestamp_config(manifest: dict[str, Any]) -> tuple[list[str], int]: + config = manifest.get("timestamp_validation") + if config is None: + return [], 80 + if not isinstance(config, dict): + raise SystemExit("error: manifest timestamp_validation must be an object") + + raw_levels = config.get("levels", []) + if not isinstance(raw_levels, list): + raise SystemExit("error: manifest timestamp_validation.levels must be a list") + + levels: list[str] = [] + for raw_level in raw_levels: + level = str(raw_level) + if level == "none": + continue + if level not in {"segment", "word", "token"}: + raise SystemExit( + f"error: unsupported manifest timestamp level {raw_level!r}" + ) + if level not in levels: + levels.append(level) + + tolerance_raw = config.get("tolerance_ms", 80) + if ( + isinstance(tolerance_raw, bool) + or not isinstance(tolerance_raw, (int, float)) + or not math.isfinite(float(tolerance_raw)) + or float(tolerance_raw) < 0.0 + or not float(tolerance_raw).is_integer() + ): + raise SystemExit("error: manifest timestamp_validation.tolerance_ms must be a nonnegative integer") + return levels, int(tolerance_raw) + + +def timestamp_artifact_error( + source: str, + timestamps: Any, + required_levels: list[str], +) -> str | None: + if not isinstance(timestamps, dict): + return f"{source} timestamps are absent or malformed" + + for level in required_levels: + rows = timestamps.get(level) + if not isinstance(rows, list) or not rows: + return f"{source} {level} timestamps are absent or empty" + previous_start = -1.0 + previous_end = -1.0 + for index, row in enumerate(rows): + if not isinstance(row, dict): + return f"{source} {level} timestamp row {index} is malformed" + label = row.get(level) + if not isinstance(label, str) or not label.strip(): + return f"{source} {level} timestamp row {index} has no label" + edges: list[float] = [] + for edge in ("start", "end"): + value = row.get(edge) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return f"{source} {level} timestamp row {index} has malformed {edge}" + value = float(value) + if not math.isfinite(value) or value < 0.0: + return f"{source} {level} timestamp row {index} has invalid {edge}" + edges.append(value) + start, end = edges + if start >= end: + return f"{source} {level} timestamp row {index} has start >= end" + if start < previous_start or end < previous_end: + return f"{source} {level} timestamp row {index} has non-monotonic edges" + previous_start = start + previous_end = end + return None + + +def compare_timestamps( + reference: Any, + cpp: Any, + required_levels: list[str], + tolerance_ms: int, + transcript_compare: str, +) -> tuple[bool, str | None]: + for source, timestamps in (("reference", reference), ("C++", cpp)): + error = timestamp_artifact_error(source, timestamps, required_levels) + if error is not None: + return False, error + + for level in required_levels: + ref_rows = reference[level] + cpp_rows = cpp[level] + if len(ref_rows) != len(cpp_rows): + return False, ( + f"{level} timestamp row count differs: " + f"reference={len(ref_rows)}, C++={len(cpp_rows)}" + ) + for index, (ref_row, cpp_row) in enumerate(zip(ref_rows, cpp_rows)): + ref_label = str(ref_row[level]) + cpp_label = str(cpp_row[level]) + if transcript_compare == "normalized": + ref_label = normalize_text_for_compare(ref_label) + cpp_label = normalize_text_for_compare(cpp_label) + elif transcript_compare == "dediarized": + ref_label = dediarize_text(ref_label) + cpp_label = dediarize_text(cpp_label) + if ref_label != cpp_label: + return False, f"{level} timestamp row {index} label differs" + for edge in ("start", "end"): + ref_ms = int(round(float(ref_row[edge]) * 1000.0)) + cpp_ms = int(round(float(cpp_row[edge]) * 1000.0)) + if abs(ref_ms - cpp_ms) > tolerance_ms: + return False, ( + f"{level} timestamp row {index} {edge} differs by " + f"{abs(ref_ms - cpp_ms)} ms" + ) + return True, None + + def write_cpp_transcript( out_dir: Path, *, @@ -318,6 +488,7 @@ def write_cpp_transcript( gguf: Path, backend: str, text: str, + timestamps: dict[str, list[dict[str, Any]]] | None = None, ) -> None: path = out_dir / "transcript.json" payload = { @@ -333,6 +504,8 @@ def write_cpp_transcript( "backend": backend, }, } + if timestamps is not None: + payload["timestamps"] = timestamps path.write_text(json.dumps(payload, indent=2) + "\n") print(f" wrote {path}", file=sys.stderr) @@ -405,6 +578,15 @@ def cmd_ref(args: argparse.Namespace) -> int: # pass overwrites the encoder pass (same values). for stage in ["encoder", "decode"]: cmd = base_args + [stage] + common_args + if args.family == "canary" and stage == "decode": + multitask = manifest.get("multitask_default") or {} + cmd += [ + "--source-lang", str(multitask.get("source_lang", "en")), + "--target-lang", str(multitask.get("target_lang", "en")), + "--task", str(multitask.get("task", "asr")), + "--pnc", str(multitask.get("pnc", "yes")), + "--toggle-timestamps", str(multitask.get("toggle_timestamps", "no")), + ] run_cmd( cmd, repo, @@ -473,6 +655,10 @@ def cmd_cpp(args: argparse.Namespace) -> int: cmd += ["--language", language] if args.family == "whisper": cmd += ["--timestamps", "none"] + if args.family == "canary": + multitask = manifest.get("multitask_default") or {} + timestamp_toggle = str(multitask.get("toggle_timestamps", "no")) + cmd += ["--timestamps", "word" if timestamp_toggle == "yes" else "none"] if args.family in ("sensevoice", "parakeet"): # The reference dumper emits the raw token stream including # control / language tags (sensevoice: language / event / @@ -510,6 +696,13 @@ def cmd_cpp(args: argparse.Namespace) -> int: raise SystemExit( f"error: cpp dump [{args.family}/{case_name}] did not emit a transcript line" ) + try: + timestamps = parse_cli_timestamps(result.stdout or "") + except ValueError as exc: + raise SystemExit( + f"error: cpp dump [{args.family}/{case_name}] emitted malformed " + f"timestamps: {exc}" + ) from exc write_cpp_transcript( out_dir, family=args.family, @@ -518,6 +711,7 @@ def cmd_cpp(args: argparse.Namespace) -> int: gguf=gguf, backend=args.backend, text=transcript, + timestamps=timestamps, ) return 0 @@ -527,6 +721,7 @@ def cmd_compare(args: argparse.Namespace) -> int: repo = find_repo_root(Path(__file__).parent) manifest = load_manifest(repo, args.family, getattr(args, "variant", None)) variant = manifest["variant"] + required_timestamp_levels, timestamp_tolerance_ms = manifest_timestamp_config(manifest) # Prefer per-variant tolerances declared in the manifest; fall # back to the family default. This lets larger variants (e.g. @@ -561,10 +756,26 @@ def cmd_compare(args: argparse.Namespace) -> int: if not cpp_dir.exists(): print(f"SKIP {case_name}: no C++ dumps at {cpp_dir}", file=sys.stderr) all_passed = False + if required_timestamp_levels: + transcript_results.append({ + "case": case_name, + "match": False, + "reason": "missing C++ dump directory", + "timestamps_match": False, + "timestamps_reason": "missing C++ dump directory", + }) continue if not ref_dir.exists(): print(f"SKIP {case_name}: no reference dumps at {ref_dir}", file=sys.stderr) all_passed = False + if required_timestamp_levels: + transcript_results.append({ + "case": case_name, + "match": False, + "reason": "missing reference dump directory", + "timestamps_match": False, + "timestamps_reason": "missing reference dump directory", + }) continue cmd = [ @@ -614,6 +825,11 @@ def cmd_compare(args: argparse.Namespace) -> int: transcript_results.append({ "case": case_name, "match": False, "reason": "missing C++ transcript artifact", + "timestamps_match": False if required_timestamp_levels else None, + "timestamps_reason": ( + "missing C++ transcript artifact" + if required_timestamp_levels else None + ), }) continue @@ -645,6 +861,36 @@ def cmd_compare(args: argparse.Namespace) -> int: else: print(f"\n Transcript: ok ({transcript_compare}) {cpp_text!r}") + if required_timestamp_levels: + timestamp_match, timestamp_reason = compare_timestamps( + ref_data.get("timestamps"), + cpp_data.get("timestamps"), + required_timestamp_levels, + timestamp_tolerance_ms, + transcript_compare, + ) + transcript_results[-1]["timestamps_match"] = timestamp_match + if not timestamp_match: + transcript_results[-1]["timestamps_reason"] = timestamp_reason + print( + f"FAIL timestamps mismatch (tolerance {timestamp_tolerance_ms} ms): " + f"{timestamp_reason}" + ) + all_passed = False + else: + print(f" Timestamps: ok (tolerance {timestamp_tolerance_ms} ms)") + elif required_timestamp_levels: + reason = "missing reference transcript artifact" + print(f"FAIL timestamps: {reason}: {ref_transcript}", file=sys.stderr) + all_passed = False + transcript_results.append({ + "case": case_name, + "match": False, + "reason": reason, + "timestamps_match": False, + "timestamps_reason": reason, + }) + if report_mode: write_report_bundle( repo=repo, @@ -823,6 +1069,12 @@ def write_report_bundle( else: summary.append(f"- reference: `{tr.get('reference', '')!r}`") summary.append(f"- c++: `{tr.get('cpp', '')!r}`") + if tr.get("timestamps_match") is not None: + summary.append( + f"- Timestamps match: **{'yes' if tr['timestamps_match'] else 'no'}**" + ) + if tr.get("timestamps_reason"): + summary.append(f"- Timestamp reason: {tr['timestamps_reason']}") summary.append("") (bundle_dir / "summary.md").write_text("\n".join(summary)) diff --git a/src/arch/canary/canary.h b/src/arch/canary/canary.h index ee905c48..d40c059b 100644 --- a/src/arch/canary/canary.h +++ b/src/arch/canary/canary.h @@ -33,6 +33,23 @@ namespace transcribe::canary { void apply_family_invariants(transcribe_model & model); +// Viterbi-align a transcript against frame-major CTC emissions. The returned +// states index the blank-interleaved target sequence at each frame. +bool viterbi_ctc_alignment(const std::vector & emissions, + int n_frames, + int vocab_size, + int blank_id, + const std::vector & text_ids, + std::vector & alignment); + +bool canary_aligner_input_too_long(int mel_n_frames, + const CanaryHParams & hp, + transcribe_timestamp_kind requested_kind); + +std::vector canary_preserve_special_tags(const Tokenizer & tokenizer, + const std::vector & generated_ids, + const std::vector & aligned_words); + // --------------------------------------------------------------------------- // KV cache for the autoregressive decoder. Same shape as cohere. // --------------------------------------------------------------------------- diff --git a/src/arch/canary/capabilities.cpp b/src/arch/canary/capabilities.cpp index 338ceafe..e91d70cf 100644 --- a/src/arch/canary/capabilities.cpp +++ b/src/arch/canary/capabilities.cpp @@ -15,8 +15,8 @@ void apply_family_invariants(transcribe_model & model) { // directions are an optional GGUF contract in stt.translation.pairs. caps.supports_translate = true; - // Timestamps out of scope. Advertise NONE; the GGUF KV - // stt.capability.timestamps overrides this once the path is wired up. + // The loader raises this to WORD after it has distinguished Flash's AED + // timestamp-token path from v2's optional auxiliary CTC aligner. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; // Feature bits: cancellation is wired at the run level. Canary diff --git a/src/arch/canary/encoder.cpp b/src/arch/canary/encoder.cpp index 744c7194..dae089f8 100644 --- a/src/arch/canary/encoder.cpp +++ b/src/arch/canary/encoder.cpp @@ -224,4 +224,73 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } +AlignerBuild build_aligner_graph(ggml_context * ctx, + const CanaryWeights & w, + const CanaryHParams & hp, + int n_mel_frames, + ggml_type kv_type, + bool use_flash, + const char * backend_name) { + AlignerBuild ab{}; + if (ctx == nullptr || !hp.aligner_present || n_mel_frames <= 0) { + return ab; + } + if (hp.aligner_subsampling_factor != 8 || hp.fe_num_mels != 128 || w.aligner.blocks.empty()) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "canary aligner: unsupported geometry (subsampling=%d num_mels=%d blocks=%zu)", + hp.aligner_subsampling_factor, hp.fe_num_mels, w.aligner.blocks.size()); + return ab; + } + + conf::ConvPolicy policy{}; + policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); + policy.direct_dw_in_pre_encode = false; + + ab.mel_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_mel_frames, hp.fe_num_mels); + ggml_set_name(ab.mel_in, "aligner.mel.in"); + ggml_set_input(ab.mel_in); + + ggml_tensor * x = conf::build_pre_encode(ctx, to_view(w.aligner.pre_encode), ab.mel_in, policy, + /*name_prefix=*/"aligner.enc.pre_encode", + /*error_tag=*/"canary aligner"); + if (x == nullptr) { + return ab; + } + + const int64_t T_enc = x->ne[1]; + const int64_t pos_len = 2 * T_enc - 1; + ab.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.aligner_d_model, pos_len); + ggml_set_name(ab.pos_emb_in, "aligner.pos_emb.in"); + ggml_set_input(ab.pos_emb_in); + + conf::BlockParams bparams{}; + bparams.d_model = hp.aligner_d_model; + bparams.n_head = hp.aligner_n_heads; + bparams.conv_kernel = hp.aligner_conv_kernel; + bparams.kv_type = kv_type; + bparams.use_flash = use_flash; + bparams.policy = policy; + bparams.conv_norm_type = conf::BlockParams::ConvNormType::BatchNorm; + + for (const CanaryBlock & block : w.aligner.blocks) { + x = conf::build_conformer_block(ctx, x, ab.pos_emb_in, to_view(block), bparams); + if (x == nullptr) { + return ab; + } + } + + ab.logits = ggml_mul_mat(ctx, w.aligner.head.weight, x); + ab.logits = ggml_add(ctx, ab.logits, w.aligner.head.bias); + ggml_set_name(ab.logits, "aligner.ctc.logits"); + ggml_set_output(ab.logits); + + ab.graph = ggml_new_graph_custom(ctx, 16384, false); + if (ab.graph == nullptr) { + return ab; + } + ggml_build_forward_expand(ab.graph, ab.logits); + return ab; +} + } // namespace transcribe::canary diff --git a/src/arch/canary/encoder.h b/src/arch/canary/encoder.h index c19e8cd9..f74c752c 100644 --- a/src/arch/canary/encoder.h +++ b/src/arch/canary/encoder.h @@ -35,6 +35,13 @@ struct EncoderBuild { ggml_cgraph * graph = nullptr; }; +struct AlignerBuild { + ggml_tensor * mel_in = nullptr; + ggml_tensor * pos_emb_in = nullptr; + ggml_tensor * logits = nullptr; // [vocab_size, T_enc] + ggml_cgraph * graph = nullptr; +}; + EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const CanaryWeights & weights, const CanaryHParams & hp, @@ -43,4 +50,12 @@ EncoderBuild build_encoder_graph(ggml_context * compute_ctx, bool use_flash = true, const char * backend_name = ""); +AlignerBuild build_aligner_graph(ggml_context * compute_ctx, + const CanaryWeights & weights, + const CanaryHParams & hp, + int n_mel_frames, + ggml_type kv_type = GGML_TYPE_COUNT, + bool use_flash = true, + const char * backend_name = ""); + } // namespace transcribe::canary diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index 15809fda..2d653d1f 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -210,6 +211,73 @@ CanaryModel::~CanaryModel() { plan.primary_kind = transcribe::BackendKind::Unknown; } +bool viterbi_ctc_alignment(const std::vector & emissions, + int n_frames, + int vocab_size, + int blank_id, + const std::vector & text_ids, + std::vector & alignment) { + alignment.clear(); + if (n_frames <= 0 || vocab_size <= 0 || blank_id < 0 || blank_id >= vocab_size || text_ids.empty() || + emissions.size() != static_cast(n_frames) * vocab_size) { + return false; + } + for (int id : text_ids) { + if (id < 0 || id >= vocab_size || id == blank_id) { + return false; + } + } + + const int n_states = static_cast(text_ids.size()) * 2 + 1; + if (n_frames < static_cast(text_ids.size())) { + return false; + } + + std::vector labels(n_states, blank_id); + for (int i = 0; i < static_cast(text_ids.size()); ++i) { + labels[2 * i + 1] = text_ids[i]; + } + + const float neg_inf = -std::numeric_limits::infinity(); + std::vector previous(n_states, neg_inf), current(n_states, neg_inf); + std::vector back(static_cast(n_frames) * n_states, 0); + previous[0] = emissions[blank_id]; + previous[1] = emissions[text_ids[0]]; + for (int frame = 1; frame < n_frames; ++frame) { + const float * row = emissions.data() + static_cast(frame) * vocab_size; + std::fill(current.begin(), current.end(), neg_inf); + for (int state = 0; state < n_states; ++state) { + float best = previous[state]; + uint8_t step = 0; + if (state >= 1 && previous[state - 1] > best) { + best = previous[state - 1]; + step = 1; + } + if (state >= 2 && labels[state] != labels[state - 2] && previous[state - 2] > best) { + best = previous[state - 2]; + step = 2; + } + current[state] = best + row[labels[state]]; + back[static_cast(frame) * n_states + state] = step; + } + previous.swap(current); + } + + const float final_score = std::max(previous[n_states - 2], previous[n_states - 1]); + if (!std::isfinite(final_score)) { + return false; + } + + int state = previous[n_states - 2] >= previous[n_states - 1] ? n_states - 2 : n_states - 1; + alignment.resize(n_frames); + alignment[n_frames - 1] = state; + for (int frame = n_frames - 1; frame > 0; --frame) { + state -= back[static_cast(frame) * n_states + state]; + alignment[frame - 1] = state; + } + return state == 0 || state == 1; +} + namespace { constexpr float kBnEps = 1e-5f; @@ -277,14 +345,19 @@ int canary_context_ceiling(int32_t n_ctx_knob, const CanaryHParams & hp) { // Fuse inference-time BatchNorm into precomputed scale + bias. Same as // parakeet/cohere — see those files for the math. transcribe_status fuse_batch_norm(CanaryModel & m) { - const size_t n_blocks = m.weights.blocks.size(); + std::vector blocks; + blocks.reserve(m.weights.blocks.size() + m.weights.aligner.blocks.size()); + for (CanaryBlock & block : m.weights.blocks) { + blocks.push_back(&block); + } + for (CanaryBlock & block : m.weights.aligner.blocks) { + blocks.push_back(&block); + } + const size_t n_blocks = blocks.size(); if (n_blocks == 0) { return TRANSCRIBE_OK; } - const int64_t d = m.hparams.enc_d_model; - const size_t tensor_bytes = static_cast(d) * sizeof(float); - const size_t ctx_size = n_blocks * 2 * ggml_tensor_overhead() + 256; ggml_init_params params = { ctx_size, nullptr, true }; m.bn_fused_ctx = ggml_init(params); @@ -293,7 +366,8 @@ transcribe_status fuse_batch_norm(CanaryModel & m) { } for (size_t i = 0; i < n_blocks; ++i) { - auto & b = m.weights.blocks[i]; + CanaryBlock & b = *blocks[i]; + const int64_t d = b.conv_bn_w->ne[0]; b.conv_bn_fused_scale = ggml_new_tensor_1d(m.bn_fused_ctx, GGML_TYPE_F32, d); b.conv_bn_fused_bias = ggml_new_tensor_1d(m.bn_fused_ctx, GGML_TYPE_F32, d); } @@ -303,11 +377,12 @@ transcribe_status fuse_batch_norm(CanaryModel & m) { return TRANSCRIBE_ERR_BACKEND; } - std::vector bn_w(d), bn_b(d), rm(d), rv(d); - std::vector fused_s(d), fused_b(d); - for (size_t i = 0; i < n_blocks; ++i) { - auto & b = m.weights.blocks[i]; + CanaryBlock & b = *blocks[i]; + const int64_t d = b.conv_bn_w->ne[0]; + const size_t tensor_bytes = static_cast(d) * sizeof(float); + std::vector bn_w(d), bn_b(d), rm(d), rv(d); + std::vector fused_s(d), fused_b(d); ggml_backend_tensor_get(b.conv_bn_w, bn_w.data(), 0, tensor_bytes); ggml_backend_tensor_get(b.conv_bn_b, bn_b.data(), 0, tensor_bytes); ggml_backend_tensor_get(b.conv_bn_rm, rm.data(), 0, tensor_bytes); @@ -339,6 +414,14 @@ transcribe_status promote_conv_pw_to_f32_on_cpu(CanaryModel & m) { slots.push_back({ &b.conv_pw2_w, b.conv_pw2_w }); } } + for (auto & b : m.weights.aligner.blocks) { + if (b.conv_pw1_w != nullptr && b.conv_pw1_w->type == GGML_TYPE_F16) { + slots.push_back({ &b.conv_pw1_w, b.conv_pw1_w }); + } + if (b.conv_pw2_w != nullptr && b.conv_pw2_w->type == GGML_TYPE_F16) { + slots.push_back({ &b.conv_pw2_w, b.conv_pw2_w }); + } + } return load_common::promote_conv_pw_f16_to_f32_on_cpu(m.plan, slots, "canary", &m.conv_pw_f32_ctx, &m.conv_pw_f32_buffer); } @@ -383,6 +466,13 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } + if (m->hparams.aligner_present && !m->tok.has_bpe_canonicalizer()) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "canary: timestamp aligner requires a SentencePiece BPE tokenizer with per-token scores"); + return TRANSCRIBE_ERR_GGUF; + } + m->caps.max_timestamp_kind = m->hparams.aligner_present ? TRANSCRIBE_TIMESTAMPS_WORD : TRANSCRIBE_TIMESTAMPS_NONE; + // Publish the input-length ceiling now that the encoder positional span // and frontend rate are known (apply_family_invariants ran before the // hparams were read). The encoder is the binding INPUT limit for canary. @@ -679,6 +769,284 @@ bool translation_pair_allowed(const CanaryHParams & hp, const char * src, const return false; } +struct CanaryTimedWord { + std::string text; + int64_t t0_ms = 0; + int64_t t1_ms = 0; + int first_token = 0; + int n_tokens = 0; +}; + +struct CanaryWordRef { + std::string text; + int first_token = 0; + int n_tokens = 0; +}; + +std::string trim_ascii_space(std::string text) { + size_t begin = 0; + while (begin < text.size() && (text[begin] == ' ' || text[begin] == '\t' || text[begin] == '\n')) { + ++begin; + } + size_t end = text.size(); + while (end > begin && (text[end - 1] == ' ' || text[end - 1] == '\t' || text[end - 1] == '\n')) { + --end; + } + return text.substr(begin, end - begin); +} + +std::vector canary_text_ids(const CanaryModel & cm, const std::vector & generated_ids) { + std::vector text_ids; + text_ids.reserve(generated_ids.size()); + for (int id : generated_ids) { + if (!cm.tok.is_control(id)) { + text_ids.push_back(id); + } + } + return text_ids; +} + +std::vector canary_word_refs(const CanaryModel & cm, const std::vector & text_ids) { + static constexpr char kSpSpace[] = "\xE2\x96\x81"; + std::vector refs; + int begin = 0; + for (int i = 0; i < static_cast(text_ids.size()); ++i) { + const std::string & piece = cm.tok.token(text_ids[i]); + const bool starts_word = piece.size() >= 3 && std::memcmp(piece.data(), kSpSpace, 3) == 0; + if (starts_word && i > begin) { + std::string text = trim_ascii_space(cm.tok.decode(text_ids.data() + begin, i - begin)); + if (!text.empty()) { + refs.push_back({ std::move(text), begin, i - begin }); + } + begin = i; + } + } + if (begin < static_cast(text_ids.size())) { + std::string text = + trim_ascii_space(cm.tok.decode(text_ids.data() + begin, static_cast(text_ids.size()) - begin)); + if (!text.empty()) { + refs.push_back({ std::move(text), begin, static_cast(text_ids.size()) - begin }); + } + } + return refs; +} + +void assemble_timed_result(CanarySession & cc, + const std::vector & timed_words, + const std::vector & display_words, + transcribe_timestamp_kind requested_kind) { + if (timed_words.empty()) { + cc.result_kind = TRANSCRIBE_TIMESTAMPS_NONE; + return; + } + + cc.segments.clear(); + cc.words.clear(); + + const bool emit_words = requested_kind != TRANSCRIBE_TIMESTAMPS_SEGMENT; + int seg_begin = 0; + for (int i = 0; i < static_cast(timed_words.size()); ++i) { + const std::string & word_text = timed_words[i].text; + const bool boundary = + !word_text.empty() && (word_text.back() == '.' || word_text.back() == '?' || word_text.back() == '!'); + if (!boundary && i + 1 < static_cast(timed_words.size())) { + continue; + } + + transcribe_session::SegmentEntry seg{}; + seg.t0_ms = timed_words[seg_begin].t0_ms; + seg.t1_ms = timed_words[i].t1_ms; + seg.first_word = emit_words ? static_cast(cc.words.size()) : 0; + seg.n_words = emit_words ? i - seg_begin + 1 : 0; + for (int j = seg_begin; j <= i; ++j) { + const std::string & display_text = + display_words.size() == timed_words.size() ? display_words[j] : timed_words[j].text; + if (!seg.text.empty()) { + seg.text.push_back(' '); + } + seg.text += display_text; + if (emit_words) { + transcribe_session::WordEntry word{}; + word.text = display_text; + word.t0_ms = timed_words[j].t0_ms; + word.t1_ms = timed_words[j].t1_ms; + word.seg_index = static_cast(cc.segments.size()); + word.first_token = 0; + word.n_tokens = 0; + cc.words.push_back(std::move(word)); + } + } + cc.segments.push_back(std::move(seg)); + seg_begin = i + 1; + } + cc.result_kind = emit_words ? TRANSCRIBE_TIMESTAMPS_WORD : TRANSCRIBE_TIMESTAMPS_SEGMENT; +} + +transcribe_status run_ctc_aligner(CanarySession & cc, + CanaryModel & cm, + const std::vector & text_ids, + std::vector & timed_words) { + timed_words.clear(); + const std::vector refs = canary_word_refs(cm, text_ids); + if (text_ids.empty() || refs.empty()) { + return TRANSCRIBE_OK; + } + + if (cc.compute_ctx != nullptr) { + ggml_free(cc.compute_ctx); + cc.compute_ctx = nullptr; + } + ggml_init_params ip{}; + ip.mem_size = 8 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + cc.compute_ctx = ggml_init(ip); + if (cc.compute_ctx == nullptr) { + return TRANSCRIBE_ERR_OOM; + } + + ggml_type kv_type = GGML_TYPE_COUNT; + if (cc.kv_type == TRANSCRIBE_KV_TYPE_F32) { + kv_type = GGML_TYPE_F32; + } else if (cc.kv_type == TRANSCRIBE_KV_TYPE_F16) { + kv_type = GGML_TYPE_F16; + } + AlignerBuild ab = build_aligner_graph(cc.compute_ctx, cm.weights, cm.hparams, + static_cast(cc.mel_buf.size() / cm.hparams.fe_num_mels), kv_type, + cc.encoder_use_flash, cm.backend.c_str()); + if (ab.graph == nullptr || ab.logits == nullptr || ab.mel_in == nullptr || ab.pos_emb_in == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + + ggml_backend_sched_reset(cc.sched); + if (!ggml_backend_sched_alloc_graph(cc.sched, ab.graph)) { + return TRANSCRIBE_ERR_OOM; + } + ggml_backend_tensor_set(ab.mel_in, cc.mel_buf.data(), 0, cc.mel_buf.size() * sizeof(float)); + + const int d_model = cm.hparams.aligner_d_model; + const int pos_len = static_cast(ab.pos_emb_in->ne[1]); + const int T_enc = (pos_len + 1) / 2; + std::vector pos(static_cast(pos_len) * d_model, 0.0f); + const float ln_10000 = std::log(10000.0f); + for (int p = 0; p < pos_len; ++p) { + const float rel = static_cast((T_enc - 1) - p); + float * row = pos.data() + static_cast(p) * d_model; + for (int k = 0; k < d_model / 2; ++k) { + const float div = std::exp(static_cast(2 * k) * (-ln_10000 / d_model)); + row[2 * k] = std::sin(rel * div); + row[2 * k + 1] = std::cos(rel * div); + } + } + ggml_backend_tensor_set(ab.pos_emb_in, pos.data(), 0, pos.size() * sizeof(float)); + transcribe::configure_sched_n_threads(cc.sched, cc.n_threads); + if (ggml_backend_sched_graph_compute(cc.sched, ab.graph) != GGML_STATUS_SUCCESS) { + return TRANSCRIBE_ERR_GGUF; + } + + const int vocab = cm.hparams.aligner_vocab_size; + const int blank = cm.hparams.aligner_blank_id; + if (ab.logits->ne[0] != vocab || ab.logits->ne[1] != T_enc) { + return TRANSCRIBE_ERR_GGUF; + } + std::vector logits(static_cast(vocab) * T_enc); + ggml_backend_tensor_get(ab.logits, logits.data(), 0, logits.size() * sizeof(float)); + + std::vector alignment; + if (!viterbi_ctc_alignment(logits, T_enc, vocab, blank, text_ids, alignment)) { + return TRANSCRIBE_OK; + } + + std::vector token_first(text_ids.size(), -1); + std::vector token_last(text_ids.size(), -1); + for (int t = 0; t < T_enc; ++t) { + const int state = alignment[t]; + if ((state & 1) == 0) { + continue; + } + const int token = state / 2; + if (token_first[token] < 0) { + token_first[token] = t; + } + token_last[token] = t; + } + + const int64_t frame_ms = static_cast(cm.hparams.fe_hop_length) * cm.hparams.aligner_subsampling_factor * + 1000 / cm.hparams.fe_sample_rate; + for (const CanaryWordRef & ref : refs) { + const int first = ref.first_token; + const int last = ref.first_token + ref.n_tokens - 1; + if (first < 0 || last >= static_cast(token_first.size()) || token_first[first] < 0 || + token_last[last] < 0) { + continue; + } + int end_frame = token_last[last] + 1; + if (last > first) { + const std::string punctuation = trim_ascii_space(cm.tok.decode(&text_ids[last], 1)); + if ((punctuation == "," || punctuation == "." || punctuation == "!" || punctuation == "?") && + token_last[last - 1] >= 0) { + end_frame = token_last[last - 1] + 1; + } + } + timed_words.push_back( + { ref.text, token_first[first] * frame_ms, end_frame * frame_ms, ref.first_token, ref.n_tokens }); + } + return TRANSCRIBE_OK; +} + +transcribe_timestamp_kind canary_timestamp_kind_for_run(const transcribe_run_params * params) { + const transcribe_timestamp_kind requested = params == nullptr ? TRANSCRIBE_TIMESTAMPS_AUTO : params->timestamps; + if (params != nullptr && params->task == TRANSCRIBE_TASK_TRANSLATE && requested == TRANSCRIBE_TIMESTAMPS_AUTO) { + return TRANSCRIBE_TIMESTAMPS_NONE; + } + return requested; +} + +transcribe_status run_validate(const transcribe_session *, const transcribe_run_params * params) { + if (params != nullptr && params->task == TRANSCRIBE_TASK_TRANSLATE && + params->timestamps != TRANSCRIBE_TIMESTAMPS_NONE && params->timestamps != TRANSCRIBE_TIMESTAMPS_AUTO) { + return TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS; + } + return TRANSCRIBE_OK; +} + +transcribe_status apply_canary_timestamps(CanarySession & cc, + CanaryModel & cm, + const std::vector & generated_ids, + const transcribe_run_params * params, + transcribe_timestamp_kind requested) { + if (requested == TRANSCRIBE_TIMESTAMPS_NONE || cm.caps.max_timestamp_kind == TRANSCRIBE_TIMESTAMPS_NONE) { + return TRANSCRIBE_OK; + } + + const int64_t t_align_start = ggml_time_us(); + std::vector timed_words; + const std::vector decoded_ids = canary_text_ids(cm, generated_ids); + std::vector canonical_ids_i32; + const transcribe_status canonical_st = + cm.tok.canonicalize_bpe(decoded_ids.data(), static_cast(decoded_ids.size()), canonical_ids_i32); + if (canonical_st != TRANSCRIBE_OK) { + cc.t_decode_us += ggml_time_us() - t_align_start; + return canonical_st; + } + const std::vector canonical_ids(canonical_ids_i32.begin(), canonical_ids_i32.end()); + const transcribe_status st = run_ctc_aligner(cc, cm, canonical_ids, timed_words); + cc.t_decode_us += ggml_time_us() - t_align_start; + if (st != TRANSCRIBE_OK) { + return st; + } + std::vector display_words; + display_words.reserve(timed_words.size()); + for (const CanaryTimedWord & word : timed_words) { + display_words.push_back(word.text); + } + if (params != nullptr && params->keep_special_tags) { + display_words = canary_preserve_special_tags(cm.tok, generated_ids, display_words); + } + assemble_timed_result(cc, timed_words, display_words, requested); + return TRANSCRIBE_OK; +} + transcribe_status run(transcribe_session * session, const float * pcm, int n_samples, @@ -716,6 +1084,16 @@ transcribe_status run(transcribe_session * session, } cc->t_mel_us = ggml_time_us() - t_mel_start; + const transcribe_timestamp_kind requested_kind = canary_timestamp_kind_for_run(params); + if (canary_aligner_input_too_long(mel_n_frames, cm->hparams, requested_kind)) { + const int t_enc_pred = canary_predict_t_enc(mel_n_frames, cm->hparams.aligner_subsampling_factor); + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "canary run: input too long for requested timestamps - %d aligner frames exceed " + "the supported %d", + t_enc_pred, cm->hparams.aligner_pos_emb_max_len); + return TRANSCRIBE_ERR_INPUT_TOO_LONG; + } + // Input-length gate: T_enc must stay within enc_pos_emb_max_len or the // runtime pos table aliases past the trained range. T_enc is a // deterministic function of mel_n_frames, so reject an over-length clip @@ -1014,6 +1392,8 @@ transcribe_status run(transcribe_session * session, cc->kv_cache.cross_populated = true; } + std::vector generated_ids; + // Prompt pass + autoregressive decode. { if (!new_compute_ctx(4 * 1024 * 1024)) { @@ -1115,7 +1495,6 @@ transcribe_status run(transcribe_session * session, } } - std::vector generated_ids; if (next_token != eos_id) { generated_ids.push_back(next_token); } @@ -1351,6 +1730,13 @@ transcribe_status run(transcribe_session * session, commit_result(); } + if (cc->has_result) { + if (const transcribe_status st = apply_canary_timestamps(*cc, *cm, generated_ids, params, requested_kind); + st != TRANSCRIBE_OK) { + return st; + } + } + // Partial transcript committed above; a truncated decode returns the hard // OUTPUT_TRUNCATED status (the result stays readable, like an aborted run). return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; @@ -1468,13 +1854,17 @@ transcribe_status run_batch_serial(CanarySession * cc, const int * n_samples, int n, const transcribe_run_params * params) { + bool any_truncated = false; for (int i = 0; i < n; ++i) { if (cc->poll_abort()) { + cc->was_truncated = any_truncated; return TRANSCRIBE_ERR_ABORTED; } + cc->was_truncated = false; 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) { + any_truncated |= cc->was_truncated; + if (st == TRANSCRIBE_OK || (st == TRANSCRIBE_ERR_OUTPUT_TRUNCATED && cc->has_result)) { cc->batch_results.push_back(cc->capture_result(st)); } else { transcribe_session::ResultSet rs; @@ -1482,6 +1872,7 @@ transcribe_status run_batch_serial(CanarySession * cc, cc->batch_results.push_back(std::move(rs)); } } + cc->was_truncated = any_truncated; return TRANSCRIBE_OK; } @@ -1505,6 +1896,9 @@ transcribe_status run_batch(transcribe_session * session, if (n == 1 || !cc->decoder_use_flash || !primary_is_gpu || transcribe::debug::enabled()) { return run_batch_serial(cc, pcm, n_samples, n, params); } + if (cm->hparams.aligner_present && canary_timestamp_kind_for_run(params) != TRANSCRIBE_TIMESTAMPS_NONE) { + return run_batch_serial(cc, pcm, n_samples, n, params); + } transcribe::debug::init(); const auto & hp = cm->hparams; @@ -1815,6 +2209,50 @@ transcribe_status run_batch(transcribe_session * session, } // namespace +bool canary_aligner_input_too_long(int mel_n_frames, + const CanaryHParams & hp, + transcribe_timestamp_kind requested_kind) { + if (!hp.aligner_present || requested_kind == TRANSCRIBE_TIMESTAMPS_NONE || hp.aligner_pos_emb_max_len <= 0) { + return false; + } + return canary_predict_t_enc(mel_n_frames, hp.aligner_subsampling_factor) > hp.aligner_pos_emb_max_len; +} + +std::vector canary_preserve_special_tags(const Tokenizer & tokenizer, + const std::vector & generated_ids, + const std::vector & aligned_words) { + static constexpr char kSpSpace[] = "\xE2\x96\x81"; + + std::vector words = aligned_words; + if (words.empty()) { + return words; + } + + int current_word = -1; + std::string leading_tags; + for (int id : generated_ids) { + if (tokenizer.is_control(id)) { + const std::string tag = tokenizer.decode(&id, 1); + if (current_word < 0) { + leading_tags += tag; + } else { + words[std::min(current_word, static_cast(words.size()) - 1)] += tag; + } + continue; + } + + const std::string & piece = tokenizer.token(id); + const bool starts_word = piece.size() >= 3 && std::memcmp(piece.data(), kSpSpace, 3) == 0; + if (current_word < 0 || starts_word) { + ++current_word; + } + } + if (!leading_tags.empty()) { + words[0] = leading_tags + " " + words[0]; + } + return words; +} + extern const Arch arch = { /* .name = */ "canary", /* .load = */ load, @@ -1827,6 +2265,7 @@ extern const Arch arch = { /* .stream_finalize = */ nullptr, /* .stream_reset = */ nullptr, /* .accepts_ext_kind = */ nullptr, + /* .run_validate = */ run_validate, }; } // namespace transcribe::canary diff --git a/src/arch/canary/weights.cpp b/src/arch/canary/weights.cpp index 11d1d2a9..12efc664 100644 --- a/src/arch/canary/weights.cpp +++ b/src/arch/canary/weights.cpp @@ -309,6 +309,63 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, CanaryHParams & return st; } + if (auto st = read_optional_bool_kv(gguf, "stt.canary.aligner.present", kFamilyTag, false, hp.aligner_present); + st != TRANSCRIBE_OK) { + return st; + } + if (hp.aligner_present) { + if (auto st = + read_required_u32_kv(gguf, "stt.canary.aligner.encoder.n_layers", kFamilyTag, hp.aligner_n_layers); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.d_model", kFamilyTag, hp.aligner_d_model); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.n_heads", kFamilyTag, hp.aligner_n_heads); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.d_ff", kFamilyTag, hp.aligner_d_ff); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.conv_kernel", kFamilyTag, + hp.aligner_conv_kernel); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.subsampling_factor", kFamilyTag, + hp.aligner_subsampling_factor); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.subsampling_channels", kFamilyTag, + hp.aligner_subsampling_channels); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.encoder.pos_emb_max_len", kFamilyTag, + hp.aligner_pos_emb_max_len); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_string_kv(gguf, "stt.canary.aligner.encoder.conv_norm_type", kFamilyTag, + hp.aligner_conv_norm_type); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.vocab_size", kFamilyTag, hp.aligner_vocab_size); + st != TRANSCRIBE_OK) { + return st; + } + if (auto st = read_required_u32_kv(gguf, "stt.canary.aligner.blank_id", kFamilyTag, hp.aligner_blank_id); + st != TRANSCRIBE_OK) { + return st; + } + } + // Cross-field invariants. if (hp.enc_n_layers <= 0 || hp.enc_d_model <= 0 || hp.enc_n_heads <= 0 || hp.enc_d_ff <= 0 || hp.enc_conv_kernel <= 0 || hp.enc_subsampling_factor <= 0 || hp.enc_subsampling_channels <= 0) { @@ -366,6 +423,31 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, CanaryHParams & hp.prompt_format.c_str()); return TRANSCRIBE_ERR_GGUF; } + if (hp.aligner_present) { + if (hp.aligner_n_layers <= 0 || hp.aligner_d_model <= 0 || hp.aligner_n_heads <= 0 || hp.aligner_d_ff <= 0 || + hp.aligner_conv_kernel <= 0 || hp.aligner_subsampling_factor <= 0 || hp.aligner_subsampling_channels <= 0 || + hp.aligner_pos_emb_max_len <= 0 || hp.aligner_vocab_size <= 0 || hp.aligner_blank_id < 0 || + hp.aligner_blank_id >= hp.aligner_vocab_size) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary: invalid auxiliary aligner hparams"); + return TRANSCRIBE_ERR_GGUF; + } + if (hp.aligner_d_model % hp.aligner_n_heads != 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary: aligner d_model (%d) not divisible by n_heads (%d)", + hp.aligner_d_model, hp.aligner_n_heads); + return TRANSCRIBE_ERR_GGUF; + } + if (hp.aligner_subsampling_factor != 8 || hp.aligner_conv_norm_type != "batch_norm") { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary: unsupported aligner geometry (subsampling=%d conv_norm=%s)", + hp.aligner_subsampling_factor, hp.aligner_conv_norm_type.c_str()); + return TRANSCRIBE_ERR_GGUF; + } + if (hp.aligner_vocab_size != hp.dec_vocab_size + 1 || hp.aligner_blank_id != hp.dec_vocab_size) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "canary: aligner vocab/blank (%d/%d) must match AED vocab+blank (%d/%d)", hp.aligner_vocab_size, + hp.aligner_blank_id, hp.dec_vocab_size + 1, hp.dec_vocab_size); + return TRANSCRIBE_ERR_GGUF; + } + } return TRANSCRIBE_OK; } @@ -560,6 +642,72 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, const CanaryHPar GET_LIN(weights.head.weight, "dec.head.weight", d_dec, vocab); GET_F32(weights.head.bias, "dec.head.bias", vocab); + if (hp.aligner_present) { + const int64_t ad = hp.aligner_d_model; + const int64_t aff = hp.aligner_d_ff; + const int64_t aheads = hp.aligner_n_heads; + const int64_t ahead = hp.aligner_head_dim(); + const int64_t ak = hp.aligner_conv_kernel; + const int64_t ach = hp.aligner_subsampling_channels; + const int64_t apre_in = ach * (hp.fe_num_mels / hp.aligner_subsampling_factor); + const int64_t avocab = hp.aligner_vocab_size; + auto & aligner = weights.aligner; + + GET_CONV(aligner.pre_encode.conv0_w, "aligner.enc.pre_encode.conv.0.weight", 3, 3, 1, ach); + GET_F32(aligner.pre_encode.conv0_b, "aligner.enc.pre_encode.conv.0.bias", ach); + GET_CONV(aligner.pre_encode.conv2_w, "aligner.enc.pre_encode.conv.2.weight", 3, 3, 1, ach); + GET_F32(aligner.pre_encode.conv2_b, "aligner.enc.pre_encode.conv.2.bias", ach); + GET_CONV(aligner.pre_encode.conv3_w, "aligner.enc.pre_encode.conv.3.weight", 1, 1, ach, ach); + GET_F32(aligner.pre_encode.conv3_b, "aligner.enc.pre_encode.conv.3.bias", ach); + GET_CONV(aligner.pre_encode.conv5_w, "aligner.enc.pre_encode.conv.5.weight", 3, 3, 1, ach); + GET_F32(aligner.pre_encode.conv5_b, "aligner.enc.pre_encode.conv.5.bias", ach); + GET_CONV(aligner.pre_encode.conv6_w, "aligner.enc.pre_encode.conv.6.weight", 1, 1, ach, ach); + GET_F32(aligner.pre_encode.conv6_b, "aligner.enc.pre_encode.conv.6.bias", ach); + GET_LIN(aligner.pre_encode.out_w, "aligner.enc.pre_encode.out.weight", apre_in, ad); + GET_F32(aligner.pre_encode.out_b, "aligner.enc.pre_encode.out.bias", ad); + + aligner.blocks.assign(hp.aligner_n_layers, CanaryBlock{}); + for (int i = 0; i < hp.aligner_n_layers; ++i) { + auto & b = aligner.blocks[i]; + + GET_F32(b.norm_ff1_w, lname("aligner.enc.blocks.%d.norm_ff1.weight", i), ad); + GET_F32(b.norm_ff1_b, lname("aligner.enc.blocks.%d.norm_ff1.bias", i), ad); + GET_LIN(b.ff1_lin1_w, lname("aligner.enc.blocks.%d.ff1.linear1.weight", i), ad, aff); + GET_LIN(b.ff1_lin2_w, lname("aligner.enc.blocks.%d.ff1.linear2.weight", i), aff, ad); + + GET_F32(b.norm_attn_w, lname("aligner.enc.blocks.%d.norm_attn.weight", i), ad); + GET_F32(b.norm_attn_b, lname("aligner.enc.blocks.%d.norm_attn.bias", i), ad); + GET_LIN(b.attn_q_w, lname("aligner.enc.blocks.%d.attn.linear_q.weight", i), ad, ad); + GET_LIN(b.attn_k_w, lname("aligner.enc.blocks.%d.attn.linear_k.weight", i), ad, ad); + GET_LIN(b.attn_v_w, lname("aligner.enc.blocks.%d.attn.linear_v.weight", i), ad, ad); + GET_LIN(b.attn_out_w, lname("aligner.enc.blocks.%d.attn.linear_out.weight", i), ad, ad); + GET_LIN(b.attn_pos_w, lname("aligner.enc.blocks.%d.attn.linear_pos.weight", i), ad, ad); + GET_F32(b.attn_pos_u, lname("aligner.enc.blocks.%d.attn.pos_bias_u", i), ahead, aheads); + GET_F32(b.attn_pos_v, lname("aligner.enc.blocks.%d.attn.pos_bias_v", i), ahead, aheads); + + GET_F32(b.norm_conv_w, lname("aligner.enc.blocks.%d.norm_conv.weight", i), ad); + GET_F32(b.norm_conv_b, lname("aligner.enc.blocks.%d.norm_conv.bias", i), ad); + GET_CONV(b.conv_pw1_w, lname("aligner.enc.blocks.%d.conv.pointwise1.weight", i), 1, ad, 2 * ad); + GET_CONV(b.conv_dw_w, lname("aligner.enc.blocks.%d.conv.depthwise.weight", i), ak, 1, ad); + GET_CONV(b.conv_pw2_w, lname("aligner.enc.blocks.%d.conv.pointwise2.weight", i), 1, ad, ad); + GET_F32(b.conv_bn_w, lname("aligner.enc.blocks.%d.conv.bn.weight", i), ad); + GET_F32(b.conv_bn_b, lname("aligner.enc.blocks.%d.conv.bn.bias", i), ad); + GET_F32(b.conv_bn_rm, lname("aligner.enc.blocks.%d.conv.bn.running_mean", i), ad); + GET_F32(b.conv_bn_rv, lname("aligner.enc.blocks.%d.conv.bn.running_var", i), ad); + + GET_F32(b.norm_ff2_w, lname("aligner.enc.blocks.%d.norm_ff2.weight", i), ad); + GET_F32(b.norm_ff2_b, lname("aligner.enc.blocks.%d.norm_ff2.bias", i), ad); + GET_LIN(b.ff2_lin1_w, lname("aligner.enc.blocks.%d.ff2.linear1.weight", i), ad, aff); + GET_LIN(b.ff2_lin2_w, lname("aligner.enc.blocks.%d.ff2.linear2.weight", i), aff, ad); + + GET_F32(b.norm_out_w, lname("aligner.enc.blocks.%d.norm_out.weight", i), ad); + GET_F32(b.norm_out_b, lname("aligner.enc.blocks.%d.norm_out.bias", i), ad); + } + + GET_LIN(aligner.head.weight, "aligner.head.weight", ad, avocab); + GET_F32(aligner.head.bias, "aligner.head.bias", avocab); + } + return TRANSCRIBE_OK; } diff --git a/src/arch/canary/weights.h b/src/arch/canary/weights.h index 99aff864..489598a4 100644 --- a/src/arch/canary/weights.h +++ b/src/arch/canary/weights.h @@ -104,9 +104,27 @@ struct CanaryHParams { float fe_f_max = 0.0f; std::string fe_pad_mode; + // Canary-1B-v2 carries an auxiliary CTC Conformer used only for forced + // alignment. Flash variants leave this absent and emit timestamp tokens + // from the AED decoder instead. + bool aligner_present = false; + int32_t aligner_n_layers = 0; + int32_t aligner_d_model = 0; + int32_t aligner_n_heads = 0; + int32_t aligner_d_ff = 0; + int32_t aligner_conv_kernel = 0; + int32_t aligner_subsampling_factor = 0; + int32_t aligner_subsampling_channels = 0; + int32_t aligner_pos_emb_max_len = 0; + int32_t aligner_vocab_size = 0; + int32_t aligner_blank_id = -1; + std::string aligner_conv_norm_type; + int32_t enc_head_dim() const { return enc_n_heads > 0 ? enc_d_model / enc_n_heads : 0; } int32_t dec_head_dim() const { return dec_n_heads > 0 ? dec_d_model / dec_n_heads : 0; } + + int32_t aligner_head_dim() const { return aligner_n_heads > 0 ? aligner_d_model / aligner_n_heads : 0; } }; transcribe_status read_canary_hparams(const gguf_context * gguf, CanaryHParams & hp); @@ -241,6 +259,17 @@ struct CanaryHead { ggml_tensor * bias = nullptr; // [vocab_size] }; +struct CanaryAlignerHead { + ggml_tensor * weight = nullptr; // [d_model, vocab_size] + ggml_tensor * bias = nullptr; // [vocab_size] +}; + +struct CanaryAlignerWeights { + CanaryPreEncode pre_encode; + std::vector blocks; + CanaryAlignerHead head; +}; + struct CanaryWeights { CanaryPreEncode pre_encode; std::vector blocks; @@ -249,6 +278,7 @@ struct CanaryWeights { std::vector dec_blocks; CanaryDecFinal dec_final; CanaryHead head; + CanaryAlignerWeights aligner; }; transcribe_status build_canary_weights(ggml_context * ctx_meta, const CanaryHParams & hp, CanaryWeights & weights); diff --git a/src/transcribe-tokenizer.cpp b/src/transcribe-tokenizer.cpp index b46c5647..748a1bbb 100644 --- a/src/transcribe-tokenizer.cpp +++ b/src/transcribe-tokenizer.cpp @@ -18,8 +18,10 @@ #include "transcribe-unicode.h" #include +#include #include #include +#include #include #include #include @@ -310,6 +312,180 @@ struct BigramGreater { } }; +struct SentencePieceBigram { + int left = -1; + int right = -1; + float score = 0.0f; + size_t size = 0; +}; + +struct SentencePieceBigramLess { + bool operator()(const SentencePieceBigram & a, const SentencePieceBigram & b) const { + if (a.score != b.score) { + return a.score < b.score; + } + return a.left > b.left; + } +}; + +std::string normalize_decoded_sentencepiece(const std::string & text) { + std::string normalized; + normalized.reserve(text.size() + k_sp_space_len); + + bool at_word_start = true; + for (size_t offset = 0; offset < text.size();) { + const unsigned char byte = static_cast(text[offset]); + if (byte < 0x80 && std::isspace(byte)) { + at_word_start = true; + ++offset; + continue; + } + if (at_word_start) { + normalized.append(k_sp_space, k_sp_space_len); + at_word_start = false; + } + const size_t len = std::min(text.size() - offset, unicode::len_utf8(text[offset])); + normalized.append(text, offset, len); + offset += len; + } + return normalized; +} + +transcribe_status encode_sentencepiece_bpe(const std::string & normalized, + const std::vector & tokens, + const std::vector & scores, + const std::vector & token_types, + const std::unordered_map & piece_to_id, + int unk_id, + std::vector & out_ids) { + constexpr int32_t k_token_type_normal = 1; + constexpr int32_t k_token_type_unused = 5; + + out_ids.clear(); + if (normalized.empty()) { + return TRANSCRIBE_OK; + } + if (scores.size() != tokens.size()) { + return TRANSCRIBE_ERR_GGUF; + } + + auto token_type = [&](int id) -> int32_t { + if (token_types.empty()) { + return k_token_type_normal; + } + return token_types[static_cast(id)]; + }; + auto mergeable = [&](int id) { + if (id < 0 || static_cast(id) >= tokens.size() || id == unk_id) { + return false; + } + const int32_t type = token_type(id); + return type == k_token_type_normal || type == k_token_type_unused; + }; + + std::vector symbols; + symbols.reserve(normalized.size()); + int index = 0; + size_t offset = 0; + while (offset < normalized.size()) { + const size_t len = std::min(normalized.size() - offset, unicode::len_utf8(normalized[offset])); + Symbol symbol; + symbol.text = normalized.data() + offset; + symbol.n = len; + symbol.prev = index - 1; + symbol.next = offset + len == normalized.size() ? -1 : index + 1; + symbols.push_back(symbol); + offset += len; + ++index; + } + + std::priority_queue, SentencePieceBigramLess> queue; + std::unordered_map> reverse_unused_merge; + auto push_bigram = [&](int left, int right) { + if (left < 0 || right < 0) { + return; + } + const std::string left_text(symbols[left].text, symbols[left].n); + const std::string right_text(symbols[right].text, symbols[right].n); + const std::string piece = left_text + right_text; + const auto it = piece_to_id.find(piece); + if (it == piece_to_id.end() || !mergeable(it->second)) { + return; + } + queue.push({ left, right, scores[static_cast(it->second)], piece.size() }); + if (token_type(it->second) == k_token_type_unused) { + reverse_unused_merge[piece] = std::make_pair(left_text, right_text); + } + }; + + for (int i = 1; i < static_cast(symbols.size()); ++i) { + push_bigram(i - 1, i); + } + while (!queue.empty()) { + const SentencePieceBigram bigram = queue.top(); + queue.pop(); + Symbol & left = symbols[bigram.left]; + Symbol & right = symbols[bigram.right]; + if (left.n == 0 || right.n == 0 || left.n + right.n != bigram.size) { + continue; + } + left.n += right.n; + right.n = 0; + left.next = right.next; + if (right.next >= 0) { + symbols[right.next].prev = bigram.left; + } + push_bigram(left.prev, bigram.left); + push_bigram(bigram.left, left.next); + } + + const char * hex = "0123456789ABCDEF"; + bool emit_failed = false; + std::function emit_piece; + emit_piece = [&](const std::string & piece, int depth) { + const auto it = piece_to_id.find(piece); + if (it != piece_to_id.end()) { + const int id = it->second; + if (token_type(id) == k_token_type_unused && depth < 100) { + const auto reverse = reverse_unused_merge.find(piece); + if (reverse != reverse_unused_merge.end()) { + emit_piece(reverse->second.first, depth + 1); + emit_piece(reverse->second.second, depth + 1); + return; + } + } + if (mergeable(id)) { + out_ids.push_back(id); + return; + } + } + + std::vector byte_ids; + byte_ids.reserve(piece.size()); + for (unsigned char byte : piece) { + const char byte_piece[] = { '<', '0', 'x', hex[byte >> 4], hex[byte & 15], '>', '\0' }; + const auto byte_it = piece_to_id.find(byte_piece); + if (byte_it == piece_to_id.end()) { + byte_ids.clear(); + break; + } + byte_ids.push_back(byte_it->second); + } + if (!byte_ids.empty()) { + out_ids.insert(out_ids.end(), byte_ids.begin(), byte_ids.end()); + } else if (unk_id >= 0 && static_cast(unk_id) < tokens.size()) { + out_ids.push_back(unk_id); + } else { + emit_failed = true; + } + }; + + for (int i = 0; i != -1; i = symbols[i].next) { + emit_piece(std::string(symbols[i].text, symbols[i].n), 0); + } + return emit_failed ? TRANSCRIBE_ERR_GGUF : TRANSCRIBE_OK; +} + } // namespace // Tiktoken-style encoder for raw-bytes vocabularies (whisper.cpp .bin): @@ -651,6 +827,18 @@ transcribe_status Tokenizer::encode(const std::string & text, std::vector & out_ids) const { + out_ids.clear(); + if (model_ != "bpe") { + return TRANSCRIBE_ERR_NOT_IMPLEMENTED; + } + if (ids == nullptr || n <= 0) { + return TRANSCRIBE_OK; + } + const std::string normalized = normalize_decoded_sentencepiece(decode(ids, n)); + return encode_sentencepiece_bpe(normalized, tokens_, scores_, token_type_, piece_to_id_, unk_id_, out_ids); +} + // --------------------------------------------------------------------------- // load() - populate the vocab / merges / special ids from GGUF. // --------------------------------------------------------------------------- diff --git a/src/transcribe-tokenizer.h b/src/transcribe-tokenizer.h index fb35823d..4b8ff47b 100644 --- a/src/transcribe-tokenizer.h +++ b/src/transcribe-tokenizer.h @@ -20,6 +20,8 @@ // - encode(): UTF-8 text -> token ids ("gpt2" only; others return // NOT_IMPLEMENTED). Runs the pretokenizer, byte-level encodes each // pretoken, then greedily applies BPE merges in rank order. +// - canonicalize_bpe(): decode a SentencePiece BPE hypothesis, normalize +// its whitespace, and re-tokenize it with the loaded piece scores. #pragma once @@ -183,6 +185,18 @@ class Tokenizer { // merges (the encoder needs them). transcribe_status encode(const std::string & text, std::vector & out_ids) const; + // Canonicalize an already-decoded SentencePiece BPE hypothesis. This is + // intentionally narrower than encode(): decoded model output has already + // passed through the model's SentencePiece normalizer, so only the + // whitespace normalization needed to reconstruct the escaped-space input + // is performed here. "unigram" and all non-BPE models return + // TRANSCRIBE_ERR_NOT_IMPLEMENTED. + transcribe_status canonicalize_bpe(const int * ids, int n, std::vector & out_ids) const; + + // True when canonicalize_bpe() has the SentencePiece BPE model and score + // table it needs. Decode-only BPE tokenizers may legitimately omit scores. + bool has_bpe_canonicalizer() const { return model_ == "bpe" && scores_.size() == tokens_.size(); } + // Identification + special token ids. -1 if the corresponding key // was absent from the GGUF. const std::string & model_type() const { return model_; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bec7501c..19d1a86c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -671,6 +671,26 @@ if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) SKIP_RETURN_CODE 77) endif() +# Canary 1B v2 timestamp result contracts. The timestamp-enabled GGUF is +# required; an optional legacy GGUF also exercises unsupported requests. +if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) + add_executable(transcribe_canary_timestamp_real_smoke + canary_timestamp_real_smoke.cpp) + + target_link_libraries(transcribe_canary_timestamp_real_smoke + PRIVATE transcribe transcribe-common-example) + + target_compile_definitions(transcribe_canary_timestamp_real_smoke PRIVATE + "TRANSCRIBE_TEST_SAMPLES_DIR=\"${CMAKE_SOURCE_DIR}/samples\"") + + transcribe_apply_warnings(transcribe_canary_timestamp_real_smoke) + + add_test(NAME transcribe_canary_timestamp_real_smoke + COMMAND transcribe_canary_timestamp_real_smoke) + set_tests_properties(transcribe_canary_timestamp_real_smoke PROPERTIES + SKIP_RETURN_CODE 77) +endif() + # ----------------------------------------------------------------------------- # Decoder + result accessor accuracy test (real-model gated) # ----------------------------------------------------------------------------- @@ -1152,7 +1172,7 @@ add_test(NAME transcribe_mel_unit COMMAND transcribe_mel_unit) add_executable(transcribe_tokenizer_decode_only_unit tokenizer_decode_only_unit.cpp) -target_link_libraries(transcribe_tokenizer_decode_only_unit PRIVATE transcribe) +target_link_libraries(transcribe_tokenizer_decode_only_unit PRIVATE transcribe ggml) target_include_directories(transcribe_tokenizer_decode_only_unit PRIVATE ${CMAKE_SOURCE_DIR}/src) @@ -1162,6 +1182,23 @@ transcribe_apply_warnings(transcribe_tokenizer_decode_only_unit) add_test(NAME transcribe_tokenizer_decode_only_unit COMMAND transcribe_tokenizer_decode_only_unit) +# ----------------------------------------------------------------------------- +# Canary timestamp alignment unit test +# ----------------------------------------------------------------------------- + +add_executable(transcribe_canary_timestamp_unit + canary_timestamp_unit.cpp) + +target_link_libraries(transcribe_canary_timestamp_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_canary_timestamp_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_canary_timestamp_unit) + +add_test(NAME transcribe_canary_timestamp_unit + COMMAND transcribe_canary_timestamp_unit) + # ----------------------------------------------------------------------------- # Whisper .bin parser unit test # ----------------------------------------------------------------------------- diff --git a/tests/canary_timestamp_real_smoke.cpp b/tests/canary_timestamp_real_smoke.cpp new file mode 100644 index 00000000..2ef77cfa --- /dev/null +++ b/tests/canary_timestamp_real_smoke.cpp @@ -0,0 +1,338 @@ +// canary_timestamp_real_smoke.cpp - real-model gated public-API smoke for +// Canary 1B v2 timestamp behavior. + +#include "transcribe.h" +#include "wav.h" + +#include + +#include +#include +#include +#include +#include + +#ifndef TRANSCRIBE_TEST_SAMPLES_DIR +# define TRANSCRIBE_TEST_SAMPLES_DIR "samples" +#endif + +namespace { + +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) + +#define CHECK_EQ_INT(actual, expected) \ + do { \ + const long long _a = static_cast(actual); \ + const long long _e = static_cast(expected); \ + if (_a != _e) { \ + std::fprintf(stderr, "FAIL %s:%d: %s = %lld, expected %lld\n", __FILE__, __LINE__, #actual, _a, _e); \ + ++g_failures; \ + } \ + } while (0) + +bool file_exists(const std::string & path) { + struct stat st{}; + return ::stat(path.c_str(), &st) == 0; +} + +transcribe_model * load_model(const std::string & path) { + transcribe_model_load_params params; + transcribe_model_load_params_init(¶ms); + + transcribe_model * model = nullptr; + const transcribe_status st = transcribe_model_load_file(path.c_str(), ¶ms, &model); + if (st != TRANSCRIBE_OK || model == nullptr) { + std::fprintf(stderr, "FAIL load %s: %s\n", path.c_str(), transcribe_status_string(st)); + ++g_failures; + return nullptr; + } + return model; +} + +transcribe_session * init_session(transcribe_model * model, int n_ctx = 0) { + transcribe_session_params params; + transcribe_session_params_init(¶ms); + params.n_ctx = n_ctx; + + transcribe_session * session = nullptr; + const transcribe_status st = transcribe_session_init(model, ¶ms, &session); + if (st != TRANSCRIBE_OK || session == nullptr) { + std::fprintf(stderr, "FAIL session init: %s\n", transcribe_status_string(st)); + ++g_failures; + return nullptr; + } + return session; +} + +void check_empty_result(transcribe_session * session) { + CHECK(std::strcmp(transcribe_full_text(session), "") == 0); + CHECK_EQ_INT(transcribe_returned_timestamp_kind(session), TRANSCRIBE_TIMESTAMPS_NONE); + CHECK_EQ_INT(transcribe_n_segments(session), 0); + CHECK_EQ_INT(transcribe_n_words(session), 0); + CHECK_EQ_INT(transcribe_n_tokens(session), 0); +} + +void check_none_result(transcribe_session * session) { + CHECK(transcribe_full_text(session)[0] != '\0'); + CHECK_EQ_INT(transcribe_returned_timestamp_kind(session), TRANSCRIBE_TIMESTAMPS_NONE); + CHECK_EQ_INT(transcribe_n_segments(session), 1); + CHECK_EQ_INT(transcribe_n_words(session), 0); + CHECK_EQ_INT(transcribe_n_tokens(session), 0); + + transcribe_segment segment; + transcribe_segment_init(&segment); + CHECK_EQ_INT(transcribe_get_segment(session, 0, &segment), TRANSCRIBE_OK); + CHECK(segment.text != nullptr && segment.text[0] != '\0'); + CHECK_EQ_INT(segment.t0_ms, 0); + CHECK_EQ_INT(segment.t1_ms, 0); + CHECK_EQ_INT(segment.first_word, 0); + CHECK_EQ_INT(segment.n_words, 0); + CHECK_EQ_INT(segment.first_token, 0); + CHECK_EQ_INT(segment.n_tokens, 0); +} + +void check_word_result(transcribe_session * session) { + CHECK(transcribe_full_text(session)[0] != '\0'); + CHECK_EQ_INT(transcribe_returned_timestamp_kind(session), TRANSCRIBE_TIMESTAMPS_WORD); + + const int n_segments = transcribe_n_segments(session); + const int n_words = transcribe_n_words(session); + CHECK(n_segments > 0); + CHECK(n_words > 0); + CHECK_EQ_INT(transcribe_n_tokens(session), 0); + + int expected_first_word = 0; + int64_t previous_segment_end = -1; + for (int i = 0; i < n_segments; ++i) { + transcribe_segment segment; + transcribe_segment_init(&segment); + CHECK_EQ_INT(transcribe_get_segment(session, i, &segment), TRANSCRIBE_OK); + CHECK(segment.text != nullptr && segment.text[0] != '\0'); + CHECK(segment.t0_ms >= previous_segment_end); + CHECK(segment.t1_ms > segment.t0_ms); + CHECK_EQ_INT(segment.first_word, expected_first_word); + CHECK(segment.n_words > 0); + CHECK_EQ_INT(segment.first_token, 0); + CHECK_EQ_INT(segment.n_tokens, 0); + expected_first_word += segment.n_words; + previous_segment_end = segment.t1_ms; + } + CHECK_EQ_INT(expected_first_word, n_words); + + int64_t previous_word_end = -1; + for (int i = 0; i < n_words; ++i) { + transcribe_word word; + transcribe_word_init(&word); + CHECK_EQ_INT(transcribe_get_word(session, i, &word), TRANSCRIBE_OK); + CHECK(word.text != nullptr && word.text[0] != '\0'); + CHECK(word.t0_ms >= previous_word_end); + CHECK(word.t1_ms > word.t0_ms); + CHECK(word.seg_index >= 0 && word.seg_index < n_segments); + CHECK_EQ_INT(word.first_token, 0); + CHECK_EQ_INT(word.n_tokens, 0); + previous_word_end = word.t1_ms; + } +} + +void check_segment_result(transcribe_session * session) { + CHECK(transcribe_full_text(session)[0] != '\0'); + CHECK_EQ_INT(transcribe_returned_timestamp_kind(session), TRANSCRIBE_TIMESTAMPS_SEGMENT); + + const int n_segments = transcribe_n_segments(session); + CHECK(n_segments > 0); + CHECK_EQ_INT(transcribe_n_words(session), 0); + CHECK_EQ_INT(transcribe_n_tokens(session), 0); + + int64_t previous_end = -1; + for (int i = 0; i < n_segments; ++i) { + transcribe_segment segment; + transcribe_segment_init(&segment); + CHECK_EQ_INT(transcribe_get_segment(session, i, &segment), TRANSCRIBE_OK); + CHECK(segment.text != nullptr && segment.text[0] != '\0'); + CHECK(segment.t0_ms >= previous_end); + CHECK(segment.t1_ms > segment.t0_ms); + CHECK_EQ_INT(segment.first_word, 0); + CHECK_EQ_INT(segment.n_words, 0); + CHECK_EQ_INT(segment.first_token, 0); + CHECK_EQ_INT(segment.n_tokens, 0); + previous_end = segment.t1_ms; + } +} + +void check_batch_truncation_reset(transcribe_model * model, const std::vector & pcm) { + // Force JFK to hit a tiny decoder budget, then verify a following silent + // row does not inherit its per-utterance truncation state. + transcribe_session * session = init_session(model, 18); + if (session == nullptr) { + return; + } + + const std::vector silence(16000, 0.0f); + const float * rows[] = { pcm.data(), silence.data() }; + const int lengths[] = { static_cast(pcm.size()), static_cast(silence.size()) }; + + transcribe_run_params params; + transcribe_run_params_init(¶ms); + params.language = "en"; + params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + + CHECK_EQ_INT(transcribe_run_batch(session, rows, lengths, 2, ¶ms), TRANSCRIBE_OK); + CHECK_EQ_INT(transcribe_batch_n_results(session), 2); + CHECK_EQ_INT(transcribe_batch_status(session, 0), TRANSCRIBE_ERR_OUTPUT_TRUNCATED); + CHECK(transcribe_batch_full_text(session, 0)[0] != '\0'); + CHECK_EQ_INT(transcribe_batch_status(session, 1), TRANSCRIBE_OK); + CHECK(transcribe_was_truncated(session)); + + transcribe_session_free(session); +} + +void check_timestamp_model(const std::string & path, const std::vector & pcm) { + transcribe_model * model = load_model(path); + if (model == nullptr) { + return; + } + + CHECK(std::strcmp(transcribe_model_arch_string(model), "canary") == 0); + CHECK(std::strcmp(transcribe_model_variant_string(model), "canary-1b-v2") == 0); + + transcribe_capabilities caps; + transcribe_capabilities_init(&caps); + CHECK_EQ_INT(transcribe_model_get_capabilities(model, &caps), TRANSCRIBE_OK); + CHECK_EQ_INT(caps.native_sample_rate, 16000); + CHECK_EQ_INT(caps.max_timestamp_kind, TRANSCRIBE_TIMESTAMPS_WORD); + + transcribe_session * session = init_session(model); + if (session == nullptr) { + transcribe_model_free(model); + return; + } + + transcribe_run_params params; + transcribe_run_params_init(¶ms); + params.language = "en"; + params.timestamps = TRANSCRIBE_TIMESTAMPS_NONE; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); + check_none_result(session); + + params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); + check_word_result(session); + + params.timestamps = TRANSCRIBE_TIMESTAMPS_SEGMENT; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); + check_segment_result(session); + + params.task = TRANSCRIBE_TASK_TRANSLATE; + params.target_language = "de"; + params.timestamps = TRANSCRIBE_TIMESTAMPS_AUTO; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); + check_none_result(session); + + // Explicit timestamps are invalid for translated text because the CTC + // aligner emits source-language acoustic labels. The preflight rejection + // preserves the preceding AUTO translation result. + params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), + TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); + check_none_result(session); + + transcribe_session_free(session); + check_batch_truncation_reset(model, pcm); + transcribe_model_free(model); +} + +void check_legacy_model(const std::string & path, const std::vector & pcm) { + transcribe_model * model = load_model(path); + if (model == nullptr) { + return; + } + + CHECK(std::strcmp(transcribe_model_arch_string(model), "canary") == 0); + CHECK(std::strcmp(transcribe_model_variant_string(model), "canary-1b-v2") == 0); + + transcribe_capabilities caps; + transcribe_capabilities_init(&caps); + CHECK_EQ_INT(transcribe_model_get_capabilities(model, &caps), TRANSCRIBE_OK); + CHECK_EQ_INT(caps.max_timestamp_kind, TRANSCRIBE_TIMESTAMPS_NONE); + + transcribe_session * session = init_session(model); + if (session == nullptr) { + transcribe_model_free(model); + return; + } + + transcribe_run_params params; + transcribe_run_params_init(¶ms); + params.language = "en"; + params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), + TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); + check_empty_result(session); + + params.timestamps = TRANSCRIBE_TIMESTAMPS_SEGMENT; + CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), + TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); + check_empty_result(session); + + transcribe_session_free(session); + transcribe_model_free(model); +} + +} // namespace + +int main() { + const char * timestamp_env = std::getenv("TRANSCRIBE_CANARY_TIMESTAMP_GGUF"); + if (timestamp_env == nullptr || timestamp_env[0] == '\0' || !file_exists(timestamp_env)) { + std::fprintf(stderr, + "canary_timestamp_real_smoke: TRANSCRIBE_CANARY_TIMESTAMP_GGUF " + "unset or missing; skipping.\n"); + return 77; + } + + std::string wav_path; + if (const char * audio_env = std::getenv("TRANSCRIBE_TEST_AUDIO"); audio_env != nullptr && audio_env[0] != '\0') { + wav_path = audio_env; + } else { + wav_path = std::string(TRANSCRIBE_TEST_SAMPLES_DIR) + "/jfk.wav"; + } + if (!file_exists(wav_path)) { + std::fprintf(stderr, "canary_timestamp_real_smoke: audio not found: %s\n", wav_path.c_str()); + return EXIT_FAILURE; + } + + std::vector pcm; + std::string wav_error; + if (!transcribe_cli::load_wav_mono_16k(wav_path, pcm, wav_error) || pcm.empty()) { + std::fprintf(stderr, "canary_timestamp_real_smoke: wav load: %s\n", wav_error.c_str()); + return EXIT_FAILURE; + } + + check_timestamp_model(timestamp_env, pcm); + + const char * legacy_env = std::getenv("TRANSCRIBE_CANARY_LEGACY_GGUF"); + if (legacy_env == nullptr || legacy_env[0] == '\0') { + std::fprintf(stderr, + "canary_timestamp_real_smoke: TRANSCRIBE_CANARY_LEGACY_GGUF " + "not set; legacy checks skipped.\n"); + } else if (!file_exists(legacy_env)) { + std::fprintf(stderr, "FAIL legacy model not found: %s\n", legacy_env); + ++g_failures; + } else { + check_legacy_model(legacy_env, pcm); + } + + if (g_failures != 0) { + std::fprintf(stderr, "canary_timestamp_real_smoke: %d failures\n", g_failures); + return EXIT_FAILURE; + } + std::fprintf(stdout, "canary_timestamp_real_smoke: ok\n"); + return EXIT_SUCCESS; +} diff --git a/tests/canary_timestamp_unit.cpp b/tests/canary_timestamp_unit.cpp new file mode 100644 index 00000000..eef32079 --- /dev/null +++ b/tests/canary_timestamp_unit.cpp @@ -0,0 +1,117 @@ +// canary_timestamp_unit.cpp - synthetic CTC forced-alignment tests. + +#include "arch/canary/canary.h" +#include "gguf.h" + +#include +#include +#include + +namespace { + +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) + +std::vector emissions_for_path(const std::vector & path, int vocab_size) { + std::vector emissions(path.size() * static_cast(vocab_size), -20.0f); + for (size_t frame = 0; frame < path.size(); ++frame) { + emissions[frame * static_cast(vocab_size) + path[frame]] = 0.0f; + } + return emissions; +} + +void test_distinct_tokens() { + constexpr int vocab = 3; + constexpr int blank = 2; + const std::vector path = { blank, 0, blank, 1, blank }; + std::vector alignment; + + CHECK(transcribe::canary::viterbi_ctc_alignment(emissions_for_path(path, vocab), static_cast(path.size()), + vocab, blank, { 0, 1 }, alignment)); + CHECK(alignment == std::vector({ 0, 1, 2, 3, 4 })); +} + +void test_repeated_token_requires_blank() { + constexpr int vocab = 4; + constexpr int blank = 3; + const std::vector path = { blank, 1, blank, 1, blank, 2, blank }; + std::vector alignment; + + CHECK(transcribe::canary::viterbi_ctc_alignment(emissions_for_path(path, vocab), static_cast(path.size()), + vocab, blank, { 1, 1, 2 }, alignment)); + CHECK(alignment == std::vector({ 0, 1, 2, 3, 4, 5, 6 })); +} + +void test_invalid_inputs() { + std::vector alignment = { 42 }; + CHECK(!transcribe::canary::viterbi_ctc_alignment({}, 0, 3, 2, { 0 }, alignment)); + CHECK(alignment.empty()); + + const std::vector emissions(6, 0.0f); + CHECK(!transcribe::canary::viterbi_ctc_alignment(emissions, 2, 3, 2, { 2 }, alignment)); + CHECK(!transcribe::canary::viterbi_ctc_alignment(emissions, 2, 3, 2, { 0, 0 }, alignment)); +} + +void test_aligner_limit_only_applies_to_timestamps() { + transcribe::canary::CanaryHParams hp; + hp.aligner_present = true; + hp.aligner_subsampling_factor = 8; + hp.aligner_pos_emb_max_len = 10; + + CHECK(!transcribe::canary::canary_aligner_input_too_long(80, hp, TRANSCRIBE_TIMESTAMPS_WORD)); + CHECK(transcribe::canary::canary_aligner_input_too_long(81, hp, TRANSCRIBE_TIMESTAMPS_WORD)); + CHECK(transcribe::canary::canary_aligner_input_too_long(81, hp, TRANSCRIBE_TIMESTAMPS_AUTO)); + CHECK(!transcribe::canary::canary_aligner_input_too_long(81, hp, TRANSCRIBE_TIMESTAMPS_NONE)); + + hp.aligner_present = false; + CHECK(!transcribe::canary::canary_aligner_input_too_long(81, hp, TRANSCRIBE_TIMESTAMPS_WORD)); +} + +void test_keep_special_tags_in_aligned_words() { + const char * tokens[] = { "", "\xE2\x96\x81hello", "\xE2\x96\x81world", "<|start|>", "<|mid|>", "<|end|>" }; + const float scores[] = { 0.0f, -1.0f, -1.0f, 0.0f, 0.0f, 0.0f }; + const int32_t types[] = { 2, 1, 1, 3, 3, 3 }; + + gguf_context * gguf = gguf_init_empty(); + CHECK(gguf != nullptr); + if (gguf == nullptr) { + return; + } + gguf_set_val_str(gguf, "tokenizer.ggml.model", "bpe"); + gguf_set_arr_str(gguf, "tokenizer.ggml.tokens", tokens, sizeof(tokens) / sizeof(tokens[0])); + gguf_set_arr_data(gguf, "tokenizer.ggml.scores", GGUF_TYPE_FLOAT32, scores, sizeof(scores) / sizeof(scores[0])); + gguf_set_arr_data(gguf, "tokenizer.ggml.token_type", GGUF_TYPE_INT32, types, sizeof(types) / sizeof(types[0])); + gguf_set_val_u32(gguf, "tokenizer.ggml.unknown_token_id", 0); + + transcribe::Tokenizer tokenizer; + CHECK(tokenizer.load(gguf) == TRANSCRIBE_OK); + gguf_free(gguf); + + const std::vector words = + transcribe::canary::canary_preserve_special_tags(tokenizer, { 3, 1, 4, 2, 5 }, { "hello", "world" }); + CHECK(words == std::vector({ "<|start|> hello<|mid|>", "world<|end|>" })); +} + +} // namespace + +int main() { + test_distinct_tokens(); + test_repeated_token_requires_blank(); + test_invalid_inputs(); + test_aligner_limit_only_applies_to_timestamps(); + test_keep_special_tags_in_aligned_words(); + + if (g_failures != 0) { + std::fprintf(stderr, "FAILED: %d check(s)\n", g_failures); + return 1; + } + std::puts("canary timestamp unit: ok"); + return 0; +} diff --git a/tests/golden/canary/canary-1b-v2.manifest.json b/tests/golden/canary/canary-1b-v2.manifest.json index 438ddc34..6e07e904 100644 --- a/tests/golden/canary/canary-1b-v2.manifest.json +++ b/tests/golden/canary/canary-1b-v2.manifest.json @@ -146,12 +146,19 @@ "target_lang": "en", "task": "asr", "pnc": "yes", - "toggle_timestamps": "no", + "toggle_timestamps": "yes", "beam": 1 }, "tolerance_file": "tests/tolerances/canary.json", "cases": [ "jfk" ], - "transcript_compare": "normalized" + "transcript_compare": "normalized", + "timestamp_validation": { + "levels": [ + "word", + "segment" + ], + "tolerance_ms": 80 + } } diff --git a/tests/tokenizer_decode_only_unit.cpp b/tests/tokenizer_decode_only_unit.cpp index 21f6a01a..9b02d34a 100644 --- a/tests/tokenizer_decode_only_unit.cpp +++ b/tests/tokenizer_decode_only_unit.cpp @@ -25,6 +25,7 @@ // tiktoken-style piece_to_id_ ranks. // - DecodeOnlySpecials values are carried into bos/eos/unk/blank ids. +#include "gguf.h" #include "transcribe-tokenizer.h" #include @@ -191,12 +192,86 @@ void test_empty_vocab_rejected() { CHECK(tok.load_decode_only_raw_bytes(empty) == TRANSCRIBE_ERR_INVALID_ARG); } +void test_sentencepiece_bpe_canonicalization() { + const char * tokens[] = { + "", + "\xE2\x96\x81", + "h", + "e", + "l", + "o", + "\xE2\x96\x81h", + "\xE2\x96\x81he", + "ll", + "llo", + "\xE2\x96\x81hello", + "w", + "r", + "d", + "\xE2\x96\x81w", + "\xE2\x96\x81wo", + "\xE2\x96\x81wor", + "\xE2\x96\x81worl", + "\xE2\x96\x81world", + }; + const float scores[] = { + 0.0f, -20.0f, -20.0f, -20.0f, -20.0f, -20.0f, -9.0f, -8.0f, -7.0f, -6.0f, + -5.0f, -20.0f, -20.0f, -20.0f, -9.0f, -8.0f, -7.0f, -6.0f, -5.0f, + }; + int32_t types[sizeof(tokens) / sizeof(tokens[0])] = {}; + types[0] = 2; + for (size_t i = 1; i < sizeof(types) / sizeof(types[0]); ++i) { + types[i] = 1; + } + + gguf_context * gguf = gguf_init_empty(); + CHECK(gguf != nullptr); + if (gguf == nullptr) { + return; + } + gguf_set_val_str(gguf, "tokenizer.ggml.model", "bpe"); + gguf_set_arr_str(gguf, "tokenizer.ggml.tokens", tokens, sizeof(tokens) / sizeof(tokens[0])); + gguf_set_arr_data(gguf, "tokenizer.ggml.scores", GGUF_TYPE_FLOAT32, scores, sizeof(scores) / sizeof(scores[0])); + gguf_set_arr_data(gguf, "tokenizer.ggml.token_type", GGUF_TYPE_INT32, types, sizeof(types) / sizeof(types[0])); + gguf_set_val_u32(gguf, "tokenizer.ggml.unknown_token_id", 0); + + transcribe::Tokenizer tok; + CHECK(tok.load(gguf) == TRANSCRIBE_OK); + gguf_free(gguf); + CHECK(tok.has_bpe_canonicalizer()); + + // This valid decoded sequence is deliberately not the canonical BPE + // segmentation and contains duplicate decoded whitespace. + const int noncanonical[] = { 1, 1, 2, 3, 4, 4, 5, 1, 11, 5, 12, 4, 13, 1 }; + std::vector canonical; + CHECK(tok.canonicalize_bpe(noncanonical, static_cast(sizeof(noncanonical) / sizeof(noncanonical[0])), + canonical) == TRANSCRIBE_OK); + CHECK(canonical == std::vector({ 10, 18 })); + + // The general plain-text encoder remains unavailable for SentencePiece; + // canonicalization does not claim unigram or arbitrary normalizer support. + std::vector encoded; + CHECK(tok.encode("hello world", encoded) == TRANSCRIBE_ERR_NOT_IMPLEMENTED); + + gguf = gguf_init_empty(); + CHECK(gguf != nullptr); + if (gguf != nullptr) { + gguf_set_val_str(gguf, "tokenizer.ggml.model", "bpe"); + gguf_set_arr_str(gguf, "tokenizer.ggml.tokens", tokens, sizeof(tokens) / sizeof(tokens[0])); + transcribe::Tokenizer decode_only; + CHECK(decode_only.load(gguf) == TRANSCRIBE_OK); + CHECK(!decode_only.has_bpe_canonicalizer()); + gguf_free(gguf); + } +} + } // namespace int main() { test_gpt2_decode_only(); test_raw_bytes_decode_only(); test_empty_vocab_rejected(); + test_sentencepiece_bpe_canonicalization(); if (g_failures > 0) { std::fprintf(stderr, "FAILED: %d check(s)\n", g_failures); From e9a7b1ec32474abb58d149096bcd6dd4ddd9d792 Mon Sep 17 00:00:00 2001 From: Braulio Oliveira Date: Sat, 25 Jul 2026 14:58:14 -0300 Subject: [PATCH 2/2] canary: add long-form chunking --- docs/input-limits.md | 18 +- docs/models/canary-1b-v2.md | 5 + docs/models/canary.md | 12 +- docs/porting/families/canary.md | 4 +- examples/cli/main.cpp | 8 +- src/arch/canary/canary.h | 19 ++ src/arch/canary/model.cpp | 401 +++++++++++++++++++++++++- tests/canary_timestamp_real_smoke.cpp | 47 ++- tests/canary_timestamp_unit.cpp | 122 ++++++++ 9 files changed, 607 insertions(+), 29 deletions(-) diff --git a/docs/input-limits.md b/docs/input-limits.md index fb7a6ae0..80cfe120 100644 --- a/docs/input-limits.md +++ b/docs/input-limits.md @@ -56,9 +56,11 @@ 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`. | +| whisper, canary-1b-v2, parakeet (all variants), voxtral_realtime | none (`max_audio_ms = 0`) | Long audio is windowed internally and stitched (whisper, canary-1b-v2), processed by an unbounded/streaming encoder (parakeet, voxtral_realtime), or padded if short. No practical input limit. whisper, canary-1b-v2, and parakeet never error on input 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`. | -Whisper slices audio into 30 s windows with prev-context stitching; parakeet's +Whisper slices audio into 30 s windows with prev-context stitching; +canary-1b-v2 uses 30-40 s windows, selects low-energy acoustic boundaries, +and stitches a 1 s overlap; parakeet's conformer is effectively unbounded (the encoder positional table is recomputed per run, not a fixed wall). This holds for every parakeet variant, including the cache-aware streaming RNN-T member (`nemotron-3.5-asr-streaming-0.6b`): the @@ -71,7 +73,7 @@ need and do not have a length gate. | 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, cohere, canary except canary-1b-v2 | 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`. | These families wrap an LLM-style decoder whose context window (`audio_tokens + prompt + generation`) is the binding constraint. The number of @@ -144,10 +146,12 @@ 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 +Most chunked / unbounded families 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. canary-1b-v2 is the exception: its input remains +unbounded because each window gets a fresh decoder context, while `n_ctx` can +still lower that per-window output-token budget and cause `OUTPUT_TRUNCATED`. +voxtral_realtime is the other 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/canary-1b-v2.md b/docs/models/canary-1b-v2.md index a88a434e..96ed10c3 100644 --- a/docs/models/canary-1b-v2.md +++ b/docs/models/canary-1b-v2.md @@ -93,6 +93,11 @@ If your audio is not already 16 kHz mono WAV, convert it first: ffmpeg -i input.mp3 -ar 16000 -ac 1 output.wav ``` +There is no practical per-call audio limit. Inputs over 40 seconds are +processed in 30-40 second windows selected at low-energy acoustic boundaries, +with 1 second overlap, and stitched into one transcript; word and segment +timestamps remain on the original audio timeline. + CLI flags specific to canary: - `--pnc` / `--no-pnc` — punctuation & capitalization. Note: canary-1b-v2 diff --git a/docs/models/canary.md b/docs/models/canary.md index 693cf822..a65e1020 100644 --- a/docs/models/canary.md +++ b/docs/models/canary.md @@ -44,11 +44,13 @@ each per-variant doc has direct download links. ## Input limits -Every variant 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, shared across -the family. 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). +`canary-1b-v2` has no practical per-call limit. Audio over 40 seconds is split +into 30-40 second windows at low-energy acoustic boundaries with 1 second +overlap, then stitched by subword-token agreement; word timestamps are rebased +onto the original timeline. The other variants retain their encoder positional-table +limit of about **6.7 minutes (400 s)** and reject longer audio up front with +`TRANSCRIBE_ERR_INPUT_TOO_LONG`. See the +[input-length contract](../input-limits.md). ## Quick start diff --git a/docs/porting/families/canary.md b/docs/porting/families/canary.md index 789d2bdd..4694a646 100644 --- a/docs/porting/families/canary.md +++ b/docs/porting/families/canary.md @@ -140,7 +140,7 @@ uv run scripts/bench/run.py \ - canary-1b-v2 / 1b-flash / 180m-flash: `[, , , , ]` (5 slots, `task='asr'|'ast'`). - Output head: LM head over the concatenated SP vocabulary. Decoding is beam search by default for the original canary-1b (beam=5, length_penalty=1.0) and greedy by default for the flash variants (beam=1). - Tokenizer: concatenated SentencePiece — one SP model per language concatenated into a single vocabulary. canary-1b-v2 is 16,384 pieces; flash/180m-flash/1b vocab sizes are not stated on model cards (Stage 2 fills from .nemo). -- Audio length contract: native ≤40 s direct inference. <1 s is symmetrically zero-padded to 1 s. >40 s is handled by an external chunked inference script with 1 s overlap (canary-1b-v2 chunk len defaults to 40 s; canary-1b-flash 10 s; canary-180m-flash 10 s). **Long-form / streaming is out of scope for the v1 port.** +- Audio length contract: native ≤40 s direct inference. <1 s is symmetrically zero-padded to 1 s. Upstream handles >40 s with overlapping chunks. transcribe.cpp implements that path for canary-1b-v2 with low-energy boundaries inside a 30-40 s window, 1 s overlap, and subword-LCS stitching; streaming remains unsupported. The other variants retain the encoder positional-table input cap. ## Capabilities (from intake) @@ -177,7 +177,7 @@ See per-variant `reports/porting/canary//intake.json::known_risks`. Fam 4. **NeMo FilterbankFeatures applies preemph=0.97 BEFORE windowing/STFT.** Standard NeMo ASR frontend trap. Per-feature normalization (per mel band, across time) — not per-utterance global stats. 5. **Cross-attention from transformer decoder to encoder output.** Padding mask on the cross-attention keys must propagate from encoder input lengths after the factor-8 subsampling. 6. **Decoding default differs by variant.** canary-1b: beam=5, length_penalty=1.0. flash variants: beam=1 greedy. Mismatch with upstream's measured-config beam will move WER. -7. **Audio length contract.** Native ≤40 s; <1 s is zero-padded to 1 s; long-form needs an external chunker. Streaming is not native — out of scope for v1. +7. **Audio length contract.** Native ≤40 s; <1 s is zero-padded to 1 s. canary-1b-v2 long-form uses internal 30-40 s overlapping windows; the other variants still expose the encoder positional-table cap. Streaming is not native. 8. **Timestamps are not native to the v2 AED.** canary-1b-v2 aligns the AED transcript with a side CTC Conformer stored in the `.nemo`. The aligner tokenizer IDs must remain identical to the AED tokenizer IDs. Flash timestamp-token decoding remains out of scope. 9. **License divergence.** canary-1b is CC-BY-NC-4.0; the rest are CC-BY-4.0. Converter must surface license correctly in `general.license` so non-commercial restrictions on canary-1b ride along with the GGUF. 10. **.nemo archive distribution.** No HF `config.json` / `tokenizer_config.json` / `generation_config.json` (except canary-1b-flash's encoder-only HF Transformers shim). Converter must mirror the parakeet pattern of unpacking the .nemo before extraction. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 9ec36e2f..ef166eeb 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1297,6 +1297,10 @@ int main(int argc, char ** argv) { { struct transcribe_session_limits lim; transcribe_session_limits_init(&lim); + struct transcribe_capabilities caps; + transcribe_capabilities_init(&caps); + const bool unbounded_audio = + transcribe_model_get_capabilities(model, &caps) == TRANSCRIBE_OK && 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); @@ -1305,6 +1309,8 @@ int main(int argc, char ** argv) { (long long) (lim.max_kv_bytes >> 20)); } std::printf("\n"); + } else if (unbounded_audio) { + std::printf(" max audio: unbounded (long audio chunked internally)\n"); } else if (lim.effective_n_ctx > 0) { // Capped family whose context is too small to fit any audio // plus a prompt (e.g. an aggressively low --n-ctx). @@ -1312,8 +1318,6 @@ int main(int argc, char ** argv) { " max audio: ~0 s (context %d tok too small for " "audio + prompt)\n", lim.effective_n_ctx); - } else { - std::printf(" max audio: unbounded (long audio chunked internally)\n"); } } } diff --git a/src/arch/canary/canary.h b/src/arch/canary/canary.h index d40c059b..2cb15bbe 100644 --- a/src/arch/canary/canary.h +++ b/src/arch/canary/canary.h @@ -50,6 +50,25 @@ std::vector canary_preserve_special_tags(const Tokenizer & const std::vector & generated_ids, const std::vector & aligned_words); +struct CanaryChunkSpan { + int start_sample = 0; + int n_samples = 0; +}; + +struct CanaryTokenSeam { + int previous_keep = 0; + int current_skip = 0; + bool matched = false; +}; + +std::vector canary_long_form_chunks(int n_samples, int sample_rate); +std::vector canary_long_form_chunks(const float * pcm, int n_samples, int sample_rate); + +CanaryTokenSeam canary_token_seam(const std::vector & previous, + const std::vector & current, + int previous_search, + int current_search); + // --------------------------------------------------------------------------- // KV cache for the autoregressive decoder. Same shape as cohere. // --------------------------------------------------------------------------- diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index 2d653d1f..cf7a0508 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -282,10 +282,15 @@ namespace { constexpr float kBnEps = 1e-5f; +constexpr int kCanaryLongFormMinSeconds = 30; +constexpr int kCanaryLongFormMaxSeconds = 40; +constexpr int kCanaryLongFormOverlapMs = 1000; + // Input-length contract (see docs/input-limits.md). Two limits: // (a) INPUT — the encoder rel-pos table (enc_pos_emb_max_len, ~400 s). // T_enc must stay within it or the runtime table aliases past the -// trained range; gated up front. Drives max_audio_ms. +// trained range. Canary v2 stays well below it via 30-40 s long-form +// windows; variants without that coordinator remain gated up front. // (b) DECODER self-KV (dec_max_position) + 512 max-new cap bound the // OUTPUT length; an overrun is kept as a partial and flagged via // transcribe_was_truncated(), not rejected. @@ -318,7 +323,7 @@ int canary_predict_t_enc(int mel_n_frames, int subsampling_factor) { return t; } -// Advisory max_audio_ms: longest audio whose T_enc still fits +// Advisory max_audio_ms for variants without long-form chunking: longest audio whose T_enc still fits // enc_pos_emb_max_len, inverting the rate // ms = T_enc * subsampling_factor * hop_length * 1000 / sr // at T_enc == enc_pos_emb_max_len. Returns 0 (unknown) if any rate hparam @@ -332,6 +337,15 @@ int64_t canary_max_audio_ms(const CanaryHParams & hp) { return mel_frames * hp.fe_hop_length * 1000 / hp.fe_sample_rate; } +bool canary_supports_long_form(const CanaryModel & model) { + return model.variant == "canary-1b-v2"; +} + +bool canary_needs_long_form(const CanaryModel & model, int n_samples) { + return canary_supports_long_form(model) && model.hparams.fe_sample_rate > 0 && + n_samples > static_cast(kCanaryLongFormMaxSeconds) * model.hparams.fe_sample_rate; +} + // Effective decoder self-KV ceiling, in tokens: dec_max_position, optionally // lowered — never raised — by the caller's session n_ctx knob. int canary_context_ceiling(int32_t n_ctx_knob, const CanaryHParams & hp) { @@ -473,10 +487,9 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par } m->caps.max_timestamp_kind = m->hparams.aligner_present ? TRANSCRIBE_TIMESTAMPS_WORD : TRANSCRIBE_TIMESTAMPS_NONE; - // Publish the input-length ceiling now that the encoder positional span - // and frontend rate are known (apply_family_invariants ran before the - // hparams were read). The encoder is the binding INPUT limit for canary. - m->caps.max_audio_ms = canary_max_audio_ms(m->hparams); + // V2 windows long audio internally; older variants still expose the + // encoder positional span as their binding per-call input limit. + m->caps.max_audio_ms = canary_supports_long_form(*m) ? 0 : canary_max_audio_ms(m->hparams); // Basis for transcribe_session_get_limits. audio_from_caps = true pins // effective_max_audio_ms to the encoder bound regardless of n_ctx; the @@ -777,6 +790,13 @@ struct CanaryTimedWord { int n_tokens = 0; }; +struct CanaryChunkResult { + std::vector generated_ids; + std::vector canonical_ids; + std::vector timed_words; + std::vector display_words; +}; + struct CanaryWordRef { std::string text; int first_token = 0; @@ -1014,7 +1034,8 @@ transcribe_status apply_canary_timestamps(CanarySession & cc, CanaryModel & cm, const std::vector & generated_ids, const transcribe_run_params * params, - transcribe_timestamp_kind requested) { + transcribe_timestamp_kind requested, + CanaryChunkResult * chunk_result) { if (requested == TRANSCRIBE_TIMESTAMPS_NONE || cm.caps.max_timestamp_kind == TRANSCRIBE_TIMESTAMPS_NONE) { return TRANSCRIBE_OK; } @@ -1043,14 +1064,20 @@ transcribe_status apply_canary_timestamps(CanarySession & cc, if (params != nullptr && params->keep_special_tags) { display_words = canary_preserve_special_tags(cm.tok, generated_ids, display_words); } + if (chunk_result != nullptr) { + chunk_result->canonical_ids = canonical_ids; + chunk_result->timed_words = timed_words; + chunk_result->display_words = display_words; + } assemble_timed_result(cc, timed_words, display_words, requested); return TRANSCRIBE_OK; } -transcribe_status run(transcribe_session * session, - const float * pcm, - int n_samples, - const transcribe_run_params * params) { +transcribe_status run_chunk(transcribe_session * session, + const float * pcm, + int n_samples, + const transcribe_run_params * params, + CanaryChunkResult * chunk_result) { if (session == nullptr || pcm == nullptr || n_samples <= 0) { return TRANSCRIBE_ERR_INVALID_ARG; } @@ -1061,6 +1088,10 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_INVALID_ARG; } + if (chunk_result != nullptr) { + *chunk_result = {}; + } + if (cc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } @@ -1508,6 +1539,10 @@ transcribe_status run(transcribe_session * session, auto commit_result = [&]() { cc->t_decode_us = ggml_time_us() - t_dec_start; + if (chunk_result != nullptr) { + chunk_result->generated_ids = generated_ids; + } + if (generated_ids.empty()) { return; } @@ -1731,7 +1766,8 @@ transcribe_status run(transcribe_session * session, } if (cc->has_result) { - if (const transcribe_status st = apply_canary_timestamps(*cc, *cm, generated_ids, params, requested_kind); + if (const transcribe_status st = + apply_canary_timestamps(*cc, *cm, generated_ids, params, requested_kind, chunk_result); st != TRANSCRIBE_OK) { return st; } @@ -1742,6 +1778,206 @@ transcribe_status run(transcribe_session * session, return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } +void merge_token_ids(std::vector & merged, + const std::vector & current, + int previous_search, + int current_search) { + if (current.empty()) { + return; + } + const CanaryTokenSeam seam = canary_token_seam(merged, current, previous_search, current_search); + merged.resize(static_cast(seam.previous_keep)); + merged.insert(merged.end(), current.begin() + seam.current_skip, current.end()); +} + +void merge_timed_chunk(std::vector & merged_ids, + std::vector & merged_words, + std::vector & merged_display, + const CanaryChunkResult & chunk, + int64_t chunk_offset_ms) { + if (chunk.canonical_ids.empty() || chunk.timed_words.empty()) { + return; + } + + CanaryTokenSeam seam{ static_cast(merged_ids.size()), 0, false }; + if (!merged_words.empty()) { + const int64_t ownership_boundary = chunk_offset_ms + kCanaryLongFormOverlapMs / 2; + for (const CanaryTimedWord & word : merged_words) { + if ((word.t0_ms + word.t1_ms) / 2 >= ownership_boundary) { + seam.previous_keep = word.first_token; + break; + } + } + seam.current_skip = static_cast(chunk.canonical_ids.size()); + for (const CanaryTimedWord & word : chunk.timed_words) { + if ((word.t0_ms + word.t1_ms) / 2 >= kCanaryLongFormOverlapMs / 2) { + seam.current_skip = word.first_token; + break; + } + } + } + + while (!merged_words.empty() && merged_words.back().first_token >= seam.previous_keep) { + merged_words.pop_back(); + merged_display.pop_back(); + } + merged_ids.resize(static_cast(seam.previous_keep)); + + const int merged_base = seam.previous_keep; + merged_ids.insert(merged_ids.end(), chunk.canonical_ids.begin() + seam.current_skip, chunk.canonical_ids.end()); + for (size_t i = 0; i < chunk.timed_words.size(); ++i) { + CanaryTimedWord word = chunk.timed_words[i]; + if (word.first_token < seam.current_skip) { + continue; + } + word.first_token = merged_base + word.first_token - seam.current_skip; + word.t0_ms += chunk_offset_ms; + word.t1_ms += chunk_offset_ms; + if (!merged_words.empty()) { + if (word.t1_ms <= merged_words.back().t1_ms) { + continue; + } + word.t0_ms = std::max(word.t0_ms, merged_words.back().t1_ms); + } + if (word.t1_ms <= word.t0_ms) { + continue; + } + merged_words.push_back(std::move(word)); + merged_display.push_back(i < chunk.display_words.size() ? chunk.display_words[i] : chunk.timed_words[i].text); + } +} + +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 || !canary_needs_long_form(*cm, n_samples)) { + return run_chunk(session, pcm, n_samples, params, nullptr); + } + + const std::vector chunks = canary_long_form_chunks(pcm, n_samples, cm->hparams.fe_sample_rate); + if (chunks.size() <= 1) { + return run_chunk(session, pcm, n_samples, params, nullptr); + } + + transcribe_run_params chunk_params; + transcribe_run_params_init(&chunk_params); + if (params != nullptr) { + chunk_params = *params; + } + const transcribe_timestamp_kind requested_kind = canary_timestamp_kind_for_run(params); + if (requested_kind == TRANSCRIBE_TIMESTAMPS_SEGMENT) { + chunk_params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + } + + const int frame_ms = + cm->hparams.fe_sample_rate > 0 ? + cm->hparams.fe_hop_length * cm->hparams.enc_subsampling_factor * 1000 / cm->hparams.fe_sample_rate : + 0; + const int delay_frames = std::max(1, frame_ms > 0 ? kCanaryLongFormOverlapMs / frame_ms : 12); + const int previous_search = delay_frames * 2; + const int current_search = std::max(1, delay_frames * 3 / 5); + + std::vector merged_generated; + std::vector merged_text; + std::vector merged_aligned; + std::vector merged_words; + std::vector merged_display; + int64_t total_mel_us = 0; + int64_t total_encode_us = 0; + int64_t total_decode_us = 0; + bool any_result = false; + bool any_truncated = false; + bool timestamps_complete = true; + transcribe_status final_status = TRANSCRIBE_OK; + const bool wants_timestamps = + requested_kind != TRANSCRIBE_TIMESTAMPS_NONE && cm->caps.max_timestamp_kind != TRANSCRIBE_TIMESTAMPS_NONE; + + for (const CanaryChunkSpan & span : chunks) { + if (cc->poll_abort()) { + final_status = TRANSCRIBE_ERR_ABORTED; + break; + } + + cc->was_truncated = false; + cc->t_mel_us = 0; + cc->t_encode_us = 0; + cc->t_decode_us = 0; + CanaryChunkResult chunk_result; + const transcribe_status st = + run_chunk(cc, pcm + span.start_sample, span.n_samples, &chunk_params, &chunk_result); + total_mel_us += cc->t_mel_us; + total_encode_us += cc->t_encode_us; + total_decode_us += cc->t_decode_us; + any_truncated |= cc->was_truncated; + + const bool chunk_has_timestamps = !chunk_result.canonical_ids.empty() && !chunk_result.timed_words.empty(); + const bool keep_unaligned_chunk = !wants_timestamps || st != TRANSCRIBE_ERR_ABORTED || merged_words.empty(); + if (!chunk_result.generated_ids.empty() && (chunk_has_timestamps || keep_unaligned_chunk)) { + any_result = true; + merge_token_ids(merged_generated, chunk_result.generated_ids, previous_search, current_search); + const std::vector clean_ids = params != nullptr && params->keep_special_tags ? + chunk_result.generated_ids : + canary_text_ids(*cm, chunk_result.generated_ids); + merge_token_ids(merged_text, clean_ids, previous_search, current_search); + if (chunk_has_timestamps) { + const int64_t chunk_offset_ms = + static_cast(span.start_sample) * 1000 / cm->hparams.fe_sample_rate; + merge_timed_chunk(merged_aligned, merged_words, merged_display, chunk_result, chunk_offset_ms); + } else if (wants_timestamps) { + timestamps_complete = false; + } + } + + if (st != TRANSCRIBE_OK) { + final_status = st; + break; + } + } + + cc->clear_result(); + cc->t_mel_us = total_mel_us; + cc->t_encode_us = total_encode_us; + cc->t_decode_us = total_decode_us; + cc->was_truncated = any_truncated; + if (final_status != TRANSCRIBE_OK && final_status != TRANSCRIBE_ERR_ABORTED && + final_status != TRANSCRIBE_ERR_OUTPUT_TRUNCATED) { + return final_status; + } + if (any_result) { + cc->raw_text = cm->tok.decode(merged_generated.data(), static_cast(merged_generated.size())); + if (wants_timestamps && timestamps_complete && !merged_words.empty()) { + for (const std::string & word : merged_display) { + if (!cc->full_text.empty()) { + cc->full_text.push_back(' '); + } + cc->full_text += word; + } + assemble_timed_result(*cc, merged_words, merged_display, requested_kind); + } else { + cc->full_text = cm->tok.decode(merged_text.data(), static_cast(merged_text.size())); + if (!cc->full_text.empty() && cc->full_text.front() == ' ') { + cc->full_text.erase(cc->full_text.begin()); + } + transcribe_session::SegmentEntry segment{}; + segment.text = cc->full_text; + cc->segments.push_back(std::move(segment)); + cc->result_kind = TRANSCRIBE_TIMESTAMPS_NONE; + } + cc->has_result = true; + } + + if (final_status == TRANSCRIBE_OK && any_truncated) { + return TRANSCRIBE_ERR_OUTPUT_TRUNCATED; + } + return final_status; +} + // =========================================================================== // Offline batched decode (transcribe_run_batch). Encoder stays serial per // utterance (compute-bound); mel parallel; the autoregressive decode is batched. @@ -1890,6 +2126,12 @@ transcribe_status run_batch(transcribe_session * session, return TRANSCRIBE_ERR_INVALID_ARG; } + for (int i = 0; i < n; ++i) { + if (n_samples[i] > 0 && canary_needs_long_form(*cm, n_samples[i])) { + return run_batch_serial(cc, pcm, n_samples, n, params); + } + } + 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; @@ -2209,6 +2451,141 @@ transcribe_status run_batch(transcribe_session * session, } // namespace +std::vector canary_long_form_chunks(int n_samples, int sample_rate) { + return canary_long_form_chunks(nullptr, n_samples, sample_rate); +} + +std::vector canary_long_form_chunks(const float * pcm, int n_samples, int sample_rate) { + if (n_samples <= 0 || sample_rate <= 0) { + return {}; + } + if (n_samples <= static_cast(kCanaryLongFormMaxSeconds) * sample_rate) { + return { + { 0, n_samples } + }; + } + + const int64_t overlap = static_cast(kCanaryLongFormOverlapMs) * sample_rate / 1000; + if (pcm != nullptr) { + const int64_t min_chunk = static_cast(kCanaryLongFormMinSeconds) * sample_rate; + const int64_t max_chunk = static_cast(kCanaryLongFormMaxSeconds) * sample_rate; + const int64_t window = std::max(1, sample_rate / 5); + const int64_t half = window / 2; + const int64_t right = window - half; + const int64_t stride = std::max(1, sample_rate / 50); + + std::vector chunks; + int64_t start = 0; + while (static_cast(n_samples) - start > max_chunk) { + const int64_t first_center = start + min_chunk; + const int64_t last_center = std::min(start + max_chunk, static_cast(n_samples) - right); + int64_t best_center = first_center; + double best_energy = std::numeric_limits::infinity(); + + int64_t window_begin = first_center - half; + int64_t window_end = window_begin + window; + double energy = 0.0; + for (int64_t i = window_begin; i < window_end; ++i) { + energy += static_cast(pcm[i]) * pcm[i]; + } + for (int64_t center = first_center; center <= last_center; center += stride) { + if (center > first_center) { + const int64_t next_begin = center - half; + const int64_t next_end = next_begin + window; + for (int64_t i = window_begin; i < next_begin; ++i) { + energy -= static_cast(pcm[i]) * pcm[i]; + } + for (int64_t i = window_end; i < next_end; ++i) { + energy += static_cast(pcm[i]) * pcm[i]; + } + window_begin = next_begin; + window_end = next_end; + } + if (energy <= best_energy) { + best_energy = energy; + best_center = center; + } + } + + chunks.push_back({ static_cast(start), static_cast(best_center - start) }); + start = best_center - overlap; + } + chunks.push_back({ static_cast(start), n_samples - static_cast(start) }); + return chunks; + } + + int64_t chunk_size = static_cast(kCanaryLongFormMinSeconds) * sample_rate; + int64_t best_tail = 0; + for (int seconds = kCanaryLongFormMinSeconds; seconds <= kCanaryLongFormMaxSeconds; ++seconds) { + const int64_t candidate = static_cast(seconds) * sample_rate; + const int64_t step = candidate - overlap; + if (step <= 0 || candidate > n_samples) { + continue; + } + const int64_t n_chunks = (static_cast(n_samples) + step - 1) / step; + const int64_t tail = n_samples - (n_chunks - 1) * step; + if (tail > best_tail) { + best_tail = tail; + chunk_size = candidate; + } + } + + const int64_t step = chunk_size - overlap; + std::vector chunks; + for (int64_t start = 0; start + overlap < n_samples; start += step) { + chunks.push_back({ static_cast(start), + static_cast(std::min(chunk_size, static_cast(n_samples) - start)) }); + } + return chunks; +} + +CanaryTokenSeam canary_token_seam(const std::vector & previous, + const std::vector & current, + int previous_search, + int current_search) { + CanaryTokenSeam 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_prev_end = static_cast(previous.size()); + 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 int prev_end = i + length; + const int curr_end = j + length; + // Preserve the previous hypothesis. Only trim current tokens that + // agree with its true suffix; repeated interior text is ambiguous. + const bool at_previous_end = prev_end == static_cast(previous.size()); + const bool accepted = at_previous_end && (length >= 2 || (length == 1 && j == 0)); + if (!accepted) { + continue; + } + if (length > best_length || + (length == best_length && length > 0 && + (prev_end > best_prev_end || (prev_end == best_prev_end && curr_end < best_curr_end)))) { + best_length = length; + best_prev_end = prev_end; + best_curr_end = curr_end; + } + } + } + if (best_length > 0) { + seam.previous_keep = static_cast(previous.size()); + seam.current_skip = best_curr_end; + seam.matched = true; + } + return seam; +} + bool canary_aligner_input_too_long(int mel_n_frames, const CanaryHParams & hp, transcribe_timestamp_kind requested_kind) { diff --git a/tests/canary_timestamp_real_smoke.cpp b/tests/canary_timestamp_real_smoke.cpp index 2ef77cfa..8f1bbdd8 100644 --- a/tests/canary_timestamp_real_smoke.cpp +++ b/tests/canary_timestamp_real_smoke.cpp @@ -176,7 +176,7 @@ void check_batch_truncation_reset(transcribe_model * model, const std::vector silence(16000, 0.0f); - const float * rows[] = { pcm.data(), silence.data() }; + const float * rows[] = { pcm.data(), silence.data() }; const int lengths[] = { static_cast(pcm.size()), static_cast(silence.size()) }; transcribe_run_params params; @@ -208,6 +208,7 @@ void check_timestamp_model(const std::string & path, const std::vector & CHECK_EQ_INT(transcribe_model_get_capabilities(model, &caps), TRANSCRIBE_OK); CHECK_EQ_INT(caps.native_sample_rate, 16000); CHECK_EQ_INT(caps.max_timestamp_kind, TRANSCRIBE_TIMESTAMPS_WORD); + CHECK_EQ_INT(caps.max_audio_ms, 0); transcribe_session * session = init_session(model); if (session == nullptr) { @@ -215,6 +216,11 @@ void check_timestamp_model(const std::string & path, const std::vector & return; } + transcribe_session_limits limits; + transcribe_session_limits_init(&limits); + CHECK_EQ_INT(transcribe_session_get_limits(session, &limits), TRANSCRIBE_OK); + CHECK_EQ_INT(limits.effective_max_audio_ms, 0); + transcribe_run_params params; transcribe_run_params_init(¶ms); params.language = "en"; @@ -225,6 +231,7 @@ void check_timestamp_model(const std::string & path, const std::vector & params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); check_word_result(session); + const int short_word_count = transcribe_n_words(session); params.timestamps = TRANSCRIBE_TIMESTAMPS_SEGMENT; CHECK_EQ_INT(transcribe_run(session, pcm.data(), static_cast(pcm.size()), ¶ms), TRANSCRIBE_OK); @@ -244,6 +251,34 @@ void check_timestamp_model(const std::string & path, const std::vector & TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); check_none_result(session); + std::vector long_pcm; + long_pcm.reserve(pcm.size() * 5); + for (int i = 0; i < 5; ++i) { + long_pcm.insert(long_pcm.end(), pcm.begin(), pcm.end()); + } + params.task = TRANSCRIBE_TASK_TRANSCRIBE; + params.target_language = nullptr; + params.timestamps = TRANSCRIBE_TIMESTAMPS_WORD; + CHECK_EQ_INT(transcribe_run(session, long_pcm.data(), static_cast(long_pcm.size()), ¶ms), TRANSCRIBE_OK); + check_word_result(session); + CHECK_EQ_INT(transcribe_n_words(session), short_word_count * 5); + std::string joined_words; + for (int i = 0; i < transcribe_n_words(session); ++i) { + transcribe_word word; + transcribe_word_init(&word); + CHECK_EQ_INT(transcribe_get_word(session, i, &word), TRANSCRIBE_OK); + if (!joined_words.empty()) { + joined_words.push_back(' '); + } + joined_words += word.text; + } + CHECK(joined_words == transcribe_full_text(session)); + transcribe_word last_word; + transcribe_word_init(&last_word); + CHECK_EQ_INT(transcribe_get_word(session, transcribe_n_words(session) - 1, &last_word), TRANSCRIBE_OK); + CHECK(last_word.t1_ms > 40000); + CHECK(last_word.t1_ms <= static_cast(long_pcm.size()) * 1000 / 16000); + transcribe_session_free(session); check_batch_truncation_reset(model, pcm); transcribe_model_free(model); @@ -262,6 +297,7 @@ void check_legacy_model(const std::string & path, const std::vector & pcm transcribe_capabilities_init(&caps); CHECK_EQ_INT(transcribe_model_get_capabilities(model, &caps), TRANSCRIBE_OK); CHECK_EQ_INT(caps.max_timestamp_kind, TRANSCRIBE_TIMESTAMPS_NONE); + CHECK_EQ_INT(caps.max_audio_ms, 0); transcribe_session * session = init_session(model); if (session == nullptr) { @@ -282,6 +318,15 @@ void check_legacy_model(const std::string & path, const std::vector & pcm TRANSCRIBE_ERR_UNSUPPORTED_TIMESTAMPS); check_empty_result(session); + std::vector long_pcm; + long_pcm.reserve(pcm.size() * 5); + for (int i = 0; i < 5; ++i) { + long_pcm.insert(long_pcm.end(), pcm.begin(), pcm.end()); + } + params.timestamps = TRANSCRIBE_TIMESTAMPS_NONE; + CHECK_EQ_INT(transcribe_run(session, long_pcm.data(), static_cast(long_pcm.size()), ¶ms), TRANSCRIBE_OK); + check_none_result(session); + transcribe_session_free(session); transcribe_model_free(model); } diff --git a/tests/canary_timestamp_unit.cpp b/tests/canary_timestamp_unit.cpp index eef32079..e607d625 100644 --- a/tests/canary_timestamp_unit.cpp +++ b/tests/canary_timestamp_unit.cpp @@ -3,7 +3,9 @@ #include "arch/canary/canary.h" #include "gguf.h" +#include #include +#include #include #include @@ -99,6 +101,123 @@ void test_keep_special_tags_in_aligned_words() { CHECK(words == std::vector({ "<|start|> hello<|mid|>", "world<|end|>" })); } +void test_long_form_chunk_schedule() { + constexpr int sample_rate = 16000; + + const std::vector short_chunks = + transcribe::canary::canary_long_form_chunks(40 * sample_rate, sample_rate); + CHECK(short_chunks.size() == 1); + if (short_chunks.size() == 1) { + CHECK(short_chunks[0].start_sample == 0); + CHECK(short_chunks[0].n_samples == 40 * sample_rate); + } + + const int barely_long_samples = 40 * sample_rate + 1; + const std::vector barely_long_chunks = + transcribe::canary::canary_long_form_chunks(barely_long_samples, sample_rate); + CHECK(barely_long_chunks.size() == 2); + if (!barely_long_chunks.empty()) { + CHECK(static_cast(barely_long_chunks.back().start_sample) + barely_long_chunks.back().n_samples == + barely_long_samples); + } + + const std::vector chunks = + transcribe::canary::canary_long_form_chunks(80 * sample_rate, sample_rate); + CHECK(chunks.size() == 3); + if (chunks.size() == 3) { + CHECK(chunks[0].start_sample == 0); + CHECK(chunks[0].n_samples == 30 * sample_rate); + CHECK(chunks[1].start_sample == 29 * sample_rate); + CHECK(chunks[1].n_samples == 30 * sample_rate); + CHECK(chunks[2].start_sample == 58 * sample_rate); + CHECK(chunks[2].n_samples == 22 * sample_rate); + } + + const int max_samples = std::numeric_limits::max(); + const std::vector max_chunks = + transcribe::canary::canary_long_form_chunks(max_samples, sample_rate); + CHECK(!max_chunks.empty()); + if (!max_chunks.empty()) { + CHECK(max_chunks.front().start_sample == 0); + CHECK(static_cast(max_chunks.back().start_sample) + max_chunks.back().n_samples == max_samples); + for (size_t i = 0; i < max_chunks.size(); ++i) { + CHECK(max_chunks[i].start_sample >= 0); + CHECK(max_chunks[i].n_samples > 0); + CHECK(max_chunks[i].n_samples <= 40 * sample_rate); + CHECK(static_cast(max_chunks[i].start_sample) + max_chunks[i].n_samples <= max_samples); + if (i > 0) { + CHECK(max_chunks[i].start_sample > max_chunks[i - 1].start_sample); + CHECK(max_chunks[i].start_sample < + static_cast(max_chunks[i - 1].start_sample) + max_chunks[i - 1].n_samples); + } + } + } +} + +void test_long_form_token_seam() { + const transcribe::canary::CanaryTokenSeam exact = + transcribe::canary::canary_token_seam({ 1, 2, 3, 4, 5 }, { 4, 5, 6, 7 }, 5, 4); + CHECK(exact.matched); + CHECK(exact.previous_keep == 5); + CHECK(exact.current_skip == 2); + + const transcribe::canary::CanaryTokenSeam replacement = + transcribe::canary::canary_token_seam({ 1, 2, 3, 9 }, { 2, 3, 4 }, 4, 3); + CHECK(!replacement.matched); + CHECK(replacement.previous_keep == 4); + CHECK(replacement.current_skip == 0); + + const transcribe::canary::CanaryTokenSeam interior = + transcribe::canary::canary_token_seam({ 1, 2, 3, 4 }, { 9, 3, 8 }, 4, 3); + CHECK(!interior.matched); + CHECK(interior.previous_keep == 4); + CHECK(interior.current_skip == 0); + + const transcribe::canary::CanaryTokenSeam strong_interior = + transcribe::canary::canary_token_seam({ 10, 1, 2, 3, 4, 11 }, { 20, 1, 2, 3, 4, 21 }, 6, 6); + CHECK(!strong_interior.matched); + CHECK(strong_interior.previous_keep == 6); + CHECK(strong_interior.current_skip == 0); + + const transcribe::canary::CanaryTokenSeam exact_one = + transcribe::canary::canary_token_seam({ 1, 2, 3 }, { 3, 4 }, 3, 2); + CHECK(exact_one.matched); + CHECK(exact_one.previous_keep == 3); + CHECK(exact_one.current_skip == 1); + + const transcribe::canary::CanaryTokenSeam none = transcribe::canary::canary_token_seam({ 1, 2 }, { 3, 4 }, 2, 2); + CHECK(!none.matched); + CHECK(none.previous_keep == 2); + CHECK(none.current_skip == 0); +} + +void test_long_form_prefers_quiet_boundaries() { + constexpr int sample_rate = 16000; + std::vector pcm(70 * sample_rate, 0.5f); + std::fill(pcm.begin() + 32 * sample_rate, pcm.begin() + 33 * sample_rate, 0.0f); + + const std::vector chunks = + transcribe::canary::canary_long_form_chunks(pcm.data(), static_cast(pcm.size()), sample_rate); + CHECK(chunks.size() == 2); + if (chunks.size() == 2) { + CHECK(chunks[0].start_sample == 0); + CHECK(chunks[0].n_samples >= 32 * sample_rate); + CHECK(chunks[0].n_samples <= 33 * sample_rate); + CHECK(chunks[1].start_sample == chunks[0].n_samples - sample_rate); + CHECK(chunks[1].start_sample + chunks[1].n_samples == static_cast(pcm.size())); + } + + constexpr int odd_sample_rate = 8016; + std::vector odd_pcm(40 * odd_sample_rate + 1, 0.5f); + const std::vector odd_chunks = + transcribe::canary::canary_long_form_chunks(odd_pcm.data(), static_cast(odd_pcm.size()), odd_sample_rate); + CHECK(!odd_chunks.empty()); + if (!odd_chunks.empty()) { + CHECK(static_cast(odd_chunks.back().start_sample) + odd_chunks.back().n_samples == + static_cast(odd_pcm.size())); + } +} + } // namespace int main() { @@ -107,6 +226,9 @@ int main() { test_invalid_inputs(); test_aligner_limit_only_applies_to_timestamps(); test_keep_special_tags_in_aligned_words(); + test_long_form_chunk_schedule(); + test_long_form_token_seam(); + test_long_form_prefers_quiet_boundaries(); if (g_failures != 0) { std::fprintf(stderr, "FAILED: %d check(s)\n", g_failures);