[WIP] Add Support for masking output using a Context-Free-Grammar
Author: jvhoffbauer
Pull request: yes
Labels: name: WIP

# What does this PR do?
Fixes #25778 

# Review 
@gante @LysandreJik as discussed. Also, @oobabooga @ArthurZucker @jorgemcgomes feel free to look at this. 
This is in large parts inspired by [rellm](https://github.com/r2d4/rellm/tree/main) and [parserllm](https://github.com/r2d4/parserllm) by @r2d4 so I am tagging him here too for visibility.

# Discussion 
This is totally WIP still. I added a notebook showcasing how it might work when using the generation API from a low-level perspective. Please educate me if that is ok that way. Long-term I want to add tests. 
Generally, we use Lark to parse the grammar. The grammar itself would most likely come from a text file or a string but Lark also has a grammar object that one can use. I built a class [CfgParsingStepper](https://github.com/jvhoffbauer/transformers/blob/cfg_masking_logits_processor/src/transformers/generation/logits_process.py#L1721) that we can use to get the state of the parser for any input string. It will give us the terminal symbol we are processing and the regex for this symbol that we can use to find valid tokens. The [CfgLogitsProcessor](https://github.com/jvhoffbauer/transformers/blob/cfg_masking_logits_processor/src/transformers/generation/logits_process.py#L1774) gets the state every turn for every beam. We can consider persisting the state during generation which might be faster, instead of recalculating it. 

The whole codebase is still very much WIP. Most importantly, we still need to 

- Add tests (instead of a Jupyter Notebook) 
- Add error handling to make sure the parser throws an error if it is started with an invalid input 
- Connect the logits processor to the actual generation APIs. I am thinking of something like `model.generate(..., grammar=Grammar(...))`
- Think about cases where you want the grammar constraint to start only for the new generations and not for the input prompt

Anyhow, I wanted to put this out there for discussion. Let me know if it makes sense this way and whether you would like me to continue working in this way or if I should change anything! 

Comments:
Thanks for your PR! I'm prompting @gante for review once he's back from leave :)

I appreciate your effort and patience @jvhoffbauer!

> Thanks for your PR! I'm prompting @gante for review once he's back from leave :)
> 
> I appreciate your effort and patience @jvhoffbauer!

Hey @LysandreJik, did you already have time to review? 

My 2 cents: it would be desirable to have compatibility with the BNF grammars in the llama.cpp repository:

https://github.com/ggerganov/llama.cpp/tree/master/grammars

Existing work in this direction (using custom `LogitsProcessor` in the transformers library):

https://github.com/Shopify/torch-grammar

https://github.com/im-not-tom/text-generation-webui-output-template-extension/

The original llama.cpp PR:

https://github.com/ggerganov/llama.cpp/pull/1773

Hey @jvhoffbauer, I will let @gante that just came back from holiday review as soon as he has a spare cycle; he's the owner of `generate` and should be able to provide a much better review than I can :)

Thanks for your contribution

(it's not forgotten, it's in my queue -- I should be able to review it over the new few days)

Awesome! It’s just a draft. Please let me know your thoughts and I will focus on wrapping it up to a complete PR. Potentially already over the coming week.

@gante no worries, I am still here! 

I understand you do not want to make the generate function dependent on a CPU-bound tokenizer which I get. The way you explain the approach makes total sense and I trust it is the best way of integrating this functionality into `transformers`. 

Does it make sense if I start drafting out the code that could be used for such an article? 

@jvhoffbauer absolutely! You can start by drafting a stand-alone `.py` file (easier to see the diff and iterate than a notebook :) ), and when we're happy with it, I will share further pointers on integrating it into the documentation.

After this is done, I would love to invite you to write a community blog post explaining the virtues of context-free-grammar and to share a space with this technique! 💛 

Sounds great! I will prepare a draft and we can iterate. 

Is it a bit overkill to introduce the dependency of Lark to implement the grammar-constrained decoding feature ? 
It seems the llama-cpp repo simply made a standalone implementation ?  @gante  https://github.com/ggerganov/llama.cpp/tree/master/grammars



@Saibo-creator It's okay, since `lark` is stable and has no dependencies of its own :)

For this feature, I think there are several aspects to take into account while implementing:
1. how does this work with different tokenizers(bpe, unigram, wordpiece etc). For example, LLAMA-CPP's implementation is specific to llama tokenizer, I don't think it will work for wordpiece. 
2. unicode support. This is a bit tricky but important for multilingual usage, llama-cpp 's first implementation didn't take care of that and later they spotted and fixed it . https://github.com/ggerganov/llama.cpp/issues/2501
3.  possible comptatbility with other sampling methods such as top-k sampling
4. [future] comptaiblity with utf-16 ? 

> My 2 cents: it would be desirable to have compatibility with the BNF grammars in the llama.cpp repository:
> 
> https://github.com/ggerganov/llama.cpp/tree/master/grammars
> 
> Existing work in this direction (using custom `LogitsProcessor` in the transformers library):
> 
> https://github.com/Shopify/torch-grammar
> 
> https://github.com/im-not-tom/text-generation-webui-output-template-extension/
> 
> The original llama.cpp PR:
> 
> [ggerganov/llama.cpp#1773](https://github.com/ggerganov/llama.cpp/pull/1773)

See  https://github.com/huggingface/transformers/pull/27557 for an implementation compatible with llama-cpp

Is anyone still working on this? I need this functionality to test an idea. I can either do a quick fork for my own use or spend a little longer and make something worth sharing.

> Is anyone still working on this? I need this functionality to test an idea. I can either do a quick fork for my own use or spend a little longer and make something worth sharing.

Hey @RadixSeven, due to the complexity of this feature, the HF team has decided it would be best to transition it to a separate project instead of integrating it. 

Feel free to check out https://github.com/epfl-dlab/transformers-CFG, which was created from this PR.