char-fork/llama.cpp

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

Andrea Arcangelicommon/grammar: fix grammar parsing issues to prevent stack overflow and hangs (#18604)990e4d969

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
11.0 KiB406 linesraw
1#ifdef NDEBUG
2#undef NDEBUG
3#endif
4
5#include "llama.h"
6
7#include "../src/llama-grammar.h"
8
9#include <cassert>
10#include <stdexcept>
11
12int main()
13{
14    llama_grammar_parser parsed_grammar;
15
16    std::vector<std::pair<std::string, uint32_t>> expected = {
17        {"expr", 2},
18        {"expr_6", 6},
19        {"expr_7", 7},
20        {"ident", 8},
21        {"ident_10", 10},
22        {"num", 9},
23        {"num_11", 11},
24        {"root", 0},
25        {"root_1", 1},
26        {"root_5", 5},
27        {"term", 4},
28        {"ws", 3},
29        {"ws_12", 12},
30    };
31
32    std::vector<std::vector<llama_grammar_element>> expected_rules = {
33        {{LLAMA_GRETYPE_RULE_REF, 5}, {LLAMA_GRETYPE_END, 0}},
34        {
35            {LLAMA_GRETYPE_RULE_REF, 2},
36            {LLAMA_GRETYPE_CHAR, 61},
37            {LLAMA_GRETYPE_RULE_REF, 3},
38            {LLAMA_GRETYPE_RULE_REF, 4},
39            {LLAMA_GRETYPE_CHAR, 10},
40            {LLAMA_GRETYPE_END, 0},
41        },
42        {{LLAMA_GRETYPE_RULE_REF, 4}, {LLAMA_GRETYPE_RULE_REF, 7}, {LLAMA_GRETYPE_END, 0}},
43        {{LLAMA_GRETYPE_RULE_REF, 12}, {LLAMA_GRETYPE_END, 0}},
44        {
45            {LLAMA_GRETYPE_RULE_REF, 8},
46            {LLAMA_GRETYPE_ALT, 0},
47            {LLAMA_GRETYPE_RULE_REF, 9},
48            {LLAMA_GRETYPE_ALT, 0},
49            {LLAMA_GRETYPE_CHAR, 40},
50            {LLAMA_GRETYPE_RULE_REF, 3},
51            {LLAMA_GRETYPE_RULE_REF, 2},
52            {LLAMA_GRETYPE_CHAR, 41},
53            {LLAMA_GRETYPE_RULE_REF, 3},
54            {LLAMA_GRETYPE_END, 0},
55        },
56        {{LLAMA_GRETYPE_RULE_REF, 1}, {LLAMA_GRETYPE_RULE_REF, 5}, {LLAMA_GRETYPE_ALT, 0}, {LLAMA_GRETYPE_RULE_REF, 1}, {LLAMA_GRETYPE_END, 0}},
57        {
58            {LLAMA_GRETYPE_CHAR, 45},
59            {LLAMA_GRETYPE_CHAR_ALT, 43},
60            {LLAMA_GRETYPE_CHAR_ALT, 42},
61            {LLAMA_GRETYPE_CHAR_ALT, 47},
62            {LLAMA_GRETYPE_RULE_REF, 4},
63            {LLAMA_GRETYPE_END, 0},
64        },
65        {{LLAMA_GRETYPE_RULE_REF, 6}, {LLAMA_GRETYPE_RULE_REF, 7}, {LLAMA_GRETYPE_ALT, 0}, {LLAMA_GRETYPE_END, 0}},
66        {
67            {LLAMA_GRETYPE_CHAR, 97},
68            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 122},
69            {LLAMA_GRETYPE_RULE_REF, 10},
70            {LLAMA_GRETYPE_RULE_REF, 3},
71            {LLAMA_GRETYPE_END, 0},
72        },
73        {{LLAMA_GRETYPE_RULE_REF, 11}, {LLAMA_GRETYPE_RULE_REF, 3}, {LLAMA_GRETYPE_END, 0}},
74        {
75            {LLAMA_GRETYPE_CHAR, 97},
76            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 122},
77            {LLAMA_GRETYPE_CHAR_ALT, 48},
78            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},
79            {LLAMA_GRETYPE_CHAR_ALT, 95},
80            {LLAMA_GRETYPE_RULE_REF, 10},
81            {LLAMA_GRETYPE_ALT, 0},
82            {LLAMA_GRETYPE_END, 0},
83        },
84        {
85            {LLAMA_GRETYPE_CHAR, 48},
86            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},
87            {LLAMA_GRETYPE_RULE_REF, 11},
88            {LLAMA_GRETYPE_ALT, 0},
89            {LLAMA_GRETYPE_CHAR, 48},
90            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},
91            {LLAMA_GRETYPE_END, 0},
92        },
93        {
94            {LLAMA_GRETYPE_CHAR, 32},
95            {LLAMA_GRETYPE_CHAR_ALT, 9},
96            {LLAMA_GRETYPE_CHAR_ALT, 10},
97            {LLAMA_GRETYPE_RULE_REF, 12},
98            {LLAMA_GRETYPE_ALT, 0},
99            {LLAMA_GRETYPE_END, 0},
100        },
101    };
102
103    for (auto pair : expected)
104    {
105        parsed_grammar.symbol_ids[pair.first] = pair.second;
106    }
107
108    for (auto rule : expected_rules)
109    {
110        parsed_grammar.rules.emplace_back();
111        for (auto element : rule)
112        {
113            parsed_grammar.rules.back().push_back(element);
114        }
115    }
116
117    std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar.c_rules());
118
119    llama_grammar * grammar = llama_grammar_init_impl(nullptr, grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
120    if (grammar == nullptr) {
121        throw std::runtime_error("Failed to initialize llama_grammar");
122    }
123
124    std::vector<std::vector<llama_grammar_element>> expected_stacks = {
125        {
126            {LLAMA_GRETYPE_CHAR, 61},
127            {LLAMA_GRETYPE_RULE_REF, 7},
128            {LLAMA_GRETYPE_CHAR, 40},
129        },
130        {
131            {LLAMA_GRETYPE_CHAR, 61},
132            {LLAMA_GRETYPE_RULE_REF, 7},
133            {LLAMA_GRETYPE_RULE_REF, 3},
134            {LLAMA_GRETYPE_CHAR, 48},
135        },
136        {
137            {LLAMA_GRETYPE_CHAR, 61},
138            {LLAMA_GRETYPE_RULE_REF, 7},
139            {LLAMA_GRETYPE_RULE_REF, 3},
140            {LLAMA_GRETYPE_CHAR, 48},
141        },
142        {
143            {LLAMA_GRETYPE_CHAR, 61},
144            {LLAMA_GRETYPE_RULE_REF, 7},
145            {LLAMA_GRETYPE_CHAR, 97},
146        },
147        {
148            {LLAMA_GRETYPE_RULE_REF, 5},
149            {LLAMA_GRETYPE_CHAR, 61},
150            {LLAMA_GRETYPE_RULE_REF, 7},
151            {LLAMA_GRETYPE_CHAR, 40},
152        },
153        {
154            {LLAMA_GRETYPE_RULE_REF, 5},
155            {LLAMA_GRETYPE_CHAR, 61},
156            {LLAMA_GRETYPE_RULE_REF, 7},
157            {LLAMA_GRETYPE_RULE_REF, 3},
158            {LLAMA_GRETYPE_CHAR, 48},
159        },
160        {
161            {LLAMA_GRETYPE_RULE_REF, 5},
162            {LLAMA_GRETYPE_CHAR, 61},
163            {LLAMA_GRETYPE_RULE_REF, 7},
164            {LLAMA_GRETYPE_RULE_REF, 3},
165            {LLAMA_GRETYPE_CHAR, 48},
166        },
167        {
168            {LLAMA_GRETYPE_RULE_REF, 5},
169            {LLAMA_GRETYPE_CHAR, 61},
170            {LLAMA_GRETYPE_RULE_REF, 7},
171            {LLAMA_GRETYPE_CHAR, 97},
172        }};
173
174    auto index = 0;
175    for (const llama_grammar_stack & stack : llama_grammar_get_stacks(grammar))
176    {
177        // compare stack to expected_stack
178        for (uint32_t i = 0; i < stack.size(); i++)
179        {
180            const llama_grammar_element * element = stack[i];
181            const llama_grammar_element & expected_element = expected_stacks[index][i];
182
183            // pretty print error message before asserting
184            if (expected_element.type != element->type || expected_element.value != element->value)
185            {
186                fprintf(stderr, "index: %d\n", index);
187                fprintf(stderr, "expected_element: %d, %u\n", expected_element.type, expected_element.value);
188                fprintf(stderr, "actual_element: %d, %u\n", element->type, element->value);
189                fprintf(stderr, "expected_element != actual_element\n");
190            }
191
192            assert(expected_element.type == element->type && expected_element.value == element->value);
193        }
194        index++;
195    }
196
197    std::vector<llama_grammar_candidate> next_candidates;
198    next_candidates.resize(23);
199
200    for (size_t i = 0; i < 23; ++i)
201    {
202        uint32_t *cp = new uint32_t[2]; // dynamically allocate memory for code_point
203        cp[0] = 37 + i;
204        cp[1] = 0;
205        next_candidates[i] = {i, cp, {}, 0};
206    }
207
208    std::vector<std::vector<std::pair<uint32_t, uint16_t>>> expected_reject = {
209        {
210            {0, 37},
211            {1, 38},
212            {2, 39},
213            {4, 41},
214            {5, 42},
215            {6, 43},
216            {7, 44},
217            {8, 45},
218            {9, 46},
219            {10, 47},
220            {11, 48},
221            {12, 49},
222            {13, 50},
223            {14, 51},
224            {15, 52},
225            {16, 53},
226            {17, 54},
227            {18, 55},
228            {19, 56},
229            {20, 57},
230            {21, 58},
231            {22, 59},
232            {23, 60},
233        },
234        {
235            {0, 37},
236            {1, 38},
237            {2, 39},
238            {3, 40},
239            {4, 41},
240            {5, 42},
241            {6, 43},
242            {7, 44},
243            {8, 45},
244            {9, 46},
245            {10, 47},
246            {21, 58},
247            {22, 59},
248            {23, 60},
249        },
250        {
251            {0, 37},
252            {1, 38},
253            {2, 39},
254            {3, 40},
255            {4, 41},
256            {5, 42},
257            {6, 43},
258            {7, 44},
259            {8, 45},
260            {9, 46},
261            {10, 47},
262            {21, 58},
263            {22, 59},
264            {23, 60},
265        },
266        {
267            {0, 37},
268            {1, 38},
269            {2, 39},
270            {3, 40},
271            {4, 41},
272            {5, 42},
273            {6, 43},
274            {7, 44},
275            {8, 45},
276            {9, 46},
277            {10, 47},
278            {11, 48},
279            {12, 49},
280            {13, 50},
281            {14, 51},
282            {15, 52},
283            {16, 53},
284            {17, 54},
285            {18, 55},
286            {19, 56},
287            {20, 57},
288            {21, 58},
289            {22, 59},
290        },
291        {
292            {0, 37},
293            {1, 38},
294            {2, 39},
295            {4, 41},
296            {5, 42},
297            {6, 43},
298            {7, 44},
299            {8, 45},
300            {9, 46},
301            {10, 47},
302            {11, 48},
303            {12, 49},
304            {13, 50},
305            {14, 51},
306            {15, 52},
307            {16, 53},
308            {17, 54},
309            {18, 55},
310            {19, 56},
311            {20, 57},
312            {21, 58},
313            {22, 59},
314            {23, 60},
315        },
316        {
317            {0, 37},
318            {1, 38},
319            {2, 39},
320            {3, 40},
321            {4, 41},
322            {5, 42},
323            {6, 43},
324            {7, 44},
325            {8, 45},
326            {9, 46},
327            {10, 47},
328            {21, 58},
329            {22, 59},
330            {23, 60},
331        },
332        {
333            {0, 37},
334            {1, 38},
335            {2, 39},
336            {3, 40},
337            {4, 41},
338            {5, 42},
339            {6, 43},
340            {7, 44},
341            {8, 45},
342            {9, 46},
343            {10, 47},
344            {21, 58},
345            {22, 59},
346            {23, 60},
347        },
348        {
349            {0, 37},
350            {1, 38},
351            {2, 39},
352            {3, 40},
353            {4, 41},
354            {5, 42},
355            {6, 43},
356            {7, 44},
357            {8, 45},
358            {9, 46},
359            {10, 47},
360            {11, 48},
361            {12, 49},
362            {13, 50},
363            {14, 51},
364            {15, 52},
365            {16, 53},
366            {17, 54},
367            {18, 55},
368            {19, 56},
369            {20, 57},
370            {21, 58},
371            {22, 59},
372        },
373    };
374
375    std::vector<llama_grammar_candidate> rejects = llama_grammar_reject_candidates_for_stack(llama_grammar_get_rules(grammar), llama_grammar_get_stacks(grammar)[0], next_candidates);
376
377    std::vector<std::vector<llama_grammar_candidate>> all_rejects;
378
379    for (std::size_t count = 0; count < llama_grammar_get_stacks(grammar).size(); ++count)
380    {
381        rejects = llama_grammar_reject_candidates_for_stack(llama_grammar_get_rules(grammar), llama_grammar_get_stacks(grammar)[count], next_candidates);
382        all_rejects.push_back(rejects);
383    }
384
385    index = 0;
386    for (auto rej : all_rejects)
387    {
388        for (uint32_t i = 0; i < rej.size(); i++)
389        {
390            auto element = rej[i];
391            auto expected_element = expected_reject[index][i];
392            assert(element.index == expected_element.first && *element.code_points == expected_element.second);
393        }
394        index++;
395    }
396
397    for (auto &candidate : next_candidates)
398    {
399        delete[] candidate.code_points;
400        candidate.code_points = nullptr;
401    }
402
403    llama_grammar_free_impl(grammar);
404
405    return 0;
406}