common : add support for multiple end sequences in the reasoning budget sampler (#25544)

* common : extract trie/ac to a separate file

* common : support multiple token sequences in the reasoning budget sampler

* common/trie : return matched word index

* common/trie : rename "word" to "pattern"

* common/reasoning-budget : expose matched end sequence

* common/sampling : replay end sequence when reasoning budget is done

* cont : update to use multiple end sequences

* cont : clean up
This commit is contained in:
Aldehir Rojas
2026-07-25 11:58:09 +02:00
committed by GitHub
parent d67c0b4107
commit 910196f6b3
14 changed files with 497 additions and 254 deletions
+27 -8
View File
@@ -390,21 +390,40 @@ std::vector<std::unique_ptr<field>> make_llama_cmpl_schema(const common_params &
ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get<std::string>(), false, true);
}));
add((new field_str("reasoning_budget_end_tag"))
->set_desc("Token string marking the end of the reasoning budget section")
add((new field_json("reasoning_budget_end_tags"))
->add_alias("reasoning_budget_end_tag")
->set_desc("Token strings marking the end of the reasoning budget section; the first is forced when the budget expires")
->set_handler([&](field_eval_context & ctx, const json & data) {
GGML_ASSERT(ctx.vocab != nullptr);
std::string end_tag = data.at("reasoning_budget_end_tag").get<std::string>();
ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true);
ctx.params.sampling.reasoning_budget_end.clear();
if (data.contains("reasoning_budget_end_tags")) {
for (const auto & t : data.at("reasoning_budget_end_tags")) {
std::string tag = t.get<std::string>();
if (!tag.empty()) {
ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true));
}
}
} else if (data.contains("reasoning_budget_end_tag")) {
std::string tag = data.at("reasoning_budget_end_tag").get<std::string>();
if (!tag.empty()) {
ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true));
}
}
}));
add((new field_str("reasoning_budget_message"))
->set_desc("Message to prepend to the reasoning budget end tag when forcing it")
->set_handler([&](field_eval_context & ctx, const json & data) {
GGML_ASSERT(ctx.vocab != nullptr);
std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string());
std::string message = data.at("reasoning_budget_message").get<std::string>();
ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true);
if (!ctx.params.sampling.reasoning_budget_end.empty()) {
llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front();
std::string message = json_value(data, "reasoning_budget_message", std::string());
if (!message.empty()) {
llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true);
end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end());
}
ctx.params.sampling.reasoning_budget_forced = std::move(end_tag);
}
}));
add((new field_json("logit_bias"))
@@ -546,7 +565,7 @@ task_params eval_llama_cmpl_schema(
// debugging
{
auto budget = params.sampling.reasoning_budget_tokens;
SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu toks, forced=%zu toks\n",
SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n",
budget, params.sampling.generation_prompt.c_str(),
params.sampling.reasoning_budget_start.size(),
params.sampling.reasoning_budget_end.size(),