llama

1
#include "unicode.h"
2
#include "llama-grammar.h"
3

4
#include <cstdio>
5
#include <cstdlib>
6
#include <sstream>
7
#include <fstream>
8
#include <string>
9
#include <vector>
10

11
static bool llama_grammar_validate(struct llama_grammar * grammar, const std::string & input_str, size_t & error_pos, std::string & error_msg) {
12
    const auto cpts = unicode_cpts_from_utf8(input_str);
13

14
    const llama_grammar_rules  & rules      = llama_grammar_get_rules (grammar);
15
          llama_grammar_stacks & stacks_cur = llama_grammar_get_stacks(grammar);
16

17
    size_t pos = 0;
18
    for (const auto & cpt : cpts) {
19
        const llama_grammar_stacks stacks_prev = llama_grammar_get_stacks(grammar); // copy
20

21
        llama_grammar_accept(rules, stacks_prev, cpt, stacks_cur);
22

23
        if (stacks_cur.empty()) {
24
            error_pos = pos;
25
            error_msg = "Unexpected character '" + unicode_cpt_to_utf8(cpt) + "'";
26
            stacks_cur = stacks_prev;
27
            return false;
28
        }
29
        ++pos;
30
    }
31

32
    for (const auto & stack : stacks_cur) {
33
        if (stack.empty()) {
34
            return true;
35
        }
36
    }
37

38
    error_pos = pos;
39
    error_msg = "Unexpected end of input";
40
    return false;
41
}
42

43
static void print_error_message(const std::string & input_str, size_t error_pos, const std::string & error_msg) {
44
    fprintf(stdout, "Input string is invalid according to the grammar.\n");
45
    fprintf(stdout, "Error: %s at position %zu\n", error_msg.c_str(), error_pos);
46
    fprintf(stdout, "\n");
47
    fprintf(stdout, "Input string:\n");
48
    fprintf(stdout, "%s", input_str.substr(0, error_pos).c_str());
49
    if (error_pos < input_str.size()) {
50
        fprintf(stdout, "\033[1;31m%c", input_str[error_pos]);
51
        if (error_pos+1 < input_str.size()) {
52
            fprintf(stdout, "\033[0;31m%s", input_str.substr(error_pos+1).c_str());
53
        }
54
        fprintf(stdout, "\033[0m\n");
55
    }
56
}
57

58
int main(int argc, char** argv) {
59
    if (argc != 3) {
60
        fprintf(stdout, "Usage: %s <grammar_filename> <input_filename>\n", argv[0]);
61
        return 1;
62
    }
63

64
    const std::string grammar_filename = argv[1];
65
    const std::string input_filename = argv[2];
66

67
    // Read the GBNF grammar file
68
    FILE* grammar_file = fopen(grammar_filename.c_str(), "r");
69
    if (!grammar_file) {
70
        fprintf(stdout, "Failed to open grammar file: %s\n", grammar_filename.c_str());
71
        return 1;
72
    }
73

74
    std::string grammar_str;
75
    {
76
        std::ifstream grammar_file(grammar_filename);
77
        GGML_ASSERT(grammar_file.is_open() && "Failed to open grammar file");
78
        std::stringstream buffer;
79
        buffer << grammar_file.rdbuf();
80
        grammar_str = buffer.str();
81
    }
82

83
    llama_grammar * grammar = llama_grammar_init_impl(nullptr, grammar_str.c_str(), "root");
84
    if (grammar == nullptr) {
85
        throw std::runtime_error("Failed to initialize llama_grammar");
86
    }
87
    // Read the input file
88
    std::string input_str;
89
    {
90
        std::ifstream input_file(input_filename);
91
        GGML_ASSERT(input_file.is_open() && "Failed to open input file");
92
        std::stringstream buffer;
93
        buffer << input_file.rdbuf();
94
        input_str = buffer.str();
95
    }
96

97
    // Validate the input string against the grammar
98
    size_t error_pos;
99
    std::string error_msg;
100
    bool is_valid = llama_grammar_validate(grammar, input_str, error_pos, error_msg);
101

102
    if (is_valid) {
103
        fprintf(stdout, "Input string is valid according to the grammar.\n");
104
    } else {
105
        print_error_message(input_str, error_pos, error_msg);
106
    }
107

108
    // Clean up
109
    llama_grammar_free_impl(grammar);
110

111
    return 0;
112
}
113