epfl-dlab / epfl-dlab/transformers-CFG

beam search doesn't work with transformers_cfg

Open
#9 7 comments 2 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
139
Forks
23
PR merge metrics
No merged PRs in 30d

Description

```
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers_cfg.grammar_utils import IncrementalGrammarConstraint
from transformers_cfg.generation.logits_process import GrammarConstrainedLogitsProcessor

if __name__ == "__main__":
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained("gpt2")

# Load json grammar
with open("examples/grammars/json.ebnf", "r") as file:
grammar_str = file.read()
grammar = IncrementalGrammarConstraint(grammar_str, "root", tokenizer)
grammar_processor = GrammarConstrainedLogitsProcessor(grammar)

# Generate
prefix1 = "This is a valid json string for http request:"
prefix2 = "This is a valid json string for shopping cart:"
input_ids = tokenizer([prefix1, prefix2], add_special_tokens=False, return_tensors="pt", padding=True)["input_ids"]

output = model.generate(
input_ids,
do_sample=False,
max_length=50,
num_beams=1, #this can't be >1
logits_processor=[grammar_processor],
repetition_penalty=5.0,
num_return_sequences=1,
)





```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.