Skip to content

Fix TGI (Text Generation Inference) Endpoint Inference and TGI JSON Grammar Generation #502

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

#

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Fix TGI (Text Generation Inference) Endpoint Inference and TGI JSON Grammar Generation #502

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ab68a1b
fix: Lighteval communication with TGI
cpcdoy Jan 15, 2025
f442a29
fix: JSON grammar constrained generation
cpcdoy Jan 15, 2025
6bb6b13
Merge branch 'main' into fix/tgi_inference
cpcdoy Feb 7, 2025
cb13617
Merge branch 'main' into fix/tgi_inference
NathanHB Apr 8, 2025
749887e
Merge branch 'main' into fix/tgi_inference
cpcdoy Apr 9, 2025
55e3790
Merge branch 'main' into fix/tgi_inference
NathanHB Apr 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ dependencies = [

[project.optional-dependencies]
litellm = ["litellm", "diskcache"]
tgi = ["text-generation>=0.6.0"]
tgi = ["text-generation==0.7.0"]
optimum = ["optimum==1.12.0"]
quantization = ["bitsandbytes>=0.41.0", "auto-gptq>=0.4.2"]
adapters = ["peft==0.3.0"]
Expand Down
2 changes: 2 additions & 0 deletions src/lighteval/models/endpoints/endpoint_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,7 @@ async def _async_process_batch_logprob(
context=request.context if rolling else request.context + request.choice,
stop_tokens=[],
max_tokens=1,
grammar=request.generation_grammar,
)
for request in requests
]
Expand All @@ -446,6 +447,7 @@ def _process_batch_logprob(
context=request.context if rolling else request.context + request.choice,
stop_tokens=[],
max_tokens=1,
grammar=request.generation_grammar,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is the grammar in the request here while it is defined in the generation config ?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the feedback @NathanHB 🙏🏻

I did it similarly to how I've seen it done in the same file. See here for an example.
Potentially the usage could be improved in a follow-up PR.

Lmk if I'm also missing something on my end

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ?

Copy link
Author

@cpcdoy cpcdoy Apr 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)
for request in requests
]
Expand Down
21 changes: 19 additions & 2 deletions src/lighteval/models/endpoints/tgi_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def __init__(self, config: TGIModelConfig) -> None:

model_name = str(self.model_info["model_id"])
model_sha = self.model_info["model_sha"]
model_precision = self.model_info["model_dtype"]
model_precision = self.model_info.get("model_dtype")
self.model_info = ModelInfo(
model_name=model_name,
model_sha=model_sha,
Expand All @@ -105,7 +105,24 @@ def _async_process_request(
grammar=grammar,
)

generated_text = self.client.generate(prompt=context, generation_config=generation_config)
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,
)
Comment on lines +108 to +125
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this needed ?

Copy link
Author

@cpcdoy cpcdoy Feb 7, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?


return generated_text

Expand Down
4 changes: 2 additions & 2 deletions src/lighteval/models/model_input.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,9 @@ class GenerationParameters(BaseModel, extra="forbid"):
frequency_penalty: NonNegativeFloat | None = None # vllm, tgi, sglang
length_penalty: NonNegativeFloat | None = None # vllm, transformers
presence_penalty: NonNegativeFloat | None = None # vllm, sglang

max_new_tokens: NonNegativeInt | None = None # vllm, transformers, tgi, litellm, sglang
min_new_tokens: NonNegativeInt | None = None # vllm, transformers, sglang

grammar: str | None = None # tgi
seed: NonNegativeInt | None = None # vllm, tgi, litellm
stop_tokens: list[str] | None = None # vllm, transformers, tgi, litellm, sglang
temperature: NonNegativeFloat | None = None # vllm, transformers, tgi, litellm, sglang
Expand Down Expand Up @@ -208,6 +207,7 @@ def to_tgi_ie_dict(self) -> dict:
"top_k": self.top_k,
"top_p": self.top_p,
"truncate": self.truncate_prompt,
"grammar": self.grammar,
}
return {k: v for k, v in args.items() if v is not None}

Expand Down
4 changes: 1 addition & 3 deletions src/lighteval/models/model_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,7 @@ def load_model_with_tgi(config: TGIModelConfig):
raise ImportError(NO_TGI_ERROR_MSG)

logger.info(f"Load model from inference server: {config.inference_server_address}")
model = ModelClient(
address=config.inference_server_address, auth_token=config.inference_server_auth, model_id=config.model_id
)
model = ModelClient(config=config)
return model


Expand Down
Loading