canary: bias the greedy decode toward vocabulary terms

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.

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.

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.

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 | | | |
 

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);

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;
 };

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 <algorithm>
@@ -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,13 +992,41 @@ transcribe_status run(transcribe_session *          session,
 
     std::vector<int32_t> prompt_ids;
     std::vector<int32_t> 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);
             st != TRANSCRIBE_OK) {
             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<char>(c - 'a' + 'A');
+                        } else if (c >= 'A' && c <= 'Z') {
+                            c = static_cast<char>(c - 'A' + 'a');
+                        } else {
+                            break;
+                        }
+                    }
+                    std::vector<int32_t> 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,21 +1224,16 @@ transcribe_status run(transcribe_session *          session,
             std::vector<float> logits_host(static_cast<size_t>(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<size_t>(prompt_len - 1) * vocab_size;
-            float         best        = last_logits[0];
-            for (int j = 1; j < static_cast<int>(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<int> generated_ids;
         if (next_token != eos_id) {
             generated_ids.push_back(next_token);
         }
 
-        int n_past = prompt_len;
+        int                n_past = prompt_len;
+        std::vector<float> 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<size_t>(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<size_t>(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);
     }
 

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 <cstdint>
+#include <cstdlib>
+#include <vector>
+
+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<std::vector<int32_t>> 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<int> & 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<int>(j);
+            }
+        }
+        if (!active()) {
+            return best_id;
+        }
+        for (const std::vector<int32_t> & 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<int> & generated, const std::vector<int32_t> & 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

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)
 # -----------------------------------------------------------------------------