Fix TGI (Text Generation Inference) Endpoint Inference and TGI JSON Grammar Generation
#502
Conversation
cpcdoy
changed the title
TGI (Text Generation Inference) Endpoint InferenceTGI (Text Generation Inference) Endpoint Inference and TGI JSON Grammar Generation
|
Updated the PR to add support for JSON Grammar Constrained Generation for TGI |
| @@ -491,6 +492,7 @@ def _process_batch_logprob( | |||
| context=request.context if rolling else request.context + request.choice, | |||
| stop_tokens=[], | |||
| max_tokens=1, | |||
| grammar=request.generation_grammar, | |||
There was a problem hiding this comment.
why is the grammar in the request here while it is defined in the generation config ?
There was a problem hiding this comment.
Hey ! Sorry for the delay, so in your PR you are adding generation_grammar in the GenerationParameters and not in the requests. You would need to add a field in requests (defaults to None). Though maybe i'm missing something, did you test with your tasks and made sure it was using the correct grammar ?
There was a problem hiding this comment.
Hey @NathanHB, no worries! Yes, I've used my branch to run evaluations on several models and, just in case, I just reran the evaluation pipeline I wrote at the time (this is running on a base Qwen2.5-0.5B-Instruct not fine-tuned, just in case, so it's harder for it to follow my grammar to help demonstrate this), and TGI does receive my grammar in its logs like this:
2025-04-09T12:48:30.781244Z INFO compat_generate{default_return_full_text=true compute_type=Extension(ComputeType("1-nvidia-geforce-rtx-3060"))}:generate{parameters=GenerateParameters { best_of: None, temperature: None, repetition_penalty: None, frequency_penalty: None, top_k: None, top_p: None, typical_p: None, do_sample: false, max_new_tokens: Some(256), return_full_text: Some(false), stop: ["\n\n", "<|im_end|>"], truncate: None, watermark: false, details: true, decoder_input_details: true, seed: None, top_n_tokens: None, grammar: Some(Json(Object {"type": String("object"), "properties": Object {"reason": Object {"type": String("string"), "description": String("Reasoning for the classification")}, "classification": Object {"type": String("string"), "description": String("Banking transaction classification"), "enum": Array [String("label 0"), ..., String("label n"),]}, "confidence": Object {"type": String("string"), "enum": Array [String("high"), String("medium"), String("low")]}}, "required": Array [String("classification"), String("reason"), String("confidence")]})), adapter_id: None } total_time="1.430241619s" validation_time="1.408552ms" queue_time="107.404µs" inference_time="1.428725763s" time_per_token="18.085136ms" seed="None"}: text_generation_router::server: router/src/server.rs:422: Success
2025-04-09T12:48:30.787272Z INFO text_generation_router_v3::radix: backends/v3/src/radix.rs:108: Prefix 0 - Suffix 468
And the model's predictions look like this, which follow the grammar format perfectly:
"predictions":"['{\"reason\": \"...\", \"classification\": \"...\", \"confidence\": \"high\"}']
This is how I define the grammar and then I run it using the lighteval command on my script:
...
def get_bank_system_classification_grammar() -> TextGenerationInputGrammarType:
return TextGenerationInputGrammarType(
type="json",
value={
"type": "object",
"properties": {
"reason": {
"type": "string",
"description": "Reasoning for the classification",
},
"classification": {
"type": "string",
"description": "Banking transaction classification",
"enum": BANK_SYSTEM_LABELS,
},
"confidence": {"type": "string", "enum": ["high", "medium", "low"]},
},
"required": ["reason", "classification", "confidence"],
},
)
DATASET_DIR = "src/llm_tasks_eval/datasets/bank_system_classification_dataset"
BANK_SYSTEM_CLASSIFICATION_TASK = LightevalTaskConfig(
name="bank_system_classification",
prompt_function=prompt_bank_system_classification,
suite=["custom"],
hf_repo=DATASET_DIR,
hf_subset=None,
metric=[
bank_system_group,
semantic_similarity_metric,
detailed_classification_metric,
openai_comparison_metric,
],
generation_size=256,
generation_grammar=get_bank_system_classification_grammar(),
stop_sequence=["\n\n"],
trust_dataset=True,
evaluation_splits=["test"],
hf_avail_splits=["test"],
)
TASKS_TABLE = [BANK_SYSTEM_CLASSIFICATION_TASK]Note: I have censored/truncated multiple outputs and fields since they contain sensitive elements from work.
| generated_text = self.client.generate( | ||
| prompt=context, | ||
| do_sample=generation_config.do_sample or False, | ||
| max_new_tokens=generation_config.max_new_tokens, | ||
| best_of=generation_config.best_of, | ||
| repetition_penalty=generation_config.repetition_penalty, | ||
| return_full_text=generation_config.return_full_text or False, | ||
| seed=generation_config.seed, | ||
| stop_sequences=generation_config.stop, | ||
| temperature=generation_config.temperature, | ||
| top_k=generation_config.top_k, | ||
| top_p=generation_config.top_p, | ||
| truncate=generation_config.truncate, | ||
| typical_p=generation_config.typical_p, | ||
| watermark=generation_config.watermark or False, | ||
| decoder_input_details=generation_config.decoder_input_details, | ||
| grammar=generation_config.grammar, | ||
| ) |
There was a problem hiding this comment.
IIRC, this is the interface text-generation==0.7.0 is exposing now which is different since I have upgraded it from 0.6.0 in this PR.
Did you mean, is there a cleaner way to do this?
|
UP! I encountered a similar issue where the bug prevented us from using the TGI endpoint. The key issues I found are:
|
|
Exactly @naufalso , this is already solved in this PR! |
|
+1 is this going to be merged @NathanHB ? Would really like to use lighteval with locally hosted TGI, but I'm seeing the same |
Description
While implementing a custom task using
lighteval, I needed to use constrained grammar generation with TGI and it seems that TGI integration is not up-to-date and not working.Fixes for TGI Endpoint Inference
/inforoute of TGI3.0.1doesn't always return required fields such asmodel_dtype, so it was set toNoneby default if not found:AsyncClientfrom TGI has ageneratefunction that expects multiple parameters and not a structure.do_sample,return_full_textandwatermarkparameters asFalseby default since they come fromhuggingface_hubwhich accepts aNonedefault parameters but TGI doesn't accept them_async_process_requestanyway and maybe this should be fixed in another PR. Same foradapter_idfor LoRA heads.ModelClient's usage has been fixed to use theconfig: TGIModelConfigby default instead of named parametersFixes for TGI JSON Grammar Generation
text_generationto0.7.0Environment
Command
Dependencies
model_config_pathargument for TGItgi.yaml:Test Results
It works as can be seen from the logs.
TGI Logs with JSON Grammar Generation
Lighteval Logs
Note: I have anonymized parts of the logs