Skip to content

llama : minor llama_grammar refactoring #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 17, 2024
Merged

llama : minor llama_grammar refactoring #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 4 additions & 8 deletions examples/gbnf-validator/gbnf-validator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,20 +11,15 @@
static bool llama_grammar_validate(struct llama_grammar * grammar, const std::string & input_str, size_t & error_pos, std::string & error_msg) {
const auto cpts = unicode_cpts_from_utf8(input_str);

const llama_grammar_rules & rules = llama_grammar_get_rules (grammar);
llama_grammar_stacks & stacks_cur = llama_grammar_get_stacks(grammar);
llama_grammar_stacks_cache & stacks_cache = llama_grammar_get_stacks_cache(grammar);
auto & stacks_cur = llama_grammar_get_stacks(grammar);

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

llama_grammar_accept(rules, stacks_prev, cpt, stacks_cur, stacks_cache);
llama_grammar_accept(grammar, cpt);

if (stacks_cur.empty()) {
error_pos = pos;
error_msg = "Unexpected character '" + unicode_cpt_to_utf8(cpt) + "'";
stacks_cur = stacks_prev;
return false;
}
++pos;
Expand Down Expand Up @@ -83,7 +78,8 @@ int main(int argc, char** argv) {

llama_grammar * grammar = llama_grammar_init_impl(nullptr, grammar_str.c_str(), "root");
if (grammar == nullptr) {
throw std::runtime_error("Failed to initialize llama_grammar");
fprintf(stdout, "Failed to initialize llama_grammar\n");
return 1;
}
// Read the input file
std::string input_str;
Expand Down
42 changes: 19 additions & 23 deletions src/llama-grammar.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -764,7 +764,7 @@ static void llama_grammar_advance_stack_memo(
if (it != stacks_cache.end()) {
advanced_stacks = it->second;
} else {
// Advance stacks with memorization
// Advance stacks with memorization
llama_grammar_advance_stack_memo_impl(rules, stack, advanced_stacks, stacks_cache);
stacks_cache.insert(make_pair(stack, advanced_stacks));
}
Expand Down Expand Up @@ -917,20 +917,11 @@ llama_grammar_stacks & llama_grammar_get_stacks(struct llama_grammar * grammar)
return grammar->stacks;
}

llama_grammar_stacks_cache & llama_grammar_get_stacks_cache(struct llama_grammar * grammar) {
return grammar->stacks_cache;
}

void llama_grammar_accept(
const llama_grammar_rules & rules,
const llama_grammar_stacks & stacks,
const uint32_t chr,
llama_grammar_stacks & stacks_new,
llama_grammar_stacks_cache & stacks_cache) {
stacks_new.clear();
stacks_new.reserve(stacks.size());
void llama_grammar_accept(struct llama_grammar * grammar, uint32_t chr) {
llama_grammar_stacks stacks_new;
stacks_new.reserve(grammar->stacks.size());

for (const auto & stack : stacks) {
for (const auto & stack : grammar->stacks) {
if (stack.empty()) {
continue;
}
Expand All @@ -944,9 +935,11 @@ void llama_grammar_accept(
if (!llama_grammar_is_end_of_sequence(pos)) {
new_stack.push_back(pos);
}
llama_grammar_advance_stack_memo(rules, new_stack, stacks_new, stacks_cache);
llama_grammar_advance_stack_memo(grammar->rules, new_stack, stacks_new, grammar->stacks_cache);
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I really like the signature simplification that you made for llama_grammar_accept -- passing around the grammar structure alone (rather than all of the individual pieces of the structure) feels good.

Should we do that to llama_grammar_advance_stack / llama_grammar_advance_stack_memo as well?

I.E. something more like:

static void llama_grammar_advance_stack_memo(
        struct llama_grammar * grammar,
        const llama_grammar_stack  & stack,
              llama_grammar_stacks & new_stacks) {

?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should give it a try. I think I briefly tried it again, but there was some issue. Will retake a look tomorrow and see if we can simplify further.

}
}

grammar->stacks = std::move(stacks_new);
}

llama_grammar_candidates llama_grammar_reject_candidates_for_stack(
Expand Down Expand Up @@ -1062,7 +1055,7 @@ struct llama_grammar * llama_grammar_init_impl(
// Important: vec_rules has to be moved here, not copied, because stacks contains
// pointers to elements of vec_rules. If vec_rules were copied into llama_grammar
// then the pointers would be invalidated when the local vec_rules goes out of scope.
return new llama_grammar { vocab, std::move(vec_rules), std::move(stacks), {}, std::move(stacks_cache), };
return new llama_grammar { vocab, std::move(vec_rules), std::move(stacks), std::move(stacks_cache), {}, };
}

struct llama_grammar * llama_grammar_init_impl(const struct llama_vocab * vocab, const char * grammar_str, const char * grammar_root) {
Expand Down Expand Up @@ -1141,7 +1134,7 @@ struct llama_grammar * llama_grammar_init_impl(const struct llama_vocab * vocab,
// Important: vec_rules has to be moved here, not copied, because stacks contains
// pointers to elements of vec_rules. If vec_rules were copied into llama_grammar
// then the pointers would be invalidated when the local vec_rules goes out of scope.
return new llama_grammar { vocab, std::move(vec_rules), std::move(stacks), {}, std::move(stacks_cache), };
return new llama_grammar { vocab, std::move(vec_rules), std::move(stacks), std::move(stacks_cache), {}, };
}

void llama_grammar_free_impl(struct llama_grammar * grammar) {
Expand All @@ -1153,15 +1146,21 @@ void llama_grammar_free_impl(struct llama_grammar * grammar) {
}

struct llama_grammar * llama_grammar_clone_impl(const struct llama_grammar & grammar) {
llama_grammar * result = new llama_grammar { grammar.vocab, grammar.rules, grammar.stacks, grammar.partial_utf8, };
llama_grammar * result = new llama_grammar {
grammar.vocab,
grammar.rules,
grammar.stacks,
grammar.stacks_cache,
grammar.partial_utf8,
};

// redirect elements in stacks to point to new rules
for (size_t is = 0; is < result->stacks.size(); is++) {
for (size_t ie = 0; ie < result->stacks[is].size(); ie++) {
for (size_t ir0 = 0; ir0 < grammar.rules.size(); ir0++) {
for (size_t ir1 = 0; ir1 < grammar.rules[ir0].size(); ir1++) {
if (grammar.stacks[is][ie] == &grammar.rules[ir0][ir1]) {
result->stacks[is][ie] = &result->rules[ir0][ir1];
result->stacks[is][ie] = &result->rules[ir0][ir1];
}
}
}
Expand Down Expand Up @@ -1228,11 +1227,8 @@ void llama_grammar_accept_impl(struct llama_grammar & grammar, llama_token token
const auto decoded = decode_utf8(piece, grammar.partial_utf8);
const auto & code_points = decoded.first;

llama_grammar_stacks stacks_new;

for (auto it = code_points.begin(), end = code_points.end() - 1; it != end; ++it) {
llama_grammar_accept(grammar.rules, grammar.stacks, *it, stacks_new, grammar.stacks_cache);
grammar.stacks = std::move(stacks_new);
llama_grammar_accept(&grammar, *it);
}

grammar.partial_utf8 = decoded.second;
Expand Down
14 changes: 5 additions & 9 deletions src/llama-grammar.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,20 +71,15 @@ struct VectorPointerHash {

using llama_grammar_stacks_cache = std::unordered_map<llama_grammar_stack, llama_grammar_stacks, VectorPointerHash>;

// TODO: remove, needed for tests atm
const llama_grammar_rules & llama_grammar_get_rules (const struct llama_grammar * grammar);
llama_grammar_stacks & llama_grammar_get_stacks( struct llama_grammar * grammar);
llama_grammar_stacks_cache & llama_grammar_get_stacks_cache( struct llama_grammar * grammar);

// takes a set of possible pushdown stacks on a grammar, which are required to
// be positioned at a character range (see `llama_grammar_advance_stack`), and
// produces the N possible stacks if the given char is accepted at those
// positions
void llama_grammar_accept(
const llama_grammar_rules & rules,
const llama_grammar_stacks & stacks,
uint32_t chr,
llama_grammar_stacks & stacks_new,
llama_grammar_stacks_cache & stacks_cache);
void llama_grammar_accept(struct llama_grammar * grammar, uint32_t chr);

std::vector<llama_grammar_candidate> llama_grammar_reject_candidates_for_stack(
const llama_grammar_rules & rules,
Expand Down Expand Up @@ -128,10 +123,11 @@ struct llama_grammar {
const llama_grammar_rules rules; // TODO: shared ptr
llama_grammar_stacks stacks;

// buffer for partially generated UTF-8 sequence from accepted tokens
llama_partial_utf8 partial_utf8;
// cache N possible stacks from a stack
llama_grammar_stacks_cache stacks_cache;

// buffer for partially generated UTF-8 sequence from accepted tokens
llama_partial_utf8 partial_utf8;
};

//
Expand Down
10 changes: 3 additions & 7 deletions tests/test-grammar-integration.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,10 @@ static bool test_build_grammar_fails(const std::string & grammar_str) {
static bool match_string(const std::string & input, llama_grammar * grammar) {
const auto cpts = unicode_cpts_from_utf8(input);

const llama_grammar_rules & rules = llama_grammar_get_rules (grammar);
llama_grammar_stacks & stacks_cur = llama_grammar_get_stacks(grammar);
llama_grammar_stacks_cache & stacks_cache = llama_grammar_get_stacks_cache(grammar);
auto & stacks_cur = llama_grammar_get_stacks(grammar);

for (const auto & cpt : cpts) {
const llama_grammar_stacks stacks_prev = llama_grammar_get_stacks(grammar); // copy

llama_grammar_accept(rules, stacks_prev, cpt, stacks_cur, stacks_cache);
llama_grammar_accept(grammar, cpt);

if (stacks_cur.empty()) {
// no stacks means that the grammar failed to match at this point
Expand All @@ -64,7 +60,7 @@ static void test(const std::string & test_desc, const std::string & grammar_str,
auto * grammar = build_grammar(grammar_str);

// Save the original grammar stacks so that we can reset after every new string we want to test
const llama_grammar_stacks stacks_org = llama_grammar_get_stacks(grammar);
const llama_grammar_stacks stacks_org = llama_grammar_get_stacks(grammar); // copy

llama_grammar_stacks & stacks_cur = llama_grammar_get_stacks(grammar);

Expand Down
6 changes: 2 additions & 4 deletions tests/test-llama-grammar.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,12 +113,10 @@ int main()
}
}

llama_grammar * grammar = NULL;
std::vector<const llama_grammar_element *> grammar_rules(parsed_grammar.c_rules());

grammar = llama_grammar_init_impl(nullptr, grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
if (grammar == nullptr)
{
llama_grammar * grammar = llama_grammar_init_impl(nullptr, grammar_rules.data(), grammar_rules.size(), parsed_grammar.symbol_ids.at("root"));
if (grammar == nullptr) {
throw std::runtime_error("Failed to initialize llama_grammar");
}

Expand Down