1  #define LLAMA_API_INTERNAL
  2  
  3  #include "grammar-parser.h"
  4  #include "ggml.h"
  5  #include "llama.h"
  6  #include "unicode.h"
  7  
  8  #include <cstdio>
  9  #include <cstdlib>
 10  #include <sstream>
 11  #include <fstream>
 12  #include <string>
 13  #include <vector>
 14  
 15  static bool llama_sample_grammar_string(struct llama_grammar * grammar, const std::string & input_str, size_t & error_pos, std::string & error_msg) {
 16      auto decoded = decode_utf8(input_str, {});
 17      const auto & code_points = decoded.first;
 18  
 19      size_t pos = 0;
 20      for (auto it = code_points.begin(), end = code_points.end() - 1; it != end; ++it) {
 21          auto prev_stacks = grammar->stacks;
 22          llama_grammar_accept(grammar->rules, prev_stacks, *it, grammar->stacks);
 23          if (grammar->stacks.empty()) {
 24              error_pos = pos;
 25              error_msg = "Unexpected character '" + unicode_cpt_to_utf8(*it) + "'";
 26              grammar->stacks = prev_stacks;
 27              return false;
 28          }
 29          ++pos;
 30      }
 31  
 32      for (const auto & stack : grammar->stacks) {
 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      // Parse the GBNF grammar
 84      auto parsed_grammar = grammar_parser::parse(grammar_str.c_str());
 85  
 86      // will be empty (default) if there are parse errors
 87      if (parsed_grammar.rules.empty()) {
 88          fprintf(stdout, "%s: failed to parse grammar\n", __func__);
 89          return 1;
 90      }
 91  
 92      // Ensure that there is a "root" node.
 93      if (parsed_grammar.symbol_ids.find("root") == parsed_grammar.symbol_ids.end()) {
 94          fprintf(stdout, "%s: grammar does not contain a 'root' symbol\n", __func__);
 95          return 1;
 96      }
 97  
 98      std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar.c_rules());
 99  
100      // Create the LLAMA grammar
101      auto grammar = llama_grammar_init(
102              grammar_rules.data(),
103              grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
104  
105      // Read the input file
106      std::string input_str;
107      {
108          std::ifstream input_file(input_filename);
109          GGML_ASSERT(input_file.is_open() && "Failed to open input file");
110          std::stringstream buffer;
111          buffer << input_file.rdbuf();
112          input_str = buffer.str();
113      }
114  
115      // Validate the input string against the grammar
116      size_t error_pos;
117      std::string error_msg;
118      bool is_valid = llama_sample_grammar_string(grammar, input_str, error_pos, error_msg);
119  
120      if (is_valid) {
121          fprintf(stdout, "Input string is valid according to the grammar.\n");
122      } else {
123          print_error_message(input_str, error_pos, error_msg);
124      }
125  
126      // Clean up
127      llama_grammar_free(grammar);
128  
129      return 0;
130  }