epfl-dlab / epfl-dlab/transformers-CFG
beam search doesn't work with transformers_cfg
- 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.