Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
common/CMakeLists.txt
+2
View File
@@ -100,6 +100,8 @@ add_library(${TARGET}
sampling.h
speculative.cpp
speculative.h
trie.cpp
trie.h
unicode.cpp
unicode.h
jinja/lexer.cpp
common/chat.cpp
+14 -8
View File
@@ -1024,7 +1024,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
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;
@@ -1150,6 +1150,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
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 = {
@@ -1294,7 +1297,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ
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>",
@@ -1569,7 +1572,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp
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;
@@ -1703,7 +1706,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat
}
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();
@@ -1943,7 +1946,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
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>",
@@ -2160,7 +2163,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
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,
@@ -2510,7 +2513,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
};
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" },
@@ -2866,7 +2869,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
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);
common/chat.h
+1 -1
View File
@@ -274,7 +274,7 @@ struct common_chat_params {
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;
common/common.h
+6 -6
View File
@@ -284,12 +284,12 @@ struct common_params_sampling {
// 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;
common/peg-parser.cpp
+5 -153
View File
@@ -3,10 +3,10 @@
#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>
@@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) {
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};
@@ -797,7 +649,7 @@ struct parser_executor {
}
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;
@@ -824,12 +676,12 @@ struct parser_executor {
// 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);
}
@@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar(
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) {
common/reasoning-budget.cpp
+77 -39
View File
@@ -1,39 +1,52 @@
#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;
}
void reset() { pos = 0; }
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;
}
// 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 {
@@ -41,7 +54,7 @@ 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
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
// 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*/) {
@@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
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);
@@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
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;
}
@@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
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) {
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
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);
@@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s
}
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;
@@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state(
/* .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) {
@@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla
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;
common/reasoning-budget.h
+16 -10
View File
@@ -2,6 +2,8 @@
#include "llama.h"
#include "common.h"
#include <cstdint>
#include <vector>
@@ -17,30 +19,34 @@ enum common_reasoning_budget_state {
// 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);
common/sampling.cpp
+12 -1
View File
@@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
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);
@@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo
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) {
common/trie.cpp
+123
View File
@@ -0,0 +1,123 @@
#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;
}
common/trie.h
+73
View File
@@ -0,0 +1,73 @@
#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;
};