sampling.h
speculative.cpp
speculative.h
+ trie.cpp
+ trie.h
unicode.cpp
unicode.h
jinja/lexer.cpp
data.supports_thinking = true;
data.thinking_start_tag = "[THINK]";
- data.thinking_end_tag = "[/THINK]";
+ data.thinking_end_tags = {"[/THINK]"};
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
+ data.thinking_start_tag = "<|channel|>analysis<|message|>";
+ data.thinking_end_tags = {"<|end|>"};
+
// These special tokens are required to parse properly, so we include them
// even if parse_tool_calls is false.
data.preserved_tokens = {
data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4;
data.supports_thinking = true;
data.thinking_start_tag = "<|channel>thought";
- data.thinking_end_tag = "<channel|>";
+ data.thinking_end_tags = {"<channel|>"};
data.preserved_tokens = {
"<|channel>",
const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>";
data.thinking_start_tag = THINK_START;
- data.thinking_end_tag = THINK_END;
+ data.thinking_end_tags = {THINK_END};
if (inputs.has_continuation()) {
const auto & msg = inputs.continue_msg;
}
data.thinking_start_tag = THINK_START;
- data.thinking_end_tag = THINK_END;
+ data.thinking_end_tags = {THINK_END};
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.thinking_start_tag = "<think>";
- data.thinking_end_tag = "</think>";
+ data.thinking_end_tags = {"</think>"};
data.preserved_tokens = {
"｜DSML｜",
"<think>",
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = true;
data.thinking_start_tag = THINK_START;
- data.thinking_end_tag = THINK_END;
+ data.thinking_end_tags = {THINK_END};
data.preserved_tokens = {
TURN_START, TURN_END, CHATBOT, USER, SYSTEM,
THINK_START, THINK_END,
};
data.thinking_start_tag = "<think>";
- data.thinking_end_tag = "</think>";
+ data.thinking_end_tags = {"</think>"};
data.message_delimiters = {
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE;
if (auto_params.supports_thinking) {
auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start);
- auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end);
+ auto end_tag = trim_whitespace(autoparser.reasoning.end);
+ if (!end_tag.empty()) {
+ auto_params.thinking_end_tags = {std::move(end_tag)};
+ }
}
common_peg_arena arena;
arena.load(auto_params.parser);
std::string generation_prompt;
bool supports_thinking = false;
std::string thinking_start_tag; // e.g., "<think>"
- std::string thinking_end_tag; // e.g., "</think>"
+ std::vector<std::string> thinking_end_tags; // e.g., "</think>"
std::vector<common_grammar_trigger> grammar_triggers;
std::vector<std::string> preserved_tokens;
std::vector<std::string> additional_stops;
// reasoning budget sampler parameters
// these are populated by the server/CLI based on chat template params
- int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
- std::vector<llama_token> reasoning_budget_start; // start tag token sequence
- std::vector<llama_token> reasoning_budget_end; // end tag token sequence
- std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + end tag)
- std::string reasoning_budget_message; // message injected before end tag when budget exhausted
- bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
+ int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
+ std::vector<llama_token> reasoning_budget_start; // start tag token sequence
+ std::vector<llama_tokens> reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence
+ std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + first end tag)
+ std::string reasoning_budget_message; // message injected before end tag when budget exhausted
+ bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
bool backend_sampling = false;
#include "common.h"
#include "json-schema-to-grammar.h"
#include "log.h"
+#include "trie.h"
#include "unicode.h"
#include <algorithm>
-#include <deque>
#include <initializer_list>
#include <map>
#include <memory>
return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F');
}
-// Trie for matching multiple literals.
-// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
-struct trie {
- struct node {
- std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
- bool is_word;
- };
-
- std::vector<node> nodes;
-
- trie(const std::vector<std::string> & words) {
- create_node(); // root node
- for (const auto & w : words) {
- insert(w);
- }
- }
-
- enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
-
- // Check if a delimiter starts at the given position
- match_result check_at(std::string_view sv, size_t start_pos) const {
- size_t current = 0; // Start at root
- size_t pos = start_pos;
-
- // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
-
- while (pos < sv.size()) {
- auto result = common_parse_utf8_codepoint(sv, pos);
- if (result.status != utf8_parse_result::SUCCESS) {
- break;
- }
-
- auto it = nodes[current].children.find(result.codepoint);
- if (it == nodes[current].children.end()) {
- // Can't continue matching
- return match_result{match_result::NO_MATCH};
- }
-
- current = it->second;
- pos += result.bytes_consumed;
-
- // Check if we've matched a complete word
- if (nodes[current].is_word) {
- return match_result{match_result::COMPLETE_MATCH};
- }
- }
-
- // Reached end of input while still in the trie (not at root)
- if (current != 0) {
- // We're in the middle of a potential match
- return match_result{match_result::PARTIAL_MATCH};
- }
-
- // Reached end at root (no match)
- return match_result{match_result::NO_MATCH};
- }
-
- private:
- size_t create_node() {
- size_t index = nodes.size();
- nodes.emplace_back();
- return index;
- }
-
- void insert(const std::string & word) {
- size_t current = 0;
- size_t pos = 0;
- while (pos < word.length()) {
- auto result = common_parse_utf8_codepoint(word, pos);
- if (result.status != utf8_parse_result::SUCCESS) {
- break;
- }
-
- uint32_t ch = result.codepoint;
- pos += result.bytes_consumed;
-
- auto it = nodes[current].children.find(ch);
- if (it == nodes[current].children.end()) {
- size_t child = create_node();
- nodes[current].children[ch] = child;
- current = child;
- } else {
- current = it->second;
- }
- }
- nodes[current].is_word = true;
- }
-};
-
-// Aho-Corasick automaton
-struct aho_corasick {
- trie t;
- std::vector<size_t> fail; // failure links
- std::vector<size_t> order; // states in BFS order
- std::vector<bool> terminal; // match states (directly or via a suffix link)
- std::set<uint32_t> alphabet; // every character with a transition
-
- aho_corasick(const std::vector<std::string> & strings) : t(strings) {
- const auto & nodes = t.nodes;
- const size_t n = nodes.size();
-
- fail.assign(n, 0);
- order.reserve(n);
-
- std::deque<size_t> queue{ 0 };
- while (!queue.empty()) {
- size_t u = queue.front();
- queue.pop_front();
- order.push_back(u);
- for (const auto & [ch, v] : nodes[u].children) {
- if (u != 0) {
- size_t f = fail[u];
- while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
- f = fail[f];
- }
- auto it = nodes[f].children.find(ch);
- fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
- }
- queue.push_back(v);
- }
- }
-
- terminal.assign(n, false);
- for (size_t u : order) {
- terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]);
- }
-
- for (const auto & node : nodes) {
- for (const auto & [ch, v] : node.children) {
- alphabet.insert(ch);
- }
- }
- }
-
- size_t num_states() const { return t.nodes.size(); }
- bool is_terminal(size_t s) const { return terminal[s]; }
-
- // follow failure links until a transition on `ch` exists.
- size_t next(size_t state, uint32_t ch) const {
- const auto & nodes = t.nodes;
- while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
- state = fail[state];
- }
- auto it = nodes[state].children.find(ch);
- return it != nodes[state].children.end() ? it->second : 0;
- }
-};
-
static std::pair<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) {
if (pos + hex_count > str.length()) {
return {0, 0};
}
common_peg_parse_result operator()(const common_peg_until_parser & p) const {
- trie matcher(p.delimiters);
+ common_trie matcher(p.delimiters);
// Scan input and check for delimiters
size_t pos = start_pos;
// Check if a delimiter starts at this position
auto match = matcher.check_at(ctx.input, pos);
- if (match == trie::COMPLETE_MATCH) {
+ if (match == common_trie::COMPLETE_MATCH) {
// Found a complete delimiter, return everything before it
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
}
- if (match == trie::PARTIAL_MATCH) {
+ if (match == common_trie::PARTIAL_MATCH) {
// Found a partial match extending to end of input, return everything before it
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
}
const std::map<size_t, std::vector<uint32_t>> &,
const std::vector<uint32_t> &,
const std::function<std::string(size_t)> &)> & build_rule) {
- aho_corasick ac(strings);
+ common_aho_corasick ac(strings);
auto state_name = [&](size_t s) -> std::string {
if (s == 0) {
#include "reasoning-budget.h"
#include "common.h"
+#include "trie.h"
#include "unicode.h"
#include "log.h"
+#include <algorithm>
#include <cmath>
#include <cstdint>
#include <string>
#include <vector>
struct token_matcher {
- std::vector<llama_token> tokens;
- size_t pos = 0;
+ std::vector<llama_tokens> seqs;
+ common_aho_corasick ac;
+ size_t state = 0;
- bool advance(llama_token token) {
- if (tokens.empty()) {
- return false;
- }
+ token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
- if (token == tokens[pos]) {
- pos++;
- if (pos >= tokens.size()) {
- pos = 0;
- return true;
- }
- } else {
- pos = 0;
- if (token == tokens[0]) {
- pos = 1;
+ static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
+ std::vector<llama_tokens> res;
+ for (const auto & seq : seqs) {
+ if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
+ res.push_back(seq);
}
}
- return false;
+ return res;
+ }
+
+ static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
+ common_trie t;
+ for (const auto & seq : seqs) {
+ t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
+ }
+ return t;
}
- void reset() { pos = 0; }
+ // returns the index into seqs of the longest sequence ending at this token, or -1
+ int32_t advance(llama_token token) {
+ state = ac.next(state, (uint32_t) token);
+ const int32_t p = ac.match_pattern(state);
+ if (p >= 0) {
+ state = 0;
+ }
+ return p;
+ }
+
+ void reset() { state = 0; }
};
struct common_reasoning_budget_ctx {
token_matcher start_matcher;
token_matcher end_matcher;
- std::vector<llama_token> forced_tokens;
+ llama_tokens forced_tokens;
int32_t budget; // maximum tokens in reasoning block
int32_t remaining; // tokens remaining in budget
// for forcing
size_t force_pos; // next position in forced_tokens to force
+
+ int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none
};
static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {
switch (ctx->state) {
case REASONING_BUDGET_IDLE:
{
- if (ctx->start_matcher.advance(token)) {
+ if (ctx->start_matcher.advance(token) >= 0) {
ctx->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget;
COM_TRC("activated, budget=%d tokens\n", ctx->budget);
case REASONING_BUDGET_COUNTING:
case REASONING_BUDGET_WAITING_UTF8:
{
- if (ctx->end_matcher.advance(token)) {
+ const int32_t match = ctx->end_matcher.advance(token);
+ if (match >= 0) {
ctx->state = REASONING_BUDGET_DONE;
+ ctx->end_match = match;
COM_TRC("%s", "deactivated (natural end)\n");
break;
}
break;
}
case REASONING_BUDGET_FORCING:
+ {
+ // track the end sequence within forced_tokens so it is also reported on DONE
+ const int32_t match = ctx->end_matcher.advance(token);
ctx->force_pos++;
if (ctx->force_pos >= ctx->forced_tokens.size()) {
ctx->state = REASONING_BUDGET_DONE;
+ ctx->end_match = match;
COM_TRC("%s", "forced sequence complete, done\n");
}
break;
+ }
case REASONING_BUDGET_DONE:
// Re-arm on a new start tag: some models emit multiple <think> blocks
// per response, and each should get a fresh budget window.
- if (ctx->start_matcher.advance(token)) {
+ if (ctx->start_matcher.advance(token) >= 0) {
ctx->state = REASONING_BUDGET_COUNTING;
ctx->remaining = ctx->budget;
ctx->end_matcher.reset();
+ ctx->end_match = -1;
COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
if (ctx->remaining <= 0) {
ctx->start_matcher.reset();
ctx->end_matcher.reset();
ctx->force_pos = 0;
+ ctx->end_match = -1;
}
static struct llama_sampler * common_reasoning_budget_init_state(
- const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens,
- const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens,
+ const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
+ const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
int32_t budget, common_reasoning_budget_state initial_state);
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
}
static struct llama_sampler * common_reasoning_budget_init_state(
- const struct llama_vocab * vocab,
- const std::vector<llama_token> & start_tokens,
- const std::vector<llama_token> & end_tokens,
- const std::vector<llama_token> & forced_tokens,
- int32_t budget,
- common_reasoning_budget_state initial_state) {
+ const struct llama_vocab * vocab,
+ const std::vector<llama_tokens> & start_seqs,
+ const std::vector<llama_tokens> & end_seqs,
+ const llama_tokens & forced_tokens,
+ int32_t budget,
+ common_reasoning_budget_state initial_state) {
// promote COUNTING with budget <= 0 to FORCING
if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
initial_state = REASONING_BUDGET_FORCING;
/* .iface = */ &common_reasoning_budget_i,
/* .ctx = */ new common_reasoning_budget_ctx {
/* .vocab = */ vocab,
- /* .start_matcher = */ { start_tokens, 0 },
- /* .end_matcher = */ { end_tokens, 0 },
+ /* .start_matcher = */ token_matcher(start_seqs),
+ /* .end_matcher = */ token_matcher(end_seqs),
/* .forced_tokens = */ forced_tokens,
/* .budget = */ budget,
/* .remaining = */ budget,
/* .state = */ initial_state,
/* .force_pos = */ 0,
+ /* .end_match = */ -1,
}
);
}
struct llama_sampler * common_reasoning_budget_init(
- const struct llama_vocab * vocab,
- const std::vector<llama_token> & start_tokens,
- const std::vector<llama_token> & end_tokens,
- const std::vector<llama_token> & forced_tokens,
- int32_t budget,
- common_reasoning_budget_state initial_state) {
- return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state);
+ const struct llama_vocab * vocab,
+ const std::vector<llama_tokens> & start_seqs,
+ const std::vector<llama_tokens> & end_seqs,
+ const llama_tokens & forced_tokens,
+ int32_t budget,
+ common_reasoning_budget_state initial_state) {
+ return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);
}
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
}
+const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {
+ if (!smpl) {
+ return nullptr;
+ }
+
+ const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
+ if (ctx->end_match < 0) {
+ return nullptr;
+ }
+
+ return &ctx->end_matcher.seqs[ctx->end_match];
+}
+
bool common_reasoning_budget_force(struct llama_sampler * smpl) {
if (!smpl) {
return false;
#include "llama.h"
+#include "common.h"
+
#include <cstdint>
#include <vector>
// reasoning block (e.g. between <think> and </think>).
//
// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
-// IDLE: passthrough, watching for start_tokens sequence
-// COUNTING: counting down remaining tokens, watching for natural end_tokens
+// IDLE: passthrough, watching for a start sequence
+// COUNTING: counting down remaining tokens, watching for a natural end sequence
// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence
// FORCING: forces forced_tokens token-by-token (all other logits -> -inf)
// DONE: passthrough forever
//
// Parameters:
// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr)
-// start_tokens - token sequence that activates counting
-// end_tokens - token sequence for natural deactivation
+// start_seqs - token sequences, any of which activates counting
+// end_seqs - token sequences, any of which naturally deactivates
// forced_tokens - token sequence forced when budget expires
// budget - max tokens allowed in the reasoning block
// initial_state - initial state
//
struct llama_sampler * common_reasoning_budget_init(
- const struct llama_vocab * vocab,
- const std::vector<llama_token> & start_tokens,
- const std::vector<llama_token> & end_tokens,
- const std::vector<llama_token> & forced_tokens,
- int32_t budget,
- common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
+ const struct llama_vocab * vocab,
+ const std::vector<llama_tokens> & start_seqs,
+ const std::vector<llama_tokens> & end_seqs,
+ const llama_tokens & forced_tokens,
+ int32_t budget,
+ common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);
+// The end sequence that transitioned the sampler to DONE, or nullptr if none
+// was recorded. Cleared when a new start sequence re-arms the sampler.
+const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl);
+
// Manually transition the reasoning budget sampler into the FORCING state.
// Returns true if the transition occurred.
bool common_reasoning_budget_force(struct llama_sampler * smpl);
if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) {
rbudget = common_reasoning_budget_init(
vocab,
- params.reasoning_budget_start,
+ {params.reasoning_budget_start},
params.reasoning_budget_end,
params.reasoning_budget_forced,
params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens);
if (gsmpl->rbudget && is_generated) {
llama_sampler_accept(gsmpl->rbudget, token);
+
+ // if done, replay end sequence which may contain a grammar trigger
+ const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE;
+ if (gsmpl->grmr && !accept_grammar && is_done) {
+ const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget);
+ if (end_seq) {
+ for (const llama_token end_token : *end_seq) {
+ llama_sampler_accept(gsmpl->grmr, end_token);
+ }
+ }
+ }
}
if (gsmpl->grmr && accept_grammar) {
--- /dev/null
+#include "trie.h"
+
+#include "unicode.h"
+
+#include <deque>
+
+common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const {
+ size_t current = 0; // Start at root
+ size_t pos = start_pos;
+
+ // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
+
+ while (pos < sv.size()) {
+ auto result = common_parse_utf8_codepoint(sv, pos);
+ if (result.status != utf8_parse_result::SUCCESS) {
+ break;
+ }
+
+ auto it = nodes[current].children.find(result.codepoint);
+ if (it == nodes[current].children.end()) {
+ // Can't continue matching
+ return match_result{match_result::NO_MATCH};
+ }
+
+ current = it->second;
+ pos += result.bytes_consumed;
+
+ // Check if we've matched a complete word
+ if (nodes[current].pattern >= 0) {
+ return match_result{match_result::COMPLETE_MATCH};
+ }
+ }
+
+ // Reached end of input while still in the trie (not at root)
+ if (current != 0) {
+ // We're in the middle of a potential match
+ return match_result{match_result::PARTIAL_MATCH};
+ }
+
+ // Reached end at root (no match)
+ return match_result{match_result::NO_MATCH};
+}
+
+int32_t common_trie::insert(const std::string & word) {
+ std::vector<uint32_t> symbols;
+ size_t pos = 0;
+ while (pos < word.length()) {
+ auto result = common_parse_utf8_codepoint(word, pos);
+ if (result.status != utf8_parse_result::SUCCESS) {
+ break;
+ }
+
+ symbols.push_back(result.codepoint);
+ pos += result.bytes_consumed;
+ }
+ return insert(symbols);
+}
+
+int32_t common_trie::insert(const std::vector<uint32_t> & symbols) {
+ size_t current = 0;
+ for (uint32_t ch : symbols) {
+ auto it = nodes[current].children.find(ch);
+ if (it == nodes[current].children.end()) {
+ size_t child = create_node();
+ nodes[current].children[ch] = child;
+ current = child;
+ } else {
+ current = it->second;
+ }
+ }
+ if (nodes[current].pattern < 0) {
+ nodes[current].pattern = n_patterns++;
+ }
+ return nodes[current].pattern;
+}
+
+common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) {
+ const auto & nodes = t.nodes;
+ const size_t n = nodes.size();
+
+ fail.assign(n, 0);
+ order.reserve(n);
+
+ std::deque<size_t> queue{ 0 };
+ while (!queue.empty()) {
+ size_t u = queue.front();
+ queue.pop_front();
+ order.push_back(u);
+ for (const auto & [ch, v] : nodes[u].children) {
+ if (u != 0) {
+ size_t f = fail[u];
+ while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
+ f = fail[f];
+ }
+ auto it = nodes[f].children.find(ch);
+ fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
+ }
+ queue.push_back(v);
+ }
+ }
+
+ // fail[u] points to a strictly shorter suffix, so the first pattern found on
+ // the fail chain (including u itself) is the longest pattern ending at u
+ match.assign(n, -1);
+ for (size_t u : order) {
+ match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1);
+ }
+
+ for (const auto & node : nodes) {
+ for (const auto & [ch, v] : node.children) {
+ alphabet.insert(ch);
+ }
+ }
+}
+
+size_t common_aho_corasick::next(size_t state, uint32_t ch) const {
+ const auto & nodes = t.nodes;
+ while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
+ state = fail[state];
+ }
+ auto it = nodes[state].children.find(ch);
+ return it != nodes[state].children.end() ? it->second : 0;
+}
--- /dev/null
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <set>
+#include <string>
+#include <string_view>
+#include <vector>
+
+// Trie for matching multiple literals.
+// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
+struct common_trie {
+ struct node {
+ std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
+ int32_t pattern = -1; // index of the pattern ending at this node, -1 if none
+ };
+
+ std::vector<node> nodes;
+
+ common_trie() {
+ create_node(); // root node
+ }
+
+ common_trie(const std::vector<std::string> & words) : common_trie() {
+ for (const auto & w : words) {
+ insert(w);
+ }
+ }
+
+ enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
+
+ // Check if a delimiter starts at the given position
+ match_result check_at(std::string_view sv, size_t start_pos) const;
+
+ // Insert a word as a sequence of Unicode codepoints, returns its pattern index
+ int32_t insert(const std::string & word);
+
+ // Insert a raw symbol sequence, returns its pattern index (insertion order,
+ // duplicates keep the first index)
+ int32_t insert(const std::vector<uint32_t> & symbols);
+
+ private:
+ int32_t n_patterns = 0;
+
+ size_t create_node() {
+ size_t index = nodes.size();
+ nodes.emplace_back();
+ return index;
+ }
+};
+
+// Aho-Corasick automaton
+struct common_aho_corasick {
+ common_trie t;
+ std::vector<size_t> fail; // failure links
+ std::vector<size_t> order; // states in BFS order
+ std::vector<int32_t> match; // longest pattern ending at each state (directly or via a suffix link), -1 if none
+ std::set<uint32_t> alphabet; // every character with a transition
+
+ common_aho_corasick(common_trie trie);
+
+ common_aho_corasick(const std::vector<std::string> & strings)
+ : common_aho_corasick(common_trie(strings)) {}
+
+ size_t num_states() const { return t.nodes.size(); }
+ bool is_terminal(size_t s) const { return match[s] >= 0; }
+
+ // index of the longest pattern ending at this state, -1 if none
+ int32_t match_pattern(size_t s) const { return match[s]; }
+
+ // follow failure links until a transition on `ch` exists.
+ size_t next(size_t state, uint32_t ch) const;
+};
// budget sampler inhibits grammar application while inside thinking blocks —
// triggers inside <think>...</think> are suppressed.
bool use_reasoning_budget_path = false;
- if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tag.empty()) {
+ if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tags.empty()) {
use_reasoning_budget_path = true;
for (const auto & trigger : parser.params_.grammar_triggers) {
if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) {
// Walk through full_input tracking thinking state; only match triggers
// when outside thinking blocks.
const auto & think_start = parser.params_.thinking_start_tag;
- const auto & think_end = parser.params_.thinking_end_tag;
+ const auto & think_ends = parser.params_.thinking_end_tags;
bool in_thinking = false;
for (size_t i = 0; i < full_input.size(); ++i) {
i += think_start.size() - 1;
continue;
}
- if (in_thinking && full_input.compare(i, think_end.size(), think_end) == 0) {
- in_thinking = false;
- i += think_end.size() - 1;
- continue;
- }
if (in_thinking) {
+ for (const auto & think_end : think_ends) {
+ if (full_input.compare(i, think_end.size(), think_end) == 0) {
+ in_thinking = false;
+ i += think_end.size() - 1;
+ break;
+ }
+ }
continue;
}
// Outside thinking — check if any trigger word starts here
static void test_reasoning_budget(
const char * test_name,
const std::vector<llama_token> & sequence,
- const std::vector<llama_token> & start_tokens,
- const std::vector<llama_token> & end_tokens,
+ const std::vector<llama_tokens> & start_seqs,
+ const std::vector<llama_tokens> & end_seqs,
const std::vector<llama_token> & forced_tokens,
int32_t budget,
common_reasoning_budget_state initial_state,
// Find the maximum token ID to ensure our vocab covers all tokens
llama_token max_token = 0;
for (auto t : sequence) max_token = std::max(max_token, t);
- for (auto t : start_tokens) max_token = std::max(max_token, t);
- for (auto t : end_tokens) max_token = std::max(max_token, t);
+ for (const auto & seq : start_seqs) {
+ for (auto t : seq) max_token = std::max(max_token, t);
+ }
+ for (const auto & seq : end_seqs) {
+ for (auto t : seq) max_token = std::max(max_token, t);
+ }
for (auto t : forced_tokens) max_token = std::max(max_token, t);
// Create a minimal sampler with mock vocabulary
// The UTF-8 boundary check will treat all tokens as complete (safe fallback)
auto * sampler = common_reasoning_budget_init(
nullptr, // vocab - not used for basic state machine tests
- start_tokens,
- end_tokens,
+ start_seqs,
+ end_seqs,
forced_tokens,
budget,
initial_state
const std::vector<llama_token> end = {101};
const std::vector<llama_token> forced = {102, 101};
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 2, REASONING_BUDGET_IDLE);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE);
llama_sampler_accept(sampler, 100); // COUNTING, remaining=2
llama_sampler_accept(sampler, 50); // COUNTING, remaining=1
const std::vector<llama_token> end = {101};
const std::vector<llama_token> forced = {102, 101};
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING);
GGML_ASSERT(get_forced_token(sampler, 102) == 102);
llama_sampler_accept(sampler, 102); // advance to the second forced token
// if COUNTING, force() succeeds and begins forcing the end sequence from the start
{
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);
llama_sampler_accept(sampler, 100); // COUNTING, remaining=5
llama_sampler_accept(sampler, 50); // COUNTING, remaining=4
// if IDLE, force() is a no-op
{
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);
GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE");
GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE);
// if DONE, force() is a no-op
{
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);
llama_sampler_accept(sampler, 100); // COUNTING
llama_sampler_accept(sampler, 101); // natural end -> DONE
// if FORCING, force() is a no-op and must not rewind the force position
{
- auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING);
+ auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING);
GGML_ASSERT(get_forced_token(sampler, 102) == 102);
llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1)
fprintf(stderr, " Test 'manual force transition' passed\n");
}
+static void test_reasoning_budget_end_match() {
+ const std::vector<llama_tokens> start = {{100}};
+ const std::vector<llama_tokens> end = {{101}, {103, 104}};
+
+ // natural end records the sequence that matched; re-arming clears it
+ {
+ auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE);
+
+ GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
+
+ llama_sampler_accept(sampler, 100); // COUNTING
+ llama_sampler_accept(sampler, 50);
+ llama_sampler_accept(sampler, 103);
+ llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE
+
+ const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
+ GGML_ASSERT(matched != nullptr);
+ GGML_ASSERT(*matched == llama_tokens({103, 104}));
+
+ llama_sampler_accept(sampler, 100); // re-arm, COUNTING
+ GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
+
+ llama_sampler_free(sampler);
+ }
+
+ // overlapping end sequences: the longest one ending at the position wins
+ {
+ const std::vector<llama_tokens> end_overlap = {{104}, {103, 104}};
+
+ auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE);
+
+ llama_sampler_accept(sampler, 100); // COUNTING
+ llama_sampler_accept(sampler, 103);
+ llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here
+
+ const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
+ GGML_ASSERT(matched != nullptr);
+ GGML_ASSERT(*matched == llama_tokens({103, 104}));
+
+ llama_sampler_free(sampler);
+ }
+
+ // forcing records the end sequence terminating forced_tokens
+ {
+ auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING);
+
+ llama_sampler_accept(sampler, 102);
+ llama_sampler_accept(sampler, 103);
+ GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
+ llama_sampler_accept(sampler, 104); // forced sequence complete, DONE
+
+ const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
+ GGML_ASSERT(matched != nullptr);
+ GGML_ASSERT(*matched == llama_tokens({103, 104}));
+
+ llama_sampler_free(sampler);
+ }
+
+ // forced_tokens not ending with a known end sequence records nothing
+ {
+ auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING);
+
+ llama_sampler_accept(sampler, 102); // forced sequence complete, DONE
+ GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);
+ GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);
+
+ llama_sampler_free(sampler);
+ }
+
+ // a null sampler is safely ignored
+ GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr);
+
+ fprintf(stderr, " Test 'matched end sequence' passed\n");
+}
+
// UTF-8 boundary detection unit test
// Tests common_utf8_is_complete() from reasoning-budget.h
static void test_utf8_boundary_detection() {
const std::vector<llama_token> forced = {102}; // forced token (not used in this test)
const std::vector<llama_token> sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more
- test_reasoning_budget("natural end before budget exhausted", sequence, start, end, forced,
+ test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced,
5, // budget of 5 tokens
REASONING_BUDGET_IDLE,
SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)
const std::vector<llama_token> forced = {102, 101}; // forced message + end
const std::vector<llama_token> sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2)
- test_reasoning_budget("budget exhausted forcing", sequence, start, end, forced,
+ test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced,
2, // budget of 2 tokens
REASONING_BUDGET_IDLE,
3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces)
const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {100, 50, 51, 52}; // start token first, then 3 tokens
- test_reasoning_budget("activate immediately budget=0", sequence, start, end, forced,
+ test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced,
0, // budget of 0 tokens
REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0
0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately)
const std::vector<llama_token> forced = {102};
const std::vector<llama_token> sequence = {50, 51, 52, 53};
- test_reasoning_budget("no start/end configured", sequence, start, end, forced,
+ test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced,
2, // budget
REASONING_BUDGET_IDLE,
SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured)
const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {50, 51, 52, 53};
- test_reasoning_budget("activate immediately with budget", sequence, start, end, forced,
+ test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced,
2, // budget of 2 tokens
REASONING_BUDGET_COUNTING,
2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces)
const std::vector<llama_token> forced = {102, 101};
const std::vector<llama_token> sequence = {100, 50, 101, 100, 60, 61, 62, 63};
- test_reasoning_budget("multi-block re-arms budget after DONE", sequence, start, end, forced,
+ test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced,
2, // budget of 2 tokens (per block)
REASONING_BUDGET_IDLE,
6, // forcing starts at i=6 (after second block exhausts at i=5)
7); // forcing continues through i=7
}
+ // Test 7: Multiple start sequences - the second sequence activates counting
+ // Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1;
+ // i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence
+ {
+ const std::vector<llama_tokens> start = {{100}, {110, 111}};
+ const std::vector<llama_tokens> end = {{101}};
+ const std::vector<llama_token> forced = {102, 101};
+ const std::vector<llama_token> sequence = {110, 111, 50, 51, 52, 53};
+
+ test_reasoning_budget("multiple start sequences", sequence, start, end, forced,
+ 2, // budget of 2 tokens
+ REASONING_BUDGET_IDLE,
+ 4, // forcing starts at i=4 (accept at i=3 depletes budget)
+ 5); // forcing continues through i=5
+ }
+
+ // Test 8: Multiple end sequences - natural end via the second sequence
+ // Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4;
+ // i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE
+ {
+ const std::vector<llama_tokens> start = {{100}};
+ const std::vector<llama_tokens> end = {{101}, {103, 104}};
+ const std::vector<llama_token> forced = {102, 101};
+ const std::vector<llama_token> sequence = {100, 50, 103, 104, 52};
+
+ test_reasoning_budget("multiple end sequences", sequence, start, end, forced,
+ 5, // budget of 5 tokens
+ REASONING_BUDGET_IDLE,
+ SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)
+ }
+
test_reasoning_budget_clone_mid_counting();
test_reasoning_budget_clone_mid_forcing();
test_reasoning_budget_force_manual();
+ test_reasoning_budget_end_match();
- printf("OK (9 tests passed)\n");
+ printf("OK (12 tests passed)\n");
printf("Testing UTF-8 boundary detection... ");
test_utf8_boundary_detection();
reasoning_budget = opt.reasoning_budget;
}
- if (!chat_params.thinking_end_tag.empty()) {
+ if (!chat_params.thinking_end_tags.empty()) {
llama_params["reasoning_budget_tokens"] = reasoning_budget;
llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag;
- llama_params["reasoning_budget_end_tag"] = chat_params.thinking_end_tag;
+ llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags;
llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message);
llama_params["reasoning_control"] = json_value(body, "reasoning_control", false);
}
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"))
// 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(),
overview / pkg / ggml / sources / llama.cpp / commitdiff