char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Clint Herrongrammar : coalesce find + insert into a single insert and adjust move/copy mechanics (#26885)7a16a6ce3

char/aliceai-t5-moe
55.0 KiB1526 linesraw
1#include "llama-grammar.h"
2
3#include "llama-impl.h"
4#include "llama-vocab.h"
5#include "llama-sampler.h"
6
7#include <cmath>
8#include <algorithm>
9#include <cstdint>
10#include <set>
11#include <stdexcept>
12
13#define MAX_REPETITION_THRESHOLD 2000
14//
15// helpers
16//
17
18// NOTE: assumes valid utf8 (but checks for overrun)
19static std::pair<uint32_t, const char *> decode_utf8(const char * src) {
20    static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 };
21    uint8_t  first_byte = static_cast<uint8_t>(*src);
22    uint8_t  highbits   = first_byte >> 4;
23    int      len        = lookup[highbits];
24    uint8_t  mask       = (1 << (8 - len)) - 1;
25    uint32_t value      = first_byte & mask;
26    const char * end    = src + len; // may overrun!
27    const char * pos    = src + 1;
28    for ( ; pos < end && *pos; pos++) {
29        value = (value << 6) + (static_cast<uint8_t>(*pos) & 0x3F);
30    }
31    return std::make_pair(value, pos);
32}
33
34static std::pair<std::vector<uint32_t>, llama_partial_utf8> decode_utf8(
35        const std::string & src,
36        llama_partial_utf8 partial_start) {
37    static const int      lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 3, 4 };
38    const char          * pos      = src.c_str();
39    std::vector<uint32_t> code_points;
40
41    // common english strings have the same number of codepoints and bytes. `+ 1` for the terminating 0.
42    code_points.reserve(src.size() + 1);
43    uint32_t value    = partial_start.value;
44    int      n_remain = partial_start.n_remain;
45
46    // continue previous decode, if applicable
47    while (*pos != 0 && n_remain > 0) {
48        uint8_t next_byte = static_cast<uint8_t>(*pos);
49        if ((next_byte >> 6) != 2) {
50            // invalid sequence, abort
51            code_points.push_back(0);
52            return std::make_pair(std::move(code_points), llama_partial_utf8{ 0, -1 });
53        }
54        value = (value << 6) + (next_byte & 0x3F);
55        ++pos;
56        --n_remain;
57    }
58
59    if (partial_start.n_remain > 0 && n_remain == 0) {
60        code_points.push_back(value);
61    }
62
63    // decode any subsequent utf-8 sequences, which may end in an incomplete one
64    while (*pos != 0) {
65        uint8_t first_byte = static_cast<uint8_t>(*pos);
66        uint8_t highbits   = first_byte >> 4;
67        n_remain   = lookup[highbits] - 1;
68
69        if (n_remain < 0) {
70            // invalid sequence, abort
71            code_points.clear();
72            code_points.push_back(0);
73            return std::make_pair(std::move(code_points), llama_partial_utf8{ 0, n_remain });
74        }
75
76        uint8_t mask  = (1 << (7 - n_remain)) - 1;
77        value = first_byte & mask;
78
79        ++pos;
80        while (*pos != 0 && n_remain > 0) {
81            value = (value << 6) + (static_cast<uint8_t>(*pos) & 0x3F);
82            ++pos;
83            --n_remain;
84        }
85        if (n_remain == 0) {
86            code_points.push_back(value);
87        }
88    }
89    code_points.push_back(0);
90
91    return std::make_pair(std::move(code_points), llama_partial_utf8{ value, n_remain });
92}
93
94static bool is_digit_char(char c) {
95    return '0' <= c && c <= '9';
96}
97
98static bool is_word_char(char c) {
99    return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '-' || is_digit_char(c);
100}
101
102static std::pair<uint32_t, const char *> parse_hex(const char * src, int size) {
103    const char * pos   = src;
104    const char * end   = src + size;
105    uint32_t     value = 0;
106    for ( ; pos < end && *pos; pos++) {
107        value <<= 4;
108        char c = *pos;
109        if ('a' <= c && c <= 'f') {
110            value += c - 'a' + 10;
111        } else if ('A' <= c && c <= 'F') {
112            value += c - 'A' + 10;
113        } else if ('0' <= c && c <= '9') {
114            value += c - '0';
115        } else {
116            break;
117        }
118    }
119    if (pos != end) {
120        throw std::runtime_error("expecting " + std::to_string(size) + " hex chars at " + src);
121    }
122    return std::make_pair(value, pos);
123}
124
125static const char * parse_space(const char * src, bool newline_ok) {
126    const char * pos = src;
127    while (*pos == ' ' || *pos == '\t' || *pos == '#' ||
128            (newline_ok && (*pos == '\r' || *pos == '\n'))) {
129        if (*pos == '#') {
130            while (*pos && *pos != '\r' && *pos != '\n') {
131                pos++;
132            }
133        } else {
134            pos++;
135        }
136    }
137    return pos;
138}
139
140static const char * parse_name(const char * src) {
141    const char * pos = src;
142    while (is_word_char(*pos)) {
143        pos++;
144    }
145    if (pos == src) {
146        throw std::runtime_error(std::string("expecting name at ") + src);
147    }
148    return pos;
149}
150
151static const char * parse_int(const char * src) {
152    const char * pos = src;
153    while (is_digit_char(*pos)) {
154        pos++;
155    }
156    if (pos == src) {
157        throw std::runtime_error(std::string("expecting integer at ") + src);
158    }
159    return pos;
160}
161
162static std::pair<uint32_t, const char *> parse_char(const char * src) {
163    if (*src == '\\') {
164        switch (src[1]) {
165            case 'x': return parse_hex(src + 2, 2);
166            case 'u': return parse_hex(src + 2, 4);
167            case 'U': return parse_hex(src + 2, 8);
168            case 't': return std::make_pair('\t', src + 2);
169            case 'r': return std::make_pair('\r', src + 2);
170            case 'n': return std::make_pair('\n', src + 2);
171            case '\\':
172            case '"':
173            case '[':
174            case ']':
175            case '-':
176                      return std::make_pair(src[1], src + 2);
177            default:
178                      throw std::runtime_error(std::string("unknown escape at ") + src);
179        }
180    } else if (*src) {
181        return decode_utf8(src);
182    }
183    throw std::runtime_error("unexpected end of input");
184}
185
186static std::pair<uint32_t, const char *> parse_token(const llama_vocab * vocab, const char * src) {
187    const char * pos = src;
188    if (*pos != '<') {
189        throw std::runtime_error(std::string("expecting '<' at ") + pos);
190    }
191    pos++;
192
193    // Parse <[id]>
194    if (*pos == '[') {
195        pos++;
196        const char * int_end = parse_int(pos);
197        uint32_t token_id = std::stoul(std::string(pos, int_end - pos));
198        pos = int_end;
199        if (*pos != ']') {
200            throw std::runtime_error(std::string("expecting ']' at ") + pos);
201        }
202        pos++;
203        if (*pos != '>') {
204            throw std::runtime_error(std::string("expecting '>' at ") + pos);
205        }
206        pos++;
207        return std::make_pair(token_id, pos);
208    }
209
210    if (vocab == nullptr) {
211        throw std::runtime_error(std::string("no vocab to parse token at ") + src);
212    }
213
214    // Parse <token> and tokenize to obtain the token id
215    while (*pos != 0 && *pos != '>') {
216        pos++;
217    }
218    if (*pos != '>') {
219        throw std::runtime_error(std::string("expecting '>' at ") + pos);
220    }
221    pos++;
222
223    llama_token tokens[2];
224    int32_t n_tokens = vocab->tokenize(src, static_cast<int32_t>(pos - src), tokens, 2, false, true);
225    if (n_tokens != 1) {
226        // must tokenize to exactly 1 token
227        throw std::runtime_error("invalid token '" + std::string(src, pos - src) + "'");
228    }
229    return std::make_pair(tokens[0], pos);
230}
231
232static void print_grammar_char(FILE * file, uint32_t c) {
233    if (0x20 <= c && c <= 0x7f) {
234        fprintf(file, "%c", static_cast<char>(c));
235    } else {
236        // cop out of encoding UTF-8
237        fprintf(file, "<U+%04X>", c);
238    }
239}
240
241static bool is_char_element(llama_grammar_element elem) {
242    switch (elem.type) {
243        case LLAMA_GRETYPE_CHAR:           return true;
244        case LLAMA_GRETYPE_CHAR_NOT:       return true;
245        case LLAMA_GRETYPE_CHAR_ALT:       return true;
246        case LLAMA_GRETYPE_CHAR_RNG_UPPER: return true;
247        case LLAMA_GRETYPE_CHAR_ANY:       return true;
248        default:                           return false;
249    }
250}
251
252static void print_rule_binary(FILE * file, const llama_grammar_rule & rule) {
253    for (auto elem : rule) {
254        switch (elem.type) {
255            case LLAMA_GRETYPE_END:            fprintf(file, "END");            break;
256            case LLAMA_GRETYPE_ALT:            fprintf(file, "ALT");            break;
257            case LLAMA_GRETYPE_RULE_REF:       fprintf(file, "RULE_REF");       break;
258            case LLAMA_GRETYPE_CHAR:           fprintf(file, "CHAR");           break;
259            case LLAMA_GRETYPE_CHAR_NOT:       fprintf(file, "CHAR_NOT");       break;
260            case LLAMA_GRETYPE_CHAR_RNG_UPPER: fprintf(file, "CHAR_RNG_UPPER"); break;
261            case LLAMA_GRETYPE_CHAR_ALT:       fprintf(file, "CHAR_ALT");       break;
262            case LLAMA_GRETYPE_CHAR_ANY:       fprintf(file, "CHAR_ANY");       break;
263            case LLAMA_GRETYPE_TOKEN:          fprintf(file, "TOKEN");          break;
264            case LLAMA_GRETYPE_TOKEN_NOT:      fprintf(file, "TOKEN_NOT");      break;
265        }
266        switch (elem.type) {
267            case LLAMA_GRETYPE_END:
268            case LLAMA_GRETYPE_ALT:
269            case LLAMA_GRETYPE_RULE_REF:
270                fprintf(file, "(%u) ", elem.value);
271                break;
272            case LLAMA_GRETYPE_CHAR:
273            case LLAMA_GRETYPE_CHAR_NOT:
274            case LLAMA_GRETYPE_CHAR_RNG_UPPER:
275            case LLAMA_GRETYPE_CHAR_ALT:
276            case LLAMA_GRETYPE_CHAR_ANY:
277                fprintf(file, "(\"");
278                print_grammar_char(file, elem.value);
279                fprintf(file, "\") ");
280                break;
281            case LLAMA_GRETYPE_TOKEN:
282                fprintf(file, "<[");
283                fprintf(file, "%u", elem.value);
284                fprintf(file, "]> ");
285                break;
286            case LLAMA_GRETYPE_TOKEN_NOT:
287                fprintf(file, "!");
288                fprintf(file, "<[");
289                fprintf(file, "%u", elem.value);
290                fprintf(file, "]> ");
291                break;
292        }
293    }
294    fprintf(file, "\n");
295}
296
297static void print_rule(
298        FILE     * file,
299        uint32_t   rule_id,
300        const llama_grammar_rule & rule,
301        const std::map<uint32_t, std::string> & symbol_id_names) {
302    if (rule.empty() || rule.back().type != LLAMA_GRETYPE_END) {
303        throw std::runtime_error(
304            "malformed rule, does not end with LLAMA_GRETYPE_END: " + std::to_string(rule_id));
305    }
306    fprintf(file, "%s ::= ", symbol_id_names.at(rule_id).c_str());
307    for (size_t i = 0, end = rule.size() - 1; i < end; i++) {
308        llama_grammar_element elem = rule[i];
309        switch (elem.type) {
310            case LLAMA_GRETYPE_END:
311                throw std::runtime_error(
312                    "unexpected end of rule: " + std::to_string(rule_id) + "," +
313                    std::to_string(i));
314            case LLAMA_GRETYPE_ALT:
315                fprintf(file, "| ");
316                break;
317            case LLAMA_GRETYPE_RULE_REF:
318                fprintf(file, "%s ", symbol_id_names.at(elem.value).c_str());
319                break;
320            case LLAMA_GRETYPE_CHAR:
321                fprintf(file, "[");
322                print_grammar_char(file, elem.value);
323                break;
324            case LLAMA_GRETYPE_CHAR_NOT:
325                fprintf(file, "[^");
326                print_grammar_char(file, elem.value);
327                break;
328            case LLAMA_GRETYPE_CHAR_RNG_UPPER:
329                if (i == 0 || !is_char_element(rule[i - 1])) {
330                    throw std::runtime_error(
331                        "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: " +
332                        std::to_string(rule_id) + "," + std::to_string(i));
333                }
334                fprintf(file, "-");
335                print_grammar_char(file, elem.value);
336                break;
337            case LLAMA_GRETYPE_CHAR_ALT:
338                if (i == 0 || !is_char_element(rule[i - 1])) {
339                    throw std::runtime_error(
340                        "LLAMA_GRETYPE_CHAR_ALT without preceding char: " +
341                        std::to_string(rule_id) + "," + std::to_string(i));
342                }
343                print_grammar_char(file, elem.value);
344                break;
345            case LLAMA_GRETYPE_CHAR_ANY:
346                fprintf(file, ".");
347                break;
348            case LLAMA_GRETYPE_TOKEN:
349                fprintf(file, "<[");
350                fprintf(file, "%u", elem.value);
351                fprintf(file, "]> ");
352                break;
353            case LLAMA_GRETYPE_TOKEN_NOT:
354                fprintf(file, "!");
355                fprintf(file, "<[");
356                fprintf(file, "%u", elem.value);
357                fprintf(file, "]> ");
358                break;
359        }
360        if (is_char_element(elem)) {
361            switch (rule[i + 1].type) {
362                case LLAMA_GRETYPE_CHAR_ALT:
363                case LLAMA_GRETYPE_CHAR_RNG_UPPER:
364                case LLAMA_GRETYPE_CHAR_ANY:
365                    break;
366                default:
367                    fprintf(file, "] ");
368            }
369        }
370    }
371    fprintf(file, "\n");
372}
373
374//
375// Regex utilities
376//
377
378size_t llama_grammar_trigger_pattern::find(const std::string & input) const {
379    auto find_start_pos = [](const std::smatch & match) {
380        // get from the first matched capturing group to the end of the string
381        size_t start = std::string::npos;
382        for (auto i = 1u; i < match.size(); i++) {
383            if (match.length(i) > 0) {
384                start = match.position(i);
385                break;
386            }
387        }
388        if (start == std::string::npos) {
389            start = match.position(0);
390        }
391        return start;
392    };
393
394    if (!pattern.empty() && pattern.front() == '^' && pattern.back() == '$') {
395        // match against the entire input
396        std::smatch match;
397        if (std::regex_match(input, match, regex)) {
398            return find_start_pos(match);
399        }
400    }
401
402    // search anywhere
403    std::smatch match;
404    if (std::regex_search(input, match, regex)) {
405        return find_start_pos(match);
406    }
407
408    return std::string::npos;
409}
410
411
412//
413// implementation
414//
415
416uint32_t llama_grammar_parser::get_symbol_id(const char * src, size_t len) {
417    uint32_t next_id = static_cast<uint32_t>(symbol_ids.size());
418    auto result = symbol_ids.emplace(std::string(src, len), next_id);
419    return result.first->second;
420}
421
422uint32_t llama_grammar_parser::generate_symbol_id(const std::string & base_name) {
423    uint32_t next_id = static_cast<uint32_t>(symbol_ids.size());
424    symbol_ids[base_name + '_' + std::to_string(next_id)] = next_id;
425    return next_id;
426}
427
428void llama_grammar_parser::add_rule(uint32_t rule_id, const llama_grammar_rule & rule) {
429    if (rules.size() <= rule_id) {
430        rules.resize(rule_id + 1);
431    }
432    rules[rule_id] = rule;
433}
434
435const char * llama_grammar_parser::parse_alternates(
436        const char        * src,
437        const std::string & rule_name,
438        uint32_t            rule_id,
439        bool                is_nested) {
440    llama_grammar_rule rule;
441    const char * pos = parse_sequence(src, rule_name, rule, is_nested);
442    while (*pos == '|') {
443        rule.push_back({LLAMA_GRETYPE_ALT, 0});
444        pos = parse_space(pos + 1, true);
445        pos = parse_sequence(pos, rule_name, rule, is_nested);
446    }
447    rule.push_back({LLAMA_GRETYPE_END, 0});
448    add_rule(rule_id, rule);
449    return pos;
450}
451
452const char * llama_grammar_parser::parse_sequence(
453        const char         * src,
454        const std::string  & rule_name,
455        llama_grammar_rule & rule,
456        bool               is_nested) {
457    size_t last_sym_start = rule.size();
458    const char * pos = src;
459    uint64_t n_prev_rules = 1;
460
461    // use UINT64_MAX as the empty value because we aligned to the proper uint64_t type so -1 can't be used
462    // (though it's technically the same as -1 now)
463    auto handle_repetitions = [&](uint64_t min_times, uint64_t max_times) {
464        bool no_max = max_times == UINT64_MAX;
465        if (last_sym_start == rule.size()) {
466            throw std::runtime_error(std::string("expecting preceding item to */+/?/{ at ") + pos);
467        }
468
469        // apply transformation to previous symbol (last_sym_start to end) according to
470        // the following rewrite rules:
471        // S{m,n} --> S S S (m times) S'(n-m)
472        //            S'(x)   ::= S S'(x-1) |
473        //            (... n-m definitions of these S' rules ...)
474        //            S'(1)   ::= S |
475        // S{m,} -->  S S S (m times) S'
476        //            S'     ::= S S' |
477        // S*     --> S{0,}
478        //        --> S'     ::= S S' |
479        // S+     --> S{1,}
480        //        --> S S'
481        //            S'     ::= S S' |
482        // S?     --> S{0,1}
483        //        --> S'
484        //            S'     ::= S |
485
486        llama_grammar_rule prev_rule(rule.begin() + last_sym_start, rule.end());
487        // Calculate the total number of rules that will be generated by this repetition
488        uint64_t total_rules = 1; // Start with 1 for the original rule
489        if (!no_max && max_times > 0) {
490            total_rules = max_times;
491        } else if (min_times > 0) {
492            total_rules = min_times;
493        }
494
495        if (n_prev_rules * total_rules > MAX_REPETITION_THRESHOLD) {
496            throw std::runtime_error("number of rules that are going to be repeated multiplied by the new repetition exceeds sane defaults, please reduce the number of repetitions or rule complexity");
497        }
498
499        if (min_times == 0) {
500            rule.resize(last_sym_start);
501        } else {
502            // Repeat the previous elements (min_times - 1) times
503            for (uint64_t i = 1; i < min_times; i++) {
504                rule.insert(rule.end(), prev_rule.begin(), prev_rule.end());
505            }
506        }
507
508        uint32_t last_rec_rule_id = 0;
509        auto n_opt = no_max ? 1 : max_times - min_times;
510
511        llama_grammar_rule rec_rule(prev_rule);
512        for (uint64_t i = 0; i < n_opt; i++) {
513            rec_rule.resize(prev_rule.size());
514            uint32_t rec_rule_id = generate_symbol_id( rule_name);
515            if (i > 0 || no_max) {
516                rec_rule.push_back({LLAMA_GRETYPE_RULE_REF, no_max ? rec_rule_id : last_rec_rule_id});
517            }
518            rec_rule.push_back({LLAMA_GRETYPE_ALT, 0});
519            rec_rule.push_back({LLAMA_GRETYPE_END, 0});
520            add_rule( rec_rule_id, rec_rule);
521            last_rec_rule_id = rec_rule_id;
522        }
523        if (n_opt > 0) {
524            rule.push_back({LLAMA_GRETYPE_RULE_REF, last_rec_rule_id});
525        }
526        n_prev_rules *= total_rules;
527        GGML_ASSERT(n_prev_rules >= 1);
528    };
529
530    while (*pos) {
531        if (*pos == '"') { // literal string
532            pos++;
533            last_sym_start = rule.size();
534            n_prev_rules = 1;
535            while (*pos != '"') {
536                if (!*pos) {
537                    throw std::runtime_error("unexpected end of input");
538                }
539                auto char_pair = parse_char(pos);
540                     pos       = char_pair.second;
541                rule.push_back({LLAMA_GRETYPE_CHAR, char_pair.first});
542            }
543            pos = parse_space(pos + 1, is_nested);
544        } else if (*pos == '[') { // char range(s)
545            pos++;
546            enum llama_gretype start_type = LLAMA_GRETYPE_CHAR;
547            if (*pos == '^') {
548                pos++;
549                start_type = LLAMA_GRETYPE_CHAR_NOT;
550            }
551            last_sym_start = rule.size();
552            n_prev_rules = 1;
553            while (*pos != ']') {
554                if (!*pos) {
555                    throw std::runtime_error("unexpected end of input");
556                }
557                auto char_pair = parse_char(pos);
558                     pos       = char_pair.second;
559                enum llama_gretype type = last_sym_start < rule.size()
560                    ? LLAMA_GRETYPE_CHAR_ALT
561                    : start_type;
562
563                rule.push_back({type, char_pair.first});
564                if (pos[0] == '-' && pos[1] != ']') {
565                    if (!pos[1]) {
566                        throw std::runtime_error("unexpected end of input");
567                    }
568                    auto endchar_pair = parse_char(pos + 1);
569                         pos          = endchar_pair.second;
570                    rule.push_back({LLAMA_GRETYPE_CHAR_RNG_UPPER, endchar_pair.first});
571                }
572            }
573            pos = parse_space(pos + 1, is_nested);
574        } else if (*pos == '<' || *pos == '!') { // token
575            auto type = LLAMA_GRETYPE_TOKEN;
576            if (*pos == '!') { // token inverse
577                type = LLAMA_GRETYPE_TOKEN_NOT;
578                pos++;
579            }
580            auto token_pair = parse_token(vocab, pos);
581            const char * token_end  = token_pair.second;
582            last_sym_start = rule.size();
583            n_prev_rules = 1;
584            rule.push_back({type, token_pair.first});
585            pos = parse_space(token_end, is_nested);
586        } else if (is_word_char(*pos)) { // rule reference
587            const char * name_end    = parse_name(pos);
588            uint32_t ref_rule_id = get_symbol_id(pos, name_end - pos);
589            pos = parse_space(name_end, is_nested);
590            last_sym_start = rule.size();
591            n_prev_rules = 1;
592            rule.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id});
593        } else if (*pos == '(') { // grouping
594            // parse nested alternates into synthesized rule
595            pos = parse_space(pos + 1, true);
596            uint32_t n_rules_before = symbol_ids.size();
597            uint32_t sub_rule_id = generate_symbol_id(rule_name);
598            pos = parse_alternates(pos, rule_name, sub_rule_id, true);
599            n_prev_rules = std::max(1u, (uint32_t)symbol_ids.size() - n_rules_before);
600            last_sym_start = rule.size();
601            // output reference to synthesized rule
602            rule.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id});
603            if (*pos != ')') {
604                throw std::runtime_error(std::string("expecting ')' at ") + pos);
605            }
606            pos = parse_space(pos + 1, is_nested);
607        } else if (*pos == '.') { // any char
608            last_sym_start = rule.size();
609            n_prev_rules = 1;
610            rule.push_back({LLAMA_GRETYPE_CHAR_ANY, 0});
611            pos = parse_space(pos + 1, is_nested);
612        } else if (*pos == '*') {
613            pos = parse_space(pos + 1, is_nested);
614            handle_repetitions(0, -1);
615        } else if (*pos == '+') {
616            pos = parse_space(pos + 1, is_nested);
617            handle_repetitions(1, -1);
618        } else if (*pos == '?') {
619            pos = parse_space(pos + 1, is_nested);
620            handle_repetitions(0, 1);
621        } else if (*pos == '{') {
622            pos = parse_space(pos + 1, is_nested);
623
624            if (!is_digit_char(*pos)) {
625                throw std::runtime_error(std::string("expecting an int at ") + pos);
626            }
627            const char * int_end = parse_int(pos);
628            uint64_t min_times = std::stoull(std::string(pos, int_end - pos));
629            pos = parse_space(int_end, is_nested);
630
631            uint64_t max_times = UINT64_MAX; // default: no max limit
632
633            if (*pos == '}') {
634                max_times = min_times;
635                pos = parse_space(pos + 1, is_nested);
636            } else if (*pos == ',') {
637                pos = parse_space(pos + 1, is_nested);
638
639                if (is_digit_char(*pos)) {
640                    const char * int_end = parse_int(pos);
641                    max_times = std::stoull(std::string(pos, int_end - pos));
642                    pos = parse_space(int_end, is_nested);
643                }
644
645                if (*pos != '}') {
646                    throw std::runtime_error(std::string("expecting '}' at ") + pos);
647                }
648                pos = parse_space(pos + 1, is_nested);
649            } else {
650                throw std::runtime_error(std::string("expecting ',' at ") + pos);
651            }
652            if (min_times > MAX_REPETITION_THRESHOLD) {
653                throw std::runtime_error(std::string("number of repetitions exceeds sane defaults, please reduce the number of repetitions"));
654            }
655            if (max_times != UINT64_MAX && max_times > MAX_REPETITION_THRESHOLD) {
656                max_times = UINT64_MAX;
657            }
658            handle_repetitions(min_times, max_times);
659        } else {
660            break;
661        }
662    }
663    return pos;
664}
665
666const char * llama_grammar_parser::parse_rule(const char * src) {
667    const char * name_end = parse_name(src);
668    const char * pos      = parse_space(name_end, false);
669    size_t       name_len = name_end - src;
670    uint32_t     rule_id  = get_symbol_id(src, name_len);
671    const std::string name(src, name_len);
672
673    if (!(pos[0] == ':' && pos[1] == ':' && pos[2] == '=')) {
674        throw std::runtime_error(std::string("expecting ::= at ") + pos);
675    }
676    pos = parse_space(pos + 3, true);
677
678    pos = parse_alternates(pos, name, rule_id, false);
679
680    if (*pos == '\r') {
681        pos += pos[1] == '\n' ? 2 : 1;
682    } else if (*pos == '\n') {
683        pos++;
684    } else if (*pos) {
685        throw std::runtime_error(std::string("expecting newline or end at ") + pos);
686    }
687    return parse_space(pos, true);
688}
689
690bool llama_grammar_parser::parse(const char * src) {
691    try {
692        const char * pos = parse_space(src, true);
693        while (*pos) {
694            pos = parse_rule(pos);
695        }
696        // Validate the state to ensure that all rules are defined
697        for (const auto & rule : rules) {
698            if (rule.empty()) {
699                throw std::runtime_error("Undefined rule");
700            }
701            for (const auto & elem : rule) {
702                if (elem.type == LLAMA_GRETYPE_RULE_REF) {
703                    // Ensure that the rule at that location exists
704                    if (elem.value >= rules.size() || rules[elem.value].empty()) {
705                        // Get the name of the rule that is missing
706                        for (const auto & kv : symbol_ids) {
707                            if (kv.second == elem.value) {
708                                throw std::runtime_error("Undefined rule identifier '" + kv.first + "'");
709                            }
710                        }
711                    }
712                }
713            }
714        }
715    } catch (const std::exception & err) {
716        fprintf(stderr, "%s: error parsing grammar: %s\n\n%s\n", __func__, err.what(), src);
717        rules.clear();
718        return false;
719    }
720
721    return true;
722}
723
724void llama_grammar_parser::print(FILE * file) {
725    try {
726        std::map<uint32_t, std::string> symbol_id_names;
727        for (const auto & kv : symbol_ids) {
728            symbol_id_names[kv.second] = kv.first;
729        }
730        for (size_t i = 0, end = rules.size(); i < end; i++) {
731            // fprintf(file, "%zu: ", i);
732            // print_rule_binary(file, rules[i]);
733            print_rule(file, uint32_t(i), rules[i], symbol_id_names);
734            // fprintf(file, "\n");
735        }
736    } catch (const std::exception & err) {
737        fprintf(stderr, "\n%s: error printing grammar: %s\n", __func__, err.what());
738    }
739}
740
741llama_grammar_stack llama_grammar_parser::c_rules() const {
742    llama_grammar_stack ret;
743    ret.reserve(rules.size());
744    for (const auto & rule : rules) {
745        ret.push_back(rule.data());
746    }
747    return ret;
748}
749
750// returns true iff pos points to the end of one of the definitions of a rule
751static bool llama_grammar_is_end_of_sequence(const llama_grammar_element * pos) {
752    switch (pos->type) {
753        case LLAMA_GRETYPE_END: return true;  // NOLINT
754        case LLAMA_GRETYPE_ALT: return true;  // NOLINT
755        default:                return false;
756    }
757}
758
759// returns true iff chr satisfies the char range at pos (regular or inverse range)
760// asserts that pos is pointing to a char range element
761static std::pair<bool, const llama_grammar_element *> llama_grammar_match_char(
762        const llama_grammar_element * pos,
763        const uint32_t                chr) {
764    bool found            = false;
765    bool is_positive_char = pos->type == LLAMA_GRETYPE_CHAR || pos->type == LLAMA_GRETYPE_CHAR_ANY;
766
767    GGML_ASSERT(is_positive_char || pos->type == LLAMA_GRETYPE_CHAR_NOT); // NOLINT
768
769    do {
770        if (pos[1].type == LLAMA_GRETYPE_CHAR_RNG_UPPER) {
771            // inclusive range, e.g. [a-z]
772            found = found || (pos->value <= chr && chr <= pos[1].value);
773            pos += 2;
774        } else if (pos->type == LLAMA_GRETYPE_CHAR_ANY) {
775            // Any character matches "."
776            found = true;
777            pos += 1;
778        } else {
779            // exact char match, e.g. [a] or "a"
780            found = found || pos->value == chr;
781            pos += 1;
782        }
783    } while (pos->type == LLAMA_GRETYPE_CHAR_ALT);
784
785    return std::make_pair(found == is_positive_char, pos);
786}
787
788// returns true iff some continuation of the given partial UTF-8 sequence could satisfy the char
789// range at pos (regular or inverse range)
790// asserts that pos is pointing to a char range element
791static bool llama_grammar_match_partial_char(
792        const llama_grammar_element * pos,
793        const llama_partial_utf8      partial_utf8) {
794    bool is_positive_char = pos->type == LLAMA_GRETYPE_CHAR || pos->type == LLAMA_GRETYPE_CHAR_ANY;
795    GGML_ASSERT(is_positive_char || pos->type == LLAMA_GRETYPE_CHAR_NOT);
796
797    uint32_t partial_value = partial_utf8.value;
798    int      n_remain      = partial_utf8.n_remain;
799
800    // invalid sequence or 7-bit char split across 2 bytes (overlong)
801    if (n_remain < 0 || (n_remain == 1 && partial_value < 2)) {
802        return false;
803    }
804
805    // range of possible code points this partial UTF-8 sequence could complete to
806    uint32_t low  = partial_value << (n_remain * 6);
807    uint32_t high = low | ((1 << (n_remain * 6)) - 1);
808
809    if (low == 0) {
810        if (n_remain == 2) {
811            low = 1 << 11;
812        } else if (n_remain == 3) {
813            low = 1 << 16;
814        }
815    }
816
817    do {
818        if (pos[1].type == LLAMA_GRETYPE_CHAR_RNG_UPPER) {
819            // inclusive range, e.g. [a-z]
820            if (pos->value <= high && low <= pos[1].value) {
821                return is_positive_char;
822            }
823            pos += 2;
824        } else if (pos->type == LLAMA_GRETYPE_CHAR_ANY) {
825            // Any character matches "."
826            return true;
827        } else {
828            // exact char match, e.g. [a] or "a"
829            if (low <= pos->value && pos->value <= high) {
830                return is_positive_char;
831            }
832            pos += 1;
833        }
834    } while (pos->type == LLAMA_GRETYPE_CHAR_ALT);
835
836    return !is_positive_char;
837}
838
839// returns true iff token matches the rule at pos (regular or inverse)
840// asserts that pos is pointing to a token element
841static bool llama_grammar_match_token(
842    const llama_grammar_element * pos,
843    const llama_token             token) {
844    GGML_ASSERT(pos->type == LLAMA_GRETYPE_TOKEN || pos->type == LLAMA_GRETYPE_TOKEN_NOT);
845    if (pos->type == LLAMA_GRETYPE_TOKEN) {
846        return pos->value == static_cast<uint32_t>(token);
847    }
848    if (pos->type == LLAMA_GRETYPE_TOKEN_NOT) {
849        return pos->value != static_cast<uint32_t>(token);
850    }
851    return false;
852}
853
854// transforms a grammar pushdown stack into N possible stacks, all ending
855// at a character range (terminal element)
856static void llama_grammar_advance_stack(
857        const llama_grammar_rules  & rules,
858        const llama_grammar_stack  & stack,
859        llama_grammar_stacks & new_stacks) {
860    std::vector<llama_grammar_stack> todo;
861    todo.push_back(stack);
862
863    auto stack_cmp = [](const llama_grammar_stack & a, const llama_grammar_stack & b) {
864        return std::lexicographical_compare(a.begin(), a.end(), b.begin(), b.end(),
865            [](const llama_grammar_element * pa, const llama_grammar_element * pb) {
866                return pa < pb;  // Compare pointer addresses
867            }
868        );
869    };
870
871    std::set<llama_grammar_stack, decltype(stack_cmp)> seen(stack_cmp);
872
873    while (!todo.empty()) {
874        llama_grammar_stack curr_stack_candidate = std::move(todo.back());
875        todo.pop_back();
876
877        auto [curr_stack_it, inserted] = seen.insert(std::move(curr_stack_candidate));
878        if (!inserted) {
879            continue;
880        }
881        const llama_grammar_stack & curr_stack = *curr_stack_it;
882
883        if (curr_stack.empty()) {
884            if (std::find(new_stacks.begin(), new_stacks.end(), curr_stack) == new_stacks.end()) {
885                new_stacks.emplace_back(curr_stack);
886            }
887            continue;
888        }
889
890        const llama_grammar_element * pos = curr_stack.back();
891
892        switch (pos->type) {
893        case LLAMA_GRETYPE_RULE_REF: {
894            const size_t                  rule_id = static_cast<size_t>(pos->value);
895            const llama_grammar_element * subpos  = rules[rule_id].data();
896            do {
897                // init new stack without the top (pos)
898                llama_grammar_stack next_stack(curr_stack.begin(), curr_stack.end() - 1);
899                if (!llama_grammar_is_end_of_sequence(pos + 1)) {
900                    // if this rule ref is followed by another element, add that to stack
901                    next_stack.push_back(pos + 1);
902                }
903                if (!llama_grammar_is_end_of_sequence(subpos)) {
904                    // if alternate is nonempty, add to stack
905                    next_stack.push_back(subpos);
906                }
907                todo.push_back(std::move(next_stack));
908                while (!llama_grammar_is_end_of_sequence(subpos)) {
909                    // scan to end of alternate def
910                    subpos++;
911                }
912                if (subpos->type == LLAMA_GRETYPE_ALT) {
913                    // there's another alternate def of this rule to process
914                    subpos++;
915                } else {
916                    break;
917                }
918            } while (true);
919            break;
920        }
921        case LLAMA_GRETYPE_CHAR:
922        case LLAMA_GRETYPE_CHAR_NOT:
923        case LLAMA_GRETYPE_CHAR_ANY:
924        case LLAMA_GRETYPE_TOKEN:
925        case LLAMA_GRETYPE_TOKEN_NOT:
926            if (std::find(new_stacks.begin(), new_stacks.end(), curr_stack) == new_stacks.end()) {
927                // only add the stack if it's not a duplicate of one we already have
928                new_stacks.emplace_back(curr_stack);
929            }
930            break;
931        default:
932            // end of alternate (LLAMA_GRETYPE_END, LLAMA_GRETYPE_ALT) or middle of char range
933            // (LLAMA_GRETYPE_CHAR_ALT, LLAMA_GRETYPE_CHAR_RNG_UPPER); stack should never be left on
934            // those
935            GGML_ABORT("fatal error");
936        }
937    }
938}
939
940static llama_grammar_candidates llama_grammar_reject_candidates(
941        const llama_grammar_rules      & rules,
942        const llama_grammar_stacks     & stacks,
943        const llama_grammar_candidates & candidates) {
944    GGML_ASSERT(!stacks.empty()); // REVIEW
945
946    if (candidates.empty()) {
947        return {};
948    }
949
950    auto rejects = llama_grammar_reject_candidates_for_stack(rules, stacks.front(), candidates);
951
952    for (size_t i = 1, size = stacks.size(); i < size; ++i) {
953        rejects = llama_grammar_reject_candidates_for_stack(rules, stacks[i], rejects);
954    }
955
956    return rejects;
957}
958
959static bool llama_grammar_detect_left_recursion(
960        const llama_grammar_rules & rules,
961        size_t rule_index,
962        std::vector<bool> * rules_visited,
963        std::vector<bool> * rules_in_progress,
964        std::vector<bool> * rules_may_be_empty) {
965    if ((*rules_in_progress)[rule_index]) {
966        return true;
967    }
968
969    (*rules_in_progress)[rule_index] = true;
970
971    const llama_grammar_rule & rule = rules[rule_index];
972
973    // First check if the rule might produce the empty string. This could be done combined with the second
974    // step but it's more readable as two steps.
975    bool at_rule_start = true;
976    for (size_t i = 0; i < rule.size(); i++) {
977        if (llama_grammar_is_end_of_sequence(&rule[i])) {
978            if (at_rule_start) {
979                (*rules_may_be_empty)[rule_index] = true;
980                break;
981            }
982            at_rule_start = true;
983        } else {
984            at_rule_start = false;
985        }
986    }
987
988    // Second, recurse into leftmost nonterminals (or next-leftmost as long as the previous nonterminal may
989    // be empty)
990    bool recurse_into_nonterminal = true;
991    for (size_t i = 0; i < rule.size(); i++) {
992        if (rule[i].type == LLAMA_GRETYPE_RULE_REF && recurse_into_nonterminal) {
993            if (llama_grammar_detect_left_recursion(rules, (size_t)rule[i].value, rules_visited, rules_in_progress, rules_may_be_empty)) {
994                return true;
995            }
996            if (!((*rules_may_be_empty)[(size_t)rule[i].value])) {
997                recurse_into_nonterminal = false;
998            }
999        } else if (llama_grammar_is_end_of_sequence(&rule[i])) {
1000            recurse_into_nonterminal = true;
1001        } else {
1002            recurse_into_nonterminal = false;
1003        }
1004    }
1005
1006    (*rules_in_progress)[rule_index] = false;
1007    (*rules_visited)[rule_index] = true;
1008
1009    return false;
1010}
1011
1012const llama_grammar_rules & llama_grammar_get_rules(const struct llama_grammar * grammar) {
1013    return grammar->rules;
1014}
1015
1016llama_grammar_stacks & llama_grammar_get_stacks(struct llama_grammar * grammar) {
1017    return grammar->stacks;
1018}
1019
1020static void llama_grammar_accept_chr(
1021        struct llama_grammar       & grammar,
1022        const llama_grammar_stack  & stack,
1023              uint32_t               chr,
1024              llama_grammar_stacks & new_stacks) {
1025    if (stack.empty()) {
1026        return;
1027    }
1028
1029    const llama_grammar_element * pos = stack.back();
1030
1031    // ignore if this turns into a token
1032    if (pos->type == LLAMA_GRETYPE_TOKEN || pos->type == LLAMA_GRETYPE_TOKEN_NOT) {
1033        return;
1034    }
1035
1036    auto match = llama_grammar_match_char(pos, chr);
1037    if (match.first) {
1038        llama_grammar_stack new_stack(stack.begin(), stack.end() - 1);
1039        if (!llama_grammar_is_end_of_sequence(match.second)) {
1040            new_stack.push_back(match.second);
1041        }
1042        llama_grammar_advance_stack(grammar.rules, new_stack, new_stacks);
1043    }
1044}
1045
1046void llama_grammar_accept(struct llama_grammar * grammar, uint32_t chr) {
1047    llama_grammar_stacks stacks_new;
1048    stacks_new.reserve(grammar->stacks.size());
1049
1050    for (const auto & stack : grammar->stacks) {
1051        llama_grammar_accept_chr(*grammar, stack, chr, stacks_new);
1052    }
1053
1054    grammar->stacks = std::move(stacks_new);
1055}
1056
1057llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
1058        const llama_grammar_rules      & rules,
1059        const llama_grammar_stack      & stack,
1060        const llama_grammar_candidates & candidates) {
1061
1062    llama_grammar_candidates rejects;
1063    rejects.reserve(candidates.size());
1064
1065    if (stack.empty()) {
1066        for (const auto & tok : candidates) {
1067            if (*tok.code_points != 0 || tok.partial_utf8.n_remain != 0) {
1068                rejects.push_back(tok);
1069            }
1070        }
1071        return rejects;
1072    }
1073
1074    const llama_grammar_element * stack_pos = stack.back();
1075
1076    // if the top of the stack is a token rule, then we only need to check the token id
1077    if (stack_pos->type == LLAMA_GRETYPE_TOKEN || stack_pos->type == LLAMA_GRETYPE_TOKEN_NOT) {
1078        for (const auto & tok : candidates) {
1079            if (*tok.code_points == 0) {
1080                // reached the end of a token consumed by char rules, reject iff it ended
1081                // in a partial response
1082                if (tok.partial_utf8.n_remain != 0) {
1083                    rejects.push_back(tok);
1084                }
1085            } else if (!llama_grammar_match_token(stack_pos, tok.id)) {
1086                rejects.push_back(tok);
1087            }
1088        }
1089        return rejects;
1090    }
1091
1092    llama_grammar_candidates next_candidates;
1093    next_candidates.reserve(candidates.size());
1094
1095    for (const auto & tok : candidates) {
1096        if (*tok.code_points == 0) {
1097            // reached end of full codepoints in token, reject iff it ended in a partial sequence
1098            // that cannot satisfy this position in grammar
1099            if (tok.partial_utf8.n_remain != 0 &&
1100                    !llama_grammar_match_partial_char(stack_pos, tok.partial_utf8)) {
1101                rejects.push_back(tok);
1102            }
1103        } else if (llama_grammar_match_char(stack_pos, *tok.code_points).first) {
1104            next_candidates.push_back({ tok.index, tok.code_points + 1, tok.partial_utf8, tok.id });
1105        } else {
1106            rejects.push_back(tok);
1107        }
1108    }
1109
1110    const auto * stack_pos_after = llama_grammar_match_char(stack_pos, 0).second;
1111
1112    // update top of stack to next element, if any
1113    llama_grammar_stack stack_after(stack.begin(), stack.end() - 1);
1114    if (!llama_grammar_is_end_of_sequence(stack_pos_after)) {
1115        stack_after.push_back(stack_pos_after);
1116    }
1117    llama_grammar_stacks next_stacks;
1118    llama_grammar_advance_stack(rules, stack_after, next_stacks);
1119
1120    auto next_rejects = llama_grammar_reject_candidates(rules, next_stacks, next_candidates);
1121    for (const auto & tok : next_rejects) {
1122        rejects.push_back({ tok.index, tok.code_points - 1, tok.partial_utf8, tok.id });
1123    }
1124
1125    return rejects;
1126}
1127
1128////////////////////
1129
1130struct llama_grammar * llama_grammar_init_impl(
1131        const struct llama_vocab * vocab,
1132        const llama_grammar_element ** rules,
1133        size_t n_rules,
1134        size_t start_rule_index) {
1135    const llama_grammar_element * pos;
1136
1137    // copy rule definitions into vectors
1138    llama_grammar_rules vec_rules(n_rules);
1139    for (size_t i = 0; i < n_rules; i++) {
1140        for (pos = rules[i]; pos->type != LLAMA_GRETYPE_END; pos++) {
1141            vec_rules[i].push_back(*pos);
1142        }
1143        vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
1144    }
1145
1146    // Validate that all rule references point to valid rules
1147    for (size_t i = 0; i < n_rules; i++) {
1148        for (const auto & elem : vec_rules[i]) {
1149            if (elem.type == LLAMA_GRETYPE_RULE_REF) {
1150                if (elem.value >= n_rules || vec_rules[elem.value].empty()) {
1151                    LLAMA_LOG_ERROR("invalid grammar: rule %zu references undefined rule %u\n", i, elem.value);
1152                    return nullptr;
1153                }
1154            }
1155        }
1156    }
1157
1158    // Check for left recursion
1159    std::vector<bool> rules_visited(n_rules);
1160    std::vector<bool> rules_in_progress(n_rules);
1161    std::vector<bool> rules_may_be_empty(n_rules);
1162    for (size_t i = 0; i < n_rules; i++) {
1163        if (rules_visited[i]) {
1164            continue;
1165        }
1166        if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
1167            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i);
1168            return nullptr;
1169        }
1170    }
1171
1172    // loop over alternates of start rule to build initial stacks
1173    llama_grammar_stacks stacks;
1174    pos = vec_rules[start_rule_index].data();
1175    do {
1176        llama_grammar_stack stack;
1177        if (!llama_grammar_is_end_of_sequence(pos)) {
1178            // if alternate is nonempty, add to stack
1179            stack.push_back(pos);
1180        }
1181        llama_grammar_advance_stack(vec_rules, stack, stacks);
1182        while (!llama_grammar_is_end_of_sequence(pos)) {
1183            // scan to end of alternate def
1184            pos++;
1185        }
1186        if (pos->type == LLAMA_GRETYPE_ALT) {
1187            // there's another alternate def of this rule to process
1188            pos++;
1189        } else {
1190            break;
1191        }
1192    } while (true);
1193
1194    // Important: vec_rules has to be moved here, not copied, because stacks contains
1195    // pointers to elements of vec_rules. If vec_rules were copied into llama_grammar
1196    // then the pointers would be invalidated when the local vec_rules goes out of scope.
1197    return new llama_grammar {
1198        vocab,
1199        std::move(vec_rules),
1200        std::move(stacks),
1201        /* .partial_utf8 = */             {},
1202        /* .lazy = */                     false,
1203        /* .awaiting_trigger = */         false,
1204        /* .trigger_buffer = */           "",
1205        /* .trigger_buffer_positions = */ {},
1206        /* .trigger_tokens = */           {},
1207        /* .trigger_patterns = */         {},
1208    };
1209}
1210
1211struct llama_grammar * llama_grammar_init_impl(
1212        const struct llama_vocab * vocab,
1213                      const char * grammar_str,
1214                      const char * grammar_root,
1215                              bool lazy,
1216                     const char ** trigger_patterns,
1217                            size_t num_trigger_patterns,
1218               const llama_token * trigger_tokens,
1219                            size_t num_trigger_tokens) {
1220    llama_grammar_parser parser(vocab);
1221
1222    // if there is a grammar, parse it
1223    // rules will be empty (default) if there are parse errors
1224    if (!parser.parse(grammar_str) || parser.rules.empty()) {
1225        LLAMA_LOG_ERROR("failed to parse grammar\n");
1226        return nullptr;
1227    }
1228
1229    // Ensure that the grammar contains the start symbol
1230    if (parser.symbol_ids.find(grammar_root) == parser.symbol_ids.end()) {
1231        LLAMA_LOG_ERROR("grammar does not contain a '%s' symbol\n", grammar_root);
1232        return nullptr;
1233    }
1234
1235    std::vector<const llama_grammar_element *> grammar_rules(parser.c_rules());
1236
1237    const size_t n_rules = grammar_rules.size();
1238    const size_t start_rule_index = parser.symbol_ids.at(grammar_root);
1239
1240    const llama_grammar_element * pos;
1241
1242    // copy rule definitions into vectors
1243    llama_grammar_rules vec_rules(n_rules);
1244    for (size_t i = 0; i < n_rules; i++) {
1245        for (pos = grammar_rules[i]; pos->type != LLAMA_GRETYPE_END; pos++) {
1246            vec_rules[i].push_back(*pos);
1247        }
1248        vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
1249    }
1250
1251    // Check for left recursion
1252    std::vector<bool> rules_visited(n_rules);
1253    std::vector<bool> rules_in_progress(n_rules);
1254    std::vector<bool> rules_may_be_empty(n_rules);
1255    for (size_t i = 0; i < n_rules; i++) {
1256        if (rules_visited[i]) {
1257            continue;
1258        }
1259        if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
1260            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i);
1261            return nullptr;
1262        }
1263    }
1264
1265    // loop over alternates of start rule to build initial stacks
1266    llama_grammar_stacks stacks;
1267    pos = vec_rules[start_rule_index].data();
1268    do {
1269        llama_grammar_stack stack;
1270        if (!llama_grammar_is_end_of_sequence(pos)) {
1271            // if alternate is nonempty, add to stack
1272            stack.push_back(pos);
1273        }
1274        llama_grammar_advance_stack(vec_rules, stack, stacks);
1275        while (!llama_grammar_is_end_of_sequence(pos)) {
1276            // scan to end of alternate def
1277            pos++;
1278        }
1279        if (pos->type == LLAMA_GRETYPE_ALT) {
1280            // there's another alternate def of this rule to process
1281            pos++;
1282        } else {
1283            break;
1284        }
1285    } while (true);
1286
1287    std::vector<llama_token>    vec_trigger_tokens;
1288    std::vector<llama_grammar_trigger_pattern> vec_trigger_patterns;
1289    for (size_t i = 0; i < num_trigger_tokens; i++) {
1290        GGML_ASSERT(trigger_tokens != nullptr);
1291        vec_trigger_tokens.push_back(trigger_tokens[i]);
1292    }
1293    for (size_t i = 0; i < num_trigger_patterns; i++) {
1294        GGML_ASSERT(trigger_patterns != nullptr);
1295        auto & trigger = vec_trigger_patterns.emplace_back();
1296        trigger.pattern = trigger_patterns[i];
1297        trigger.regex = std::regex(trigger.pattern);
1298    }
1299
1300    // Important: vec_rules has to be moved here, not copied, because stacks contains
1301    // pointers to elements of vec_rules. If vec_rules were copied into llama_grammar
1302    // then the pointers would be invalidated when the local vec_rules goes out of scope.
1303    return new llama_grammar {
1304        vocab,
1305        std::move(vec_rules),
1306        std::move(stacks),
1307        /* .partial_utf8 = */             {},
1308        /* .lazy = */                     lazy,
1309        /* .awaiting_trigger = */         lazy,
1310        /* .trigger_buffer = */           "",
1311        /* .trigger_buffer_positions = */ {},
1312        std::move(vec_trigger_tokens),
1313        std::move(vec_trigger_patterns),
1314    };
1315}
1316
1317void llama_grammar_free_impl(struct llama_grammar * grammar) {
1318    if (grammar == nullptr) {
1319        return;
1320    }
1321
1322    delete grammar;
1323}
1324
1325struct llama_grammar * llama_grammar_clone_impl(const struct llama_grammar & grammar) {
1326    auto * result = new llama_grammar {
1327        grammar.vocab,
1328        grammar.rules,
1329        grammar.stacks,
1330        grammar.partial_utf8,
1331        grammar.lazy,
1332        grammar.awaiting_trigger,
1333        grammar.trigger_buffer,
1334        grammar.trigger_buffer_positions,
1335        grammar.trigger_tokens,
1336        grammar.trigger_patterns,
1337    };
1338
1339    // redirect elements in stacks to point to new rules
1340    for (size_t is = 0; is < result->stacks.size(); is++) {
1341        for (size_t ie = 0; ie < result->stacks[is].size(); ie++) {
1342            for (size_t ir0 = 0; ir0 < grammar.rules.size(); ir0++) {
1343                for (size_t ir1 = 0; ir1 < grammar.rules[ir0].size(); ir1++) {
1344                    if (grammar.stacks[is][ie] == &grammar.rules[ir0][ir1]) {
1345                        result->stacks[is][ie] =  &result->rules[ir0][ir1];
1346                    }
1347                }
1348            }
1349        }
1350    }
1351
1352    return result;
1353}
1354
1355void llama_grammar_apply_impl(const struct llama_grammar & grammar, llama_token_data_array * cur_p) {
1356    GGML_ASSERT(grammar.vocab != nullptr);
1357
1358    if (grammar.awaiting_trigger) {
1359        return;
1360    }
1361
1362    bool allow_eog = false;
1363    for (const auto & stack : grammar.stacks) {
1364        if (stack.empty()) {
1365            allow_eog = true;
1366            break;
1367        }
1368    }
1369
1370    std::vector<std::pair<std::vector<uint32_t>, llama_partial_utf8>> candidates_decoded;
1371    candidates_decoded.reserve(cur_p->size);
1372
1373    llama_grammar_candidates candidates_grammar;
1374    candidates_grammar.reserve(cur_p->size);
1375
1376    for (size_t i = 0; i < cur_p->size; ++i) {
1377        const llama_token id      = cur_p->data[i].id;
1378        const std::string & piece = grammar.vocab->token_to_piece(id);
1379
1380        if (grammar.vocab->is_eog(id)) {
1381            if (!allow_eog) {
1382                cur_p->data[i].logit = -INFINITY;
1383            }
1384        } else if (piece.empty() || piece[0] == 0) {
1385            cur_p->data[i].logit = -INFINITY;
1386        } else {
1387            candidates_decoded.push_back(decode_utf8(piece, grammar.partial_utf8));
1388            candidates_grammar.push_back({ i, candidates_decoded.back().first.data(), candidates_decoded.back().second, id });
1389        }
1390    }
1391
1392    const auto rejects = llama_grammar_reject_candidates(grammar.rules, grammar.stacks, candidates_grammar);
1393    for (const auto & reject : rejects) {
1394        cur_p->data[reject.index].logit = -INFINITY;
1395    }
1396}
1397
1398void llama_grammar_accept_impl(struct llama_grammar & grammar, llama_token token) {
1399    GGML_ASSERT(grammar.vocab != nullptr);
1400
1401    const auto & piece = grammar.vocab->token_to_piece(token);
1402
1403    if (grammar.awaiting_trigger) {
1404        if (std::find(grammar.trigger_tokens.begin(), grammar.trigger_tokens.end(), token) != grammar.trigger_tokens.end()) {
1405            grammar.awaiting_trigger = false;
1406            grammar.trigger_buffer.clear();
1407            llama_grammar_accept_token(grammar, token, piece);
1408            LLAMA_LOG_DEBUG("Grammar triggered on token %u (`%s`)", token, piece.c_str());
1409            return;
1410        } else {
1411            auto position = std::make_pair(grammar.trigger_buffer.size(), grammar.trigger_buffer.size() + piece.size());
1412            grammar.trigger_buffer_positions.push_back(std::make_pair(token, position));
1413            grammar.trigger_buffer += piece;
1414
1415            for (const auto & trigger_pattern : grammar.trigger_patterns) {
1416                auto start = trigger_pattern.find(grammar.trigger_buffer);
1417                if (start != std::string::npos) {
1418                    grammar.awaiting_trigger = false;
1419
1420                    // replay tokens that overlap with [start, end)
1421                    for (const auto & [tok, tok_pos] : grammar.trigger_buffer_positions) {
1422                        auto [tok_start, tok_end] = tok_pos;
1423                        if (tok_end <= start) {
1424                            continue;
1425                        }
1426
1427                        size_t piece_start = (tok_start < start) ? start : tok_start; // allow for partial token pieces
1428                        size_t piece_len = tok_end - piece_start;
1429                        auto tok_piece = grammar.trigger_buffer.substr(piece_start, piece_len);
1430                        llama_grammar_accept_token(grammar, tok, tok_piece);
1431                    }
1432
1433                    auto constrained_str = grammar.trigger_buffer.substr(start);
1434                    grammar.trigger_buffer.clear();
1435                    grammar.trigger_buffer_positions.clear();
1436                    LLAMA_LOG_DEBUG("Grammar triggered on regex: '%s'\n", constrained_str.c_str());
1437                    return;
1438                }
1439            }
1440            LLAMA_LOG_DEBUG("Grammar still awaiting trigger after token %d (`%s`)\n", token, piece.c_str());
1441            return;
1442        }
1443    }
1444
1445    if (grammar.vocab->is_eog(token)) {
1446        for (const auto & stack : grammar.stacks) {
1447            if (stack.empty()) {
1448                return;
1449            }
1450        }
1451        GGML_ABORT("fatal error");
1452    }
1453
1454    llama_grammar_accept_token(grammar, token, piece);
1455}
1456
1457void llama_grammar_accept_str(struct llama_grammar & grammar, const std::string & piece) {
1458    // Note terminating 0 in decoded string
1459    const auto   decoded     = decode_utf8(piece, grammar.partial_utf8);
1460    const auto & code_points = decoded.first;
1461
1462    for (auto it = code_points.begin(), end = code_points.end() - 1; it != end; ++it) {
1463        llama_grammar_accept(&grammar, *it);
1464    }
1465
1466    grammar.partial_utf8 = decoded.second;
1467    if (grammar.stacks.empty()) {
1468        throw std::runtime_error("Unexpected empty grammar stack after accepting piece: " + piece);
1469    }
1470}
1471
1472void llama_grammar_accept_token(struct llama_grammar & grammar, llama_token token, const std::string & piece) {
1473    // Note terminating 0 in decoded string
1474    const auto   decoded     = decode_utf8(piece, grammar.partial_utf8);
1475    const auto & code_points = decoded.first;
1476
1477    llama_grammar_stacks stacks_new;
1478    stacks_new.reserve(grammar.stacks.size());
1479
1480    for (const auto & stack : grammar.stacks) {
1481        if (stack.empty()) {
1482            continue;
1483        }
1484
1485        const llama_grammar_element * pos = stack.back();
1486
1487        if (pos->type == LLAMA_GRETYPE_TOKEN || pos->type == LLAMA_GRETYPE_TOKEN_NOT) {
1488            if (llama_grammar_match_token(pos, token)) {
1489                llama_grammar_stack new_stack(stack.begin(), stack.end() - 1);
1490                if (!llama_grammar_is_end_of_sequence(pos + 1)) {
1491                    new_stack.push_back(pos + 1);
1492                }
1493                llama_grammar_advance_stack(grammar.rules, new_stack, stacks_new);
1494            }
1495        } else {
1496            llama_grammar_stacks current_stacks = {stack};
1497
1498            for (auto it = code_points.begin(), end = code_points.end() - 1; it != end; ++it) {
1499                llama_grammar_stacks next_stacks;
1500
1501                for (const auto & cur_stack : current_stacks) {
1502                    llama_grammar_accept_chr(grammar, cur_stack, *it, next_stacks);
1503                }
1504
1505                current_stacks = std::move(next_stacks);
1506                if (current_stacks.empty()) {
1507                    break;
1508                }
1509            }
1510
1511            for (auto & surviving_stack : current_stacks) {
1512                if (std::find(stacks_new.begin(), stacks_new.end(), surviving_stack) == stacks_new.end()) {
1513                    stacks_new.emplace_back(surviving_stack);
1514                }
1515            }
1516        }
1517    }
1518
1519    grammar.stacks = std::move(stacks_new);
1520    grammar.partial_utf8 = decoded.second;
1521
1522    if (grammar.stacks.empty()) {
1523        throw std::runtime_error("Unexpected empty grammar stack after accepting piece: " + piece + " (" + std::to_string(token) + ")");
1524    }
1525}
1526