diff --git a/common/common.h b/common/common.h index e7c55ae925..14889193d9 100644 --- a/common/common.h +++ b/common/common.h @@ -294,10 +294,6 @@ struct common_params_sampling { bool backend_sampling = false; - bool has_logit_bias() const { - return !logit_bias.empty(); - } - // print the parameters into a string std::string print() const; }; diff --git a/common/sampling.cpp b/common/sampling.cpp index 7b241e34f7..256ac161e2 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -310,8 +310,19 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st } } - if (params.has_logit_bias()) { - samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), params.logit_bias.size(), params.logit_bias.data())); + // logit bias: user biases + model suppress tokens (-INFINITY) + { + std::vector merged = params.logit_bias; + + int32_t n_suppress = 0; + const llama_token * suppress = llama_vocab_get_suppress_tokens(vocab, &n_suppress); + for (int32_t i = 0; i < n_suppress; ++i) { + merged.push_back({ suppress[i], -INFINITY }); + } + + if (!merged.empty()) { + samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), merged.size(), merged.data())); + } } if (params.mirostat == 0) { diff --git a/include/llama.h b/include/llama.h index 3c6d22be89..a3ceecf5ef 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1102,6 +1102,9 @@ extern "C" { LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab); LLAMA_API bool llama_vocab_get_add_sep(const struct llama_vocab * vocab); + // model-specific suppress tokens (gguf key: tokenizer.ggml.suppress_tokens) + LLAMA_API const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens); + LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab); LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab); LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab); diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9164a4dd88..443cd46408 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2578,7 +2578,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { if (suppress_idx != -1) { const int n = gguf_get_arr_n(ctx, suppress_idx); const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx); - suppress_tokens.assign(data, data + n); + // drop out-of-range ids + suppress_tokens.reserve(n); + for (int i = 0; i < n; ++i) { + const int32_t id = data[i]; + if (id >= 0 && id < (int) id_to_token.size()) { + suppress_tokens.push_back(id); + } + } } } @@ -4205,6 +4212,14 @@ bool llama_vocab_get_add_sep(const struct llama_vocab * vocab) { return vocab->get_add_sep(); } +const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens) { + const std::vector & tokens = vocab->get_suppress_tokens(); + if (n_suppress_tokens) { + *n_suppress_tokens = (int32_t) tokens.size(); + } + return tokens.data(); +} + llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab) { return vocab->token_fim_pre(); } diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index 6a96979ceb..e44f423bdb 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -142,33 +142,6 @@ static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, in idx * x->ne[0] * x->ne[1] * ggml_element_size(x)); } -// TODO @ngxson : maybe improve this in the future -class llm_graph_input_logits_bias : public llm_graph_input_i { -public: - llm_graph_input_logits_bias(const llama_vocab & vocab) { - arr.resize(vocab.n_tokens(), 0.0f); - for (llama_token id : vocab.get_suppress_tokens()) { - if (0 <= id && id < (int32_t)vocab.n_tokens()) { - arr[id] = -INFINITY; - } - } - } - virtual ~llm_graph_input_logits_bias() = default; - - void set_input(const llama_ubatch * /*ubatch*/) override { - const int64_t n_vocab = arr.size(); - ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias)); - } - - bool can_reuse(const llm_graph_params & /*params*/) override { - return true; - } - - ggml_tensor * logits_bias = nullptr; // F32 [n_vocab] - - std::vector arr; -}; - llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model), @@ -429,16 +402,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); } - // apply logits bias if needed (e.g. for gemma4_unified patch) - // this is to mirror the suppress_tokens patch on transformers, to avoid model from outputing and tokens (which is a known issue related to the checkpoint) - // TODO: maybe handle this inside the sampling system in the future - if (!model.vocab.get_suppress_tokens().empty()) { - auto inp_bias = std::make_unique(model.vocab); - inp_bias->logits_bias = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, inp_bias->arr.size()); - cur = ggml_add(ctx0, cur, inp_bias->logits_bias); - res->add_input(std::move(inp_bias)); - } - cb(cur, "result_output", -1); res->t_logits = cur;