Codeprocastinator/optimized-tinyllama-covalent

1#ifdef NDEBUG2#undef NDEBUG3#endif4 5#include "llama.h"6#include "llama-grammar.h"7 8#include <cassert>9#include <stdexcept>10 11int main()12{13    llama_grammar_parser parsed_grammar;14 15    std::vector<std::pair<std::string, uint32_t>> expected = {16        {"expr", 2},17        {"expr_6", 6},18        {"expr_7", 7},19        {"ident", 8},20        {"ident_10", 10},21        {"num", 9},22        {"num_11", 11},23        {"root", 0},24        {"root_1", 1},25        {"root_5", 5},26        {"term", 4},27        {"ws", 3},28        {"ws_12", 12},29    };30 31    std::vector<std::vector<llama_grammar_element>> expected_rules = {32        {{LLAMA_GRETYPE_RULE_REF, 5}, {LLAMA_GRETYPE_END, 0}},33        {34            {LLAMA_GRETYPE_RULE_REF, 2},35            {LLAMA_GRETYPE_CHAR, 61},36            {LLAMA_GRETYPE_RULE_REF, 3},37            {LLAMA_GRETYPE_RULE_REF, 4},38            {LLAMA_GRETYPE_CHAR, 10},39            {LLAMA_GRETYPE_END, 0},40        },41        {{LLAMA_GRETYPE_RULE_REF, 4}, {LLAMA_GRETYPE_RULE_REF, 7}, {LLAMA_GRETYPE_END, 0}},42        {{LLAMA_GRETYPE_RULE_REF, 12}, {LLAMA_GRETYPE_END, 0}},43        {44            {LLAMA_GRETYPE_RULE_REF, 8},45            {LLAMA_GRETYPE_ALT, 0},46            {LLAMA_GRETYPE_RULE_REF, 9},47            {LLAMA_GRETYPE_ALT, 0},48            {LLAMA_GRETYPE_CHAR, 40},49            {LLAMA_GRETYPE_RULE_REF, 3},50            {LLAMA_GRETYPE_RULE_REF, 2},51            {LLAMA_GRETYPE_CHAR, 41},52            {LLAMA_GRETYPE_RULE_REF, 3},53            {LLAMA_GRETYPE_END, 0},54        },55        {{LLAMA_GRETYPE_RULE_REF, 1}, {LLAMA_GRETYPE_RULE_REF, 5}, {LLAMA_GRETYPE_ALT, 0}, {LLAMA_GRETYPE_RULE_REF, 1}, {LLAMA_GRETYPE_END, 0}},56        {57            {LLAMA_GRETYPE_CHAR, 45},58            {LLAMA_GRETYPE_CHAR_ALT, 43},59            {LLAMA_GRETYPE_CHAR_ALT, 42},60            {LLAMA_GRETYPE_CHAR_ALT, 47},61            {LLAMA_GRETYPE_RULE_REF, 4},62            {LLAMA_GRETYPE_END, 0},63        },64        {{LLAMA_GRETYPE_RULE_REF, 6}, {LLAMA_GRETYPE_RULE_REF, 7}, {LLAMA_GRETYPE_ALT, 0}, {LLAMA_GRETYPE_END, 0}},65        {66            {LLAMA_GRETYPE_CHAR, 97},67            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 122},68            {LLAMA_GRETYPE_RULE_REF, 10},69            {LLAMA_GRETYPE_RULE_REF, 3},70            {LLAMA_GRETYPE_END, 0},71        },72        {{LLAMA_GRETYPE_RULE_REF, 11}, {LLAMA_GRETYPE_RULE_REF, 3}, {LLAMA_GRETYPE_END, 0}},73        {74            {LLAMA_GRETYPE_CHAR, 97},75            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 122},76            {LLAMA_GRETYPE_CHAR_ALT, 48},77            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},78            {LLAMA_GRETYPE_CHAR_ALT, 95},79            {LLAMA_GRETYPE_RULE_REF, 10},80            {LLAMA_GRETYPE_ALT, 0},81            {LLAMA_GRETYPE_END, 0},82        },83        {84            {LLAMA_GRETYPE_CHAR, 48},85            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},86            {LLAMA_GRETYPE_RULE_REF, 11},87            {LLAMA_GRETYPE_ALT, 0},88            {LLAMA_GRETYPE_CHAR, 48},89            {LLAMA_GRETYPE_CHAR_RNG_UPPER, 57},90            {LLAMA_GRETYPE_END, 0},91        },92        {93            {LLAMA_GRETYPE_CHAR, 32},94            {LLAMA_GRETYPE_CHAR_ALT, 9},95            {LLAMA_GRETYPE_CHAR_ALT, 10},96            {LLAMA_GRETYPE_RULE_REF, 12},97            {LLAMA_GRETYPE_ALT, 0},98            {LLAMA_GRETYPE_END, 0},99        },100    };101 102    for (auto pair : expected)103    {104        parsed_grammar.symbol_ids[pair.first] = pair.second;105    }106 107    for (auto rule : expected_rules)108    {109        parsed_grammar.rules.emplace_back();110        for (auto element : rule)111        {112            parsed_grammar.rules.back().push_back(element);113        }114    }115 116    std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar.c_rules());117 118    llama_grammar * grammar = llama_grammar_init_impl(nullptr, grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));119    if (grammar == nullptr) {120        throw std::runtime_error("Failed to initialize llama_grammar");121    }122 123    std::vector<std::vector<llama_grammar_element>> expected_stacks = {124        {125            {LLAMA_GRETYPE_RULE_REF, 5},126            {LLAMA_GRETYPE_CHAR, 61},127            {LLAMA_GRETYPE_RULE_REF, 7},128            {LLAMA_GRETYPE_CHAR, 97},129        },130        {131            {LLAMA_GRETYPE_RULE_REF, 5},132            {LLAMA_GRETYPE_CHAR, 61},133            {LLAMA_GRETYPE_RULE_REF, 7},134            {LLAMA_GRETYPE_RULE_REF, 3},135            {LLAMA_GRETYPE_CHAR, 48},136        },137        {138            {LLAMA_GRETYPE_RULE_REF, 5},139            {LLAMA_GRETYPE_CHAR, 61},140            {LLAMA_GRETYPE_RULE_REF, 7},141            {LLAMA_GRETYPE_RULE_REF, 3},142            {LLAMA_GRETYPE_CHAR, 48},143        },144        {145            {LLAMA_GRETYPE_RULE_REF, 5},146            {LLAMA_GRETYPE_CHAR, 61},147            {LLAMA_GRETYPE_RULE_REF, 7},148            {LLAMA_GRETYPE_CHAR, 40},149        },150        {151            {LLAMA_GRETYPE_CHAR, 61},152            {LLAMA_GRETYPE_RULE_REF, 7},153            {LLAMA_GRETYPE_CHAR, 97},154        },155        {156            {LLAMA_GRETYPE_CHAR, 61},157            {LLAMA_GRETYPE_RULE_REF, 7},158            {LLAMA_GRETYPE_RULE_REF, 3},159            {LLAMA_GRETYPE_CHAR, 48},160        },161        {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_CHAR, 61},169            {LLAMA_GRETYPE_RULE_REF, 7},170            {LLAMA_GRETYPE_CHAR, 40},171        }};172 173    auto index = 0;174    for (const llama_grammar_stack & stack : llama_grammar_get_stacks(grammar))175    {176        // compare stack to expected_stack177        for (uint32_t i = 0; i < stack.size(); i++)178        {179            const llama_grammar_element * element = stack[i];180            const llama_grammar_element & expected_element = expected_stacks[index][i];181 182            // pretty print error message before asserting183            if (expected_element.type != element->type || expected_element.value != element->value)184            {185                fprintf(stderr, "index: %d\n", index);186                fprintf(stderr, "expected_element: %d, %u\n", expected_element.type, expected_element.value);187                fprintf(stderr, "actual_element: %d, %u\n", element->type, element->value);188                fprintf(stderr, "expected_element != actual_element\n");189            }190 191            assert(expected_element.type == element->type && expected_element.value == element->value);192        }193        index++;194    }195 196    std::vector<llama_grammar_candidate> next_candidates;197    next_candidates.resize(24);198 199    for (size_t i = 0; i < 24; ++i)200    {201        uint32_t *cp = new uint32_t[2]; // dynamically allocate memory for code_point202        cp[0] = 37 + i;203        cp[1] = 0;204        next_candidates[i] = {i, cp, {}};205    }206 207    std::vector<std::vector<std::pair<uint32_t, uint16_t>>> expected_reject = {208        {209            {0, 37},210            {1, 38},211            {2, 39},212            {3, 40},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            {4, 41},271            {5, 42},272            {6, 43},273            {7, 44},274            {8, 45},275            {9, 46},276            {10, 47},277            {11, 48},278            {12, 49},279            {13, 50},280            {14, 51},281            {15, 52},282            {16, 53},283            {17, 54},284            {18, 55},285            {19, 56},286            {20, 57},287            {21, 58},288            {22, 59},289            {23, 60},290        },291        {292            {0, 37},293            {1, 38},294            {2, 39},295            {3, 40},296            {4, 41},297            {5, 42},298            {6, 43},299            {7, 44},300            {8, 45},301            {9, 46},302            {10, 47},303            {11, 48},304            {12, 49},305            {13, 50},306            {14, 51},307            {15, 52},308            {16, 53},309            {17, 54},310            {18, 55},311            {19, 56},312            {20, 57},313            {21, 58},314            {22, 59},315            {23, 60},316        },317        {318            {0, 37},319            {1, 38},320            {2, 39},321            {3, 40},322            {4, 41},323            {5, 42},324            {6, 43},325            {7, 44},326            {8, 45},327            {9, 46},328            {10, 47},329            {21, 58},330            {22, 59},331            {23, 60},332        },333        {334            {0, 37},335            {1, 38},336            {2, 39},337            {3, 40},338            {4, 41},339            {5, 42},340            {6, 43},341            {7, 44},342            {8, 45},343            {9, 46},344            {10, 47},345            {21, 58},346            {22, 59},347            {23, 60},348        },349        {350            {0, 37},351            {1, 38},352            {2, 39},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            {23, 60},373        },374    };375 376    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);377 378    std::vector<std::vector<llama_grammar_candidate>> all_rejects;379 380    for (std::size_t count = 0; count < llama_grammar_get_stacks(grammar).size(); ++count)381    {382        rejects = llama_grammar_reject_candidates_for_stack(llama_grammar_get_rules(grammar), llama_grammar_get_stacks(grammar)[count], next_candidates);383        all_rejects.push_back(rejects);384    }385 386    index = 0;387    for (auto rej : all_rejects)388    {389        for (uint32_t i = 0; i < rej.size(); i++)390        {391            auto element = rej[i];392            auto expected_element = expected_reject[index][i];393            assert(element.index == expected_element.first && *element.code_points == expected_element.second);394        }395        index++;396    }397 398    for (auto &candidate : next_candidates)399    {400        delete[] candidate.code_points;401        candidate.code_points = nullptr;402    }403 404    llama_grammar_free_impl(grammar);405 406    return 0;407}408