From 1acc074217d31dcd0fe851d3e08f9843731bcdbc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnel=20S=C3=A1nchez=20Rodr=C3=ADguez?= Date: Fri, 6 Nov 2026 16:07:32 +0100 Subject: [PATCH] canary: bias the greedy decode toward vocabulary terms Canary has no prompt slot for custom terms: the decoder-context slot made the output loop and over-insert. transcribe_run_params::vocabulary now biases the decode instead, by shallow fusion on the raw logits: a term's first token gets +strength at any step, and each following token while the transcript so far ends with the term's previous tokens. Each term is also tried with its first letter's case flipped. Strength defaults to 3; TRANSCRIBE_CANARY_VOCAB_BOOST tunes it and 0 turns it off (output then identical to no vocabulary). Measured on canary-1b-v2 over 312 phone turns answering "which area?" (8 kHz, es/de): 21 more area names read, about 5 inserted into unrelated speech, 3 lost to answers that named two areas; no added empty output, p50 latency 242 -> 258 ms. A strength of 5 started to loop. VOCABULARY is advertised on the canary2 checkpoints that encode a prefix (1b-v2, 1b-flash, 180m-flash). The static GPU step graph now also outputs the raw logits, and a batch with a vocabulary decodes serially. Co-Authored-By: Claude Opus 5.5 --- docs/models/canary-180m-flash.md | 2 +- docs/models/canary-1b-flash.md | 2 +- docs/models/canary-1b-v2.md | 2 +- docs/prompting.md | 2 +- src/arch/canary/decoder.cpp | 4 ++ src/arch/canary/decoder.h | 1 + src/arch/canary/model.cpp | 74 +++++++++++++++++++----- src/arch/canary/vocab_boost.h | 86 ++++++++++++++++++++++++++++ tests/CMakeLists.txt | 16 ++++++ tests/canary_vocab_boost_unit.cpp | 94 +++++++++++++++++++++++++++++++ 10 files changed, 264 insertions(+), 19 deletions(-) create mode 100644 src/arch/canary/vocab_boost.h create mode 100644 tests/canary_vocab_boost_unit.cpp diff --git a/docs/models/canary-180m-flash.md b/docs/models/canary-180m-flash.md index 30ff3d538..9267ad927 100644 --- a/docs/models/canary-180m-flash.md +++ b/docs/models/canary-180m-flash.md @@ -25,7 +25,7 @@ Not a streaming model. Word and segment timestamps are upstream-experimental and not exposed in the v1 port (deferred — would require porting the `_timestamps_asr_model` CTC aligner from the `.nemo` archive). -**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: any context text made this model stop early in testing. +**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: any context text made this model stop early in testing. Vocabulary (`--vocabulary`) biases the greedy decode toward the terms, as on [canary-1b-v2](canary-1b-v2.md); measured on 1b-v2 only. See NVIDIA's [model card](https://huggingface.co/nvidia/canary-180m-flash) for training data, intended use, and upstream evaluation methodology. diff --git a/docs/models/canary-1b-flash.md b/docs/models/canary-1b-flash.md index b443dcde7..4f3f52b30 100644 --- a/docs/models/canary-1b-flash.md +++ b/docs/models/canary-1b-flash.md @@ -21,7 +21,7 @@ Offline multilingual speech-to-text and translation. The model takes a - **Translation** between English and German, Spanish, or French (both directions). -**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: any context text made this model stop early in testing. +**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: any context text made this model stop early in testing. Vocabulary (`--vocabulary`) biases the greedy decode toward the terms, as on [canary-1b-v2](canary-1b-v2.md); measured on 1b-v2 only. See NVIDIA's [model card](https://huggingface.co/nvidia/canary-1b-flash) for training data, intended use, and upstream evaluation methodology. diff --git a/docs/models/canary-1b-v2.md b/docs/models/canary-1b-v2.md index 4ed8bd40f..71240b84c 100644 --- a/docs/models/canary-1b-v2.md +++ b/docs/models/canary-1b-v2.md @@ -33,7 +33,7 @@ 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. -**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: context text made the output loop and over-insert in testing. +**Prompting:** transcript prefix (`--prefix`) via NeMo's `user_prefix` turn. The decoder-context slot is not exposed: context text made the output loop and over-insert in testing. Vocabulary (`--vocabulary`) biases the greedy decode toward the terms instead (shallow fusion on the logits, strength 3, `TRANSCRIBE_CANARY_VOCAB_BOOST` to tune, `0` turns it off): on 312 phone turns answering "which area?" it read 21 more area names, about 5 of them inserted into unrelated speech. Batch runs with a vocabulary decode one utterance at a time. See NVIDIA's [model card](https://huggingface.co/nvidia/canary-1b-v2) for training data, intended use, and upstream evaluation methodology. diff --git a/docs/prompting.md b/docs/prompting.md index d908e5c2d..656480431 100644 --- a/docs/prompting.md +++ b/docs/prompting.md @@ -17,7 +17,7 @@ Probe `transcribe_model_supports()` for the matching feature bit. | Voxtral (2507) | | | yes | | | Granite 4.0-1b / 4.1-2b | yes | | | | | Granite 4.1-2b-plus | yes | | | yes | -| Canary 180m-flash / 1b-flash / 1b-v2 | | | | yes | +| Canary 180m-flash / 1b-flash / 1b-v2 | yes | | | yes | | Fun-ASR-Nano | yes | | | | | MOSS-Transcribe-Diarize | yes | | | | diff --git a/src/arch/canary/decoder.cpp b/src/arch/canary/decoder.cpp index 6683cdeef..e37a86a0d 100644 --- a/src/arch/canary/decoder.cpp +++ b/src/arch/canary/decoder.cpp @@ -614,6 +614,10 @@ StepBuild build_step_graph(ggml_context * ctx, logits = ggml_add(ctx, logits, w.head.bias); } + sb.logits_out = logits; + ggml_set_name(sb.logits_out, "step.logits_raw"); + ggml_set_output(sb.logits_out); + sb.argmax_out = ggml_argmax(ctx, logits); ggml_set_name(sb.argmax_out, "step.argmax"); ggml_set_output(sb.argmax_out); diff --git a/src/arch/canary/decoder.h b/src/arch/canary/decoder.h index b0ff35393..49e05948e 100644 --- a/src/arch/canary/decoder.h +++ b/src/arch/canary/decoder.h @@ -87,6 +87,7 @@ struct StepBuild { ggml_tensor * kv_idx_in = nullptr; // i64 [1] ggml_tensor * mask_in = nullptr; // f16 [max_n_kv, 1] ggml_tensor * argmax_out = nullptr; // i32 [1] + ggml_tensor * logits_out = nullptr; // f32 [vocab, 1], raw (vocabulary bias) int max_n_kv = 0; ggml_cgraph * graph = nullptr; }; diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index cdaf3ac64..1cf885d13 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -25,6 +25,7 @@ #include "transcribe-meta.h" #include "transcribe-prompting.h" #include "transcribe-repetition-guard.h" +#include "vocab_boost.h" #include "weights.h" #include @@ -384,9 +385,12 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par } // Transcript prefix: canary2's user_prefix (measured clean on all three // canary2 checkpoints). Needs a sub-vocab range for aggregate tokenizers. - transcribe::set_feature(m.get(), TRANSCRIBE_FEATURE_TRANSCRIPT_PREFIX, - m->hparams.prompt_format == "canary2" && - (m->hparams.tokenizer_single_sp || !m->hparams.tok_lang_codes.empty())); + // Vocabulary terms are encoded the same way and bias the greedy decode + // (vocab_boost.h). + const bool encodes_text = + m->hparams.prompt_format == "canary2" && (m->hparams.tokenizer_single_sp || !m->hparams.tok_lang_codes.empty()); + transcribe::set_feature(m.get(), TRANSCRIBE_FEATURE_TRANSCRIPT_PREFIX, encodes_text); + transcribe::set_feature(m.get(), TRANSCRIBE_FEATURE_VOCABULARY, encodes_text); // Publish the input-length ceiling now that the encoder positional span // and frontend rate are known (apply_family_invariants ran before the @@ -988,6 +992,7 @@ transcribe_status run(transcribe_session * session, std::vector prompt_ids; std::vector prefix_ids; + VocabBoost boost; if (cm->hparams.prompt_format == "canary2") { if (const transcribe_status st = encode_canary2_prefix(*cm, params != nullptr ? params->prefix : nullptr, tgt_lang, prefix_ids); @@ -995,6 +1000,33 @@ transcribe_status run(transcribe_session * session, return st; } prompt_ids = build_prompt_canary2(*cm, cm->hparams, src_id, tgt_id, task, pnc, prefix_ids); + if (params != nullptr && params->n_vocabulary > 0) { + boost.strength = vocab_boost_strength(); + // Each term as given and with its first ASCII letter's case flipped: + // the transcript starts a sentence capitalized and runs on lowercase. + for (std::string spelling : transcribe::prompting::terms(params)) { + for (int form = 0; form < 2; ++form) { + if (form == 1) { + char & c = spelling[0]; + if (c >= 'a' && c <= 'z') { + c = static_cast(c - 'a' + 'A'); + } else if (c >= 'A' && c <= 'Z') { + c = static_cast(c - 'A' + 'a'); + } else { + break; + } + } + std::vector ids; + if (const transcribe_status st = encode_canary2_prefix(*cm, spelling.c_str(), tgt_lang, ids); + st != TRANSCRIBE_OK) { + return st; + } + if (!ids.empty()) { + boost.terms.push_back(std::move(ids)); + } + } + } + } } else if (cm->hparams.prompt_format == "canary") { prompt_ids = build_prompt_canary(cm->hparams, src_id, tgt_id, task, pnc); } @@ -1183,7 +1215,7 @@ transcribe_status run(transcribe_session * session, k_gen_reserve, prompt_len, cc->kv_cache.n_ctx); int next_token = 0; - if (prompt_skip_softmax && db.argmax_out != nullptr) { + if (!boost.active() && prompt_skip_softmax && db.argmax_out != nullptr) { int32_t argmax_id = 0; ggml_backend_tensor_get(db.argmax_out, &argmax_id, 0, sizeof(int32_t)); next_token = argmax_id; @@ -1192,13 +1224,7 @@ transcribe_status run(transcribe_session * session, std::vector logits_host(static_cast(vocab_size) * prompt_len); ggml_backend_tensor_get(db.out, logits_host.data(), 0, logits_host.size() * sizeof(float)); const float * last_logits = logits_host.data() + static_cast(prompt_len - 1) * vocab_size; - float best = last_logits[0]; - for (int j = 1; j < static_cast(vocab_size); ++j) { - if (last_logits[j] > best) { - best = last_logits[j]; - next_token = j; - } - } + next_token = boost.pick(last_logits, vocab_size, {}); } std::vector generated_ids; @@ -1206,7 +1232,8 @@ transcribe_status run(transcribe_session * session, generated_ids.push_back(next_token); } - int n_past = prompt_len; + int n_past = prompt_len; + std::vector step_logits; // host copy of a step's logits, only under a vocabulary bias // Commit accumulated generated_ids as the run's segment + full // text. Called both on normal loop exit and on abort so the @@ -1361,7 +1388,14 @@ transcribe_status run(transcribe_session * session, cc->kv_cache.head = n_past; int32_t argmax_id = 0; - ggml_backend_tensor_get(sb.argmax_out, &argmax_id, 0, sizeof(int32_t)); + if (boost.active()) { + const int64_t vocab_size = sb.logits_out->ne[0]; + step_logits.resize(static_cast(vocab_size)); + ggml_backend_tensor_get(sb.logits_out, step_logits.data(), 0, step_logits.size() * sizeof(float)); + argmax_id = boost.pick(step_logits.data(), vocab_size, generated_ids); + } else { + ggml_backend_tensor_get(sb.argmax_out, &argmax_id, 0, sizeof(int32_t)); + } next_token = argmax_id; if (next_token != eos_id) { @@ -1438,7 +1472,14 @@ transcribe_status run(transcribe_session * session, cc->kv_cache.head = n_past; int32_t argmax_id = 0; - ggml_backend_tensor_get(db_step.argmax_out, &argmax_id, 0, sizeof(int32_t)); + if (boost.active()) { + const int64_t vocab_size = db_step.out->ne[0]; + step_logits.resize(static_cast(vocab_size)); + ggml_backend_tensor_get(db_step.out, step_logits.data(), 0, step_logits.size() * sizeof(float)); + argmax_id = boost.pick(step_logits.data(), vocab_size, generated_ids); + } else { + ggml_backend_tensor_get(db_step.argmax_out, &argmax_id, 0, sizeof(int32_t)); + } next_token = argmax_id; if (next_token != eos_id) { @@ -1608,7 +1649,10 @@ transcribe_status run_batch(transcribe_session * session, const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && cm->plan.primary_kind != transcribe::BackendKind::Accel && cm->plan.primary_kind != transcribe::BackendKind::Unknown; - if (n == 1 || !cc->decoder_use_flash || !primary_is_gpu || transcribe::debug::enabled()) { + // The batched decode takes a shared argmax; a vocabulary bias needs the + // per-utterance logits, so it runs serial. + const bool has_vocabulary = params != nullptr && params->n_vocabulary > 0; + if (n == 1 || !cc->decoder_use_flash || !primary_is_gpu || has_vocabulary || transcribe::debug::enabled()) { return run_batch_serial(cc, pcm, n_samples, n, params); } diff --git a/src/arch/canary/vocab_boost.h b/src/arch/canary/vocab_boost.h new file mode 100644 index 000000000..72fc736e0 --- /dev/null +++ b/src/arch/canary/vocab_boost.h @@ -0,0 +1,86 @@ +// arch/canary/vocab_boost.h - vocabulary biasing for canary's greedy decode. +// +// INTERNAL. Canary has no prompt slot for custom terms: the decoder-context +// slot makes the output loop and over-insert (docs/models/canary-1b-v2.md). +// So transcribe_run_params::vocabulary biases the decode instead, by shallow +// fusion on the raw logits: a term's first token gets +strength at any step, +// and each following token gets +strength while the transcript so far ends +// with the term's previous tokens. Measured on canary-1b-v2 over 312 phone +// turns answering "which area?" (8 kHz, es/de): strength 3 read 21 more area +// names, about 5 of them inserted into unrelated speech, against 3 lost to +// answers that named two areas; no added empty output, no added latency. +// Strength 5 started to loop and inserted more. + +#pragma once + +#include "transcribe-env.h" + +#include +#include +#include + +namespace transcribe::canary { + +inline constexpr float k_vocab_boost_default = 3.0f; + +// TRANSCRIBE_CANARY_VOCAB_BOOST, or the default when unset or not a number. +// A value <= 0 turns the bias off. +inline float vocab_boost_strength() { + const char * raw = transcribe::env::str("TRANSCRIBE_CANARY_VOCAB_BOOST"); + if (raw == nullptr) { + return k_vocab_boost_default; + } + char * end = nullptr; + const float value = std::strtof(raw, &end); + return end == raw ? k_vocab_boost_default : value; +} + +struct VocabBoost { + std::vector> terms; // token ids, one entry per spelling + float strength = 0.0f; + + bool active() const { return strength > 0.0f && !terms.empty(); } + + // The next token: the argmax of `logits` after the bias. `generated` is + // the transcript so far (prompt excluded). + int pick(const float * logits, int64_t vocab, const std::vector & generated) const { + int best_id = 0; + float best = logits[0]; + for (int64_t j = 1; j < vocab; ++j) { + if (logits[j] > best) { + best = logits[j]; + best_id = static_cast(j); + } + } + if (!active()) { + return best_id; + } + for (const std::vector & term : terms) { + for (size_t k = 0; k < term.size() && k <= generated.size(); ++k) { + if (!ends_with(generated, term, k) || term[k] < 0 || term[k] >= vocab) { + continue; + } + const float score = logits[term[k]] + strength; + if (score > best) { + best = score; + best_id = term[k]; + } + } + } + return best_id; + } + + private: + // Whether `generated` ends with the first `k` tokens of `term`. + static bool ends_with(const std::vector & generated, const std::vector & term, size_t k) { + const size_t offset = generated.size() - k; + for (size_t i = 0; i < k; ++i) { + if (generated[offset + i] != term[i]) { + return false; + } + } + return true; + } +}; + +} // namespace transcribe::canary diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 110645b02..64e48923e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -310,6 +310,22 @@ transcribe_apply_warnings(transcribe_repetition_guard_unit) add_test(NAME transcribe_repetition_guard_unit COMMAND transcribe_repetition_guard_unit) +# ----------------------------------------------------------------------------- +# Canary vocabulary bias unit test (pure host, no model) +# ----------------------------------------------------------------------------- + +add_executable(transcribe_canary_vocab_boost_unit + canary_vocab_boost_unit.cpp) + +target_link_libraries(transcribe_canary_vocab_boost_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_canary_vocab_boost_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_canary_vocab_boost_unit) + +add_test(NAME transcribe_canary_vocab_boost_unit COMMAND transcribe_canary_vocab_boost_unit) + # ----------------------------------------------------------------------------- # MOSS diarized-transcript parser unit test (pure host, no model) # ----------------------------------------------------------------------------- diff --git a/tests/canary_vocab_boost_unit.cpp b/tests/canary_vocab_boost_unit.cpp new file mode 100644 index 000000000..9d6abba90 --- /dev/null +++ b/tests/canary_vocab_boost_unit.cpp @@ -0,0 +1,94 @@ +// Canary vocabulary bias (pure host, no model). Too strong and a listed term +// replaces unrelated speech; too weak and a term the audio does carry loses to +// a near miss. See src/arch/canary/vocab_boost.h. + +#include "arch/canary/vocab_boost.h" + +#include +#include +#include + +namespace { + +int g_failures = 0; + +void expect(const char * what, bool ok) { + if (!ok) { + std::fprintf(stderr, "FAIL %s\n", what); + ++g_failures; + } +} + +void set_env(const char * key, const char * value) { +#if defined(_WIN32) + _putenv_s(key, value); +#else + setenv(key, value, 1); +#endif +} + +void unset_env(const char * key) { +#if defined(_WIN32) + _putenv_s(key, ""); +#else + unsetenv(key); +#endif +} + +constexpr int k_vocab = 8; + +std::vector logits(int top, float top_score, int other, float other_score) { + std::vector out(k_vocab, -10.0f); + out[top] = top_score; + out[other] = other_score; + return out; +} + +} // namespace + +int main(void) { + using transcribe::canary::VocabBoost; + + VocabBoost boost; + boost.terms = { + { 5, 6 } + }; // a two-token term + boost.strength = 3.0f; + + // A term's first token wins a near miss, not a clear winner. + expect("first token beats a near miss", boost.pick(logits(1, 0.0f, 5, -2.0f).data(), k_vocab, {}) == 5); + expect("first token loses to a clear winner", boost.pick(logits(1, 0.0f, 5, -4.0f).data(), k_vocab, {}) == 1); + + // The next token is boosted only right after the term's previous tokens. + expect("continuation after its prefix", boost.pick(logits(1, 0.0f, 6, -2.0f).data(), k_vocab, { 3, 5 }) == 6); + expect("no continuation elsewhere", boost.pick(logits(1, 0.0f, 6, -2.0f).data(), k_vocab, { 3, 4 }) == 1); + + // Off: plain argmax. + boost.strength = 0.0f; + expect("strength 0 is argmax", boost.pick(logits(1, 0.0f, 5, -0.5f).data(), k_vocab, {}) == 1); + expect("strength 0 is inactive", !boost.active()); + + // Out-of-range term ids are skipped, not read. + VocabBoost bad; + bad.terms = { { k_vocab + 3 } }; + bad.strength = 3.0f; + expect("out-of-range id skipped", bad.pick(logits(2, 0.0f, 4, -1.0f).data(), k_vocab, {}) == 2); + + // Strength from the environment, default when unset or not a number. + unset_env("TRANSCRIBE_CANARY_VOCAB_BOOST"); + expect("default strength", transcribe::canary::vocab_boost_strength() == transcribe::canary::k_vocab_boost_default); + set_env("TRANSCRIBE_CANARY_VOCAB_BOOST", "1.5"); + expect("env strength", transcribe::canary::vocab_boost_strength() == 1.5f); + set_env("TRANSCRIBE_CANARY_VOCAB_BOOST", "0"); + expect("env turns it off", transcribe::canary::vocab_boost_strength() == 0.0f); + set_env("TRANSCRIBE_CANARY_VOCAB_BOOST", "high"); + expect("garbage keeps the default", + transcribe::canary::vocab_boost_strength() == transcribe::canary::k_vocab_boost_default); + + if (g_failures != 0) { + std::fprintf(stderr, "%d failure(s)\n", g_failures); + return 1; + } + std::printf("canary_vocab_boost_unit: ok\n"); + return 0; +}