Custom LogitsProcessor/LogitsWarper

[loganlebanoff]

Feature request

I have my own implementation of a LogitsWarper in this format https://huggingface.co/docs/transformers/internal/generation_utils#transformers.LogitsProcessor that I would like to add to TGI. There are already warpers for temperature, top_p, etc., but I want to add my own. Note: I don't want my LogitsWarper to be merged in to the repo, since I have a vey specific use case -- I just want to know how it can be done.

Motivation

I want to have the benefits of TGI like batching and tensor parallelism while also adding custom LogitsWarper functionality.

Your contribution

I've tried editing the source files at text-generation-inference/server/text_generation_server/utils/logits_process.py etc., but I started going down a rabbit hole of files that I needed to edit.

[abhinavkulkarni]

Yes, this will be very much appreciated. This would be very useful to force the model to generate a JSON response.

Something like this would be quite useful if and when TGI decided to add support for grammar based sampling. Recently llama.cpp added support for the same.

[abhinavkulkarni]

@Narsil, @OlivierDehaene: Some guidance on this would be very much appreciated. Wether this would be a desired feature and if so, where to start in the codebase.

Thanks!

[mjsteele12]

@abhinavkulkarni I almost added grammar support as an idea in discussions before I found this thread/your comment. I am extremely interested in merging the outlines project with TGI. +1 to this idea/functionality!

[abhinavkulkarni]

Thanks @mjsteele12, here's similar project: https://github.com/Shopify/torch-grammar

It would be good to get some guidance from TGI maintainers as to how to go about this.

[abhinavkulkarni]

Here's how Llama.cpp does it. More BNF grammars from here.

[cnewton]

I'd love to have this as a feature as well @Narsil @OlivierDehaene

fwiw, vllm has an open pull request for essentially the same feature: vllm-project/vllm#535

[tripathiarpan20]

@abhinavkulkarni @mjsteele12 The idea to add custom grammar for generations in TGI is of great interest to me too, any updates on this? any way I can help?

[loganlebanoff]

I got a custom LogitsWarper working! Here's the changes I made, I hope it can help anyone else who is working on something similar.
main...loganlebanoff:text-generation-inference:main

[OlivierDehaene]

Hey!
I'm OOL on this.
What would be the minimal API change that you guys would need? If we only add logit_bias, would that be enough?

[abhinavkulkarni]

@loganlebanoff: Thanks for the great work!

Just to let you know, according to the previous thread, TGI maintainers were not interested in exposing internal token-level logic to the end user. For e.g., many users have no idea about the tokenizer vocabulary, the only interface TGI exposes is to input prompt text and get output completion text. Users are not supposed to provide token-level logit biases.

Hence it may be best to expose a functionality to take a custom grammar (in BNF or EBNF form). The grammar module can then bias the logits so as to produce an output consistent with the input grammar.

CC: @OlivierDehaene, @Narsil

[OlivierDehaene]

What's the standard? BNF? EBNF? What about Guidance?

[abhinavkulkarni]

@OlivierDehaene: I'm light on the details about Guidance, but it should be possible to integrate an EBNF grammar parser and state machine (SM) in TGI. There is a PR currently open on transformers repo (link) that adds a parser and SM in a custom logits processor.

The user simply needs to provide an EBNF grammar file.

Recently, llama.cpp integrated this functionality (link).

[lapp0]

This functionality is critical for function calling, generating code and data serialization formats such as json, yaml, and csv, along with providing other application-specific constraints on the generation.

Some commentary I've found on logit processing:

Logit processing in the context of returning top_n_tokens resulted in a performance degradation of 25% per #617 (comment)

(Although my understanding is that in certain cases grammar-constrained sampling would improve speed, as the candidate space is reduced.)

Breaking handling logit processing via a series of smaller requests for the model is marked as undesirable by Narsil #810 (comment)

Guidance is marked as a priority for the next release #505 (comment) however this was at the time of of v1.0.3, so I'm not sure whether this is still considered a priority 3 releases later.

Thanks for your hard work on TGI, would be great to hear some input on how this might be integrated.

[comaniac]

Upvote for this. Now OpenAI API allows users specify JSON response format. I believe they achieved this by controlling logit processor to follow JSON grammar (i.e., constraint semantic decoding). This, however, cannot be added to TGI for now due to the lack of custom logic processor.

[abhinavkulkarni]

I tried integrating code from this transformers PR and was able to successfully generate JSON responses from TGI.

However, this has no support for full set of Unicode characters.