diff --git a/docs/models/canary-180m-flash.md b/docs/models/canary-180m-flash.md index 30ff3d53..9267ad92 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 b443dcde..4f3f52b3 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 4ed8bd40..71240b84 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 d908e5c2..65648043 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 6683cdee..e37a86a0 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 b0ff3539..49e05948 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 cdaf3ac6..1cf885d1 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 00000000..72fc736e --- /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 110645b0..64e48923 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 00000000..9d6abba9 --- /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; +}