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

* fix: Lighteval communication with TGI

* fix: JSON grammar constrained generation

* fix: unit tests +
add: dep in extra

* fix: request var => doc var after refactor

* fix: update test to support the new grammar field

* fix: TGI endpoint with the new refactor

* update: TGI model config in examples with the latest parameters

* add: example custom task on a classification dataset to demonstrate the usage of constrained grammar generation using TGI

* add: format example task

* fix: unit test

* add: adapt the  in the yaml config to use  similarly to the other endpoints

* clean: moved new task to community_tasks

* fix: format

* clean: delete unused grammar field

* del: grammar

* add: copyright at the top

* del: langcodes dep isn't needed anymore

* add: use load from file directly in the main endpoint

* del: newlines

* add: mock HTTP request for info to TGI server

---------

Co-authored-by: Nathan Habib <30601243+NathanHB@users.noreply.github.com>

Author: cpcdoy

community_tasks/custom_task_classification_grammar_task.py

@@ -1,4 +1,6 @@
 model_parameters:
-  inference_server_address: ""
+  inference_server_address: "http://localhost:8080" # Replace with your actual TGI server address
   inference_server_auth: null
   model_name: null # Optional, only required if the TGI container was launched with model_id pointing to a local directory
+  generation_parameters:
+    temperature: 0.1

examples/model_configs/tgi_model.yaml

@@ -85,8 +85,8 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-litellm = ["litellm", "diskcache"]
-tgi = ["text-generation>=0.6.0"]
+litellm = ["litellm[caching]", "diskcache"]
+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"]

pyproject.toml

@@ -249,11 +249,8 @@ def tgi(
     """
     Evaluate models using TGI as backend.
     """
-    import yaml
-
     from lighteval.logging.evaluation_tracker import EvaluationTracker
     from lighteval.models.endpoints.tgi_model import TGIModelConfig
-    from lighteval.models.model_input import GenerationParameters
     from lighteval.pipeline import ParallelismManager, Pipeline, PipelineParameters
 
     evaluation_tracker = EvaluationTracker(
@@ -269,11 +266,7 @@ def tgi(
 
     parallelism_manager = ParallelismManager.TGI
 
-    with open(model_config_path, "r") as f:
-        config = yaml.safe_load(f)
-
-    generation_parameters = GenerationParameters(**config.get("generation", {}))
-    model_config = TGIModelConfig(**config["model"], generation_parameters=generation_parameters)
+    model_config = TGIModelConfig.from_path(model_config_path)
 
     pipeline_params = PipelineParameters(
         launcher_type=parallelism_manager,

src/lighteval/main_endpoint.py

@@ -527,6 +527,7 @@ async def _async_process_batch_logprob(self, docs: list[Doc], rolling: bool = Fa
                     context=context if rolling else context + doc.choices[0],
                     stop_tokens=[],
                     max_tokens=1,
+                    grammar=doc.generation_grammar,
                 )
                 for context, doc in zip(contexts, docs)
             ]
@@ -539,6 +540,7 @@ def _process_batch_logprob(self, docs: list[Doc], rolling: bool = False) -> list
                 context=context if rolling else context + doc.choices[0],
                 stop_tokens=[],
                 max_tokens=1,
+                grammar=doc.generation_grammar,
             )
             for context, doc in zip(contexts, docs)
         ]

src/lighteval/models/endpoints/endpoint_model.py

@@ -30,6 +30,8 @@
 
 from lighteval.models.abstract_model import ModelConfig
 from lighteval.models.endpoints.endpoint_model import InferenceEndpointModel
+from lighteval.tasks.prompt_manager import PromptManager
+from lighteval.utils.cache_management import SampleCache
 from lighteval.utils.imports import NO_TGI_ERROR_MSG, is_tgi_available
 
 
@@ -87,6 +89,7 @@ class TGIModelConfig(ModelConfig):
     inference_server_auth: str | None = None
     model_name: str | None
     model_info: dict | None = None
+    batch_size: int = 1
 
 
 # inherit from InferenceEndpointModel instead of LightevalModel since they both use the same interface, and only overwrite
@@ -110,12 +113,23 @@ def __init__(self, config: TGIModelConfig) -> None:
             raise ValueError("Error occurred when fetching info: " + str(self.model_info))
         if config.model_name:
             self.model_info["model_id"] = config.model_name
+        else:
+            # Set the model_name in config to the actual model_id from server for caching
+            config.model_name = self.model_info["model_id"]
         self.config = config
         self._tokenizer = AutoTokenizer.from_pretrained(self.model_info["model_id"])
         self._add_special_tokens = True
         self.use_async = True
         self.config.model_info = self.model_info
 
+        # Initialize prompt manager (required by parent class)
+        self.prompt_manager = PromptManager(
+            use_chat_template=True, tokenizer=self.tokenizer, system_prompt=config.system_prompt
+        )
+
+        # Initialize cache for tokenization and predictions
+        self._cache = SampleCache(config)
+
     def _async_process_request(
         self,
         context: str,
@@ -134,7 +148,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,
+        )
 
         return generated_text
 

src/lighteval/models/endpoints/tgi_model.py

@@ -109,9 +109,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
 
 

src/lighteval/models/model_loader.py

@@ -251,6 +251,12 @@ def _get_docs_from_split(self, splits: list[str], few_shots=False) -> list[Doc]:
                 item["__index"] = ix
                 doc = self.formatter(item, self.name)
                 doc.id = str(ix)
+
+                # Transfer task-level generation parameters to the document
+                doc.generation_grammar = self.generation_grammar
+                doc.generation_size = self.generation_size
+                doc.stop_sequences = self.stop_sequence
+
                 docs.append(doc)
 
         return docs

src/lighteval/tasks/lighteval_task.py

@@ -33,11 +33,12 @@ class TestTGIModelConfig:
             (
                 "examples/model_configs/tgi_model.yaml",
                 {
-                    "inference_server_address": "",
+                    "inference_server_address": "http://localhost:8080",
                     "inference_server_auth": None,
                     "model_name": None,
                     "model_info": None,
                     "system_prompt": None,
+                    "batch_size": 1,
                     "generation_parameters": {
                         "block_size": None,
                         "num_blocks": None,
@@ -52,7 +53,7 @@ class TestTGIModelConfig:
                         "repetition_penalty": None,
                         "seed": None,
                         "stop_tokens": None,
-                        "temperature": 0,
+                        "temperature": 0.1,
                         "top_k": None,
                         "top_p": None,
                         "truncate_prompt": None,

tests/models/endpoints/test_tgi_model.py

@@ -219,21 +219,27 @@ def test_cache_vllm(self, mock_create_model, mock_greedy_until, mock_loglikeliho
 
             self._test_cache(model)
 
+    @patch("requests.get")
     @patch("lighteval.models.endpoints.tgi_model.ModelClient._greedy_until")
     @patch("lighteval.models.endpoints.tgi_model.ModelClient._loglikelihood")
-    def test_cache_tgi(self, mock_greedy_until, mock_loglikelihood):
+    def test_cache_tgi(self, mock_loglikelihood, mock_greedy_until, mock_requests_get):
         from lighteval.models.endpoints.tgi_model import ModelClient, TGIModelConfig
         from lighteval.utils.imports import is_tgi_available
 
         if not is_tgi_available():
             pytest.skip("Skipping because missing the imports")
 
         # Mock TGI requests
-        mock_greedy_until.return_value = self.model_responses
         mock_loglikelihood.return_value = self.model_responses
+        mock_greedy_until.return_value = self.model_responses
+
+        # Mock HTTP info request
+        mock_requests_get.return_value.json.return_value = {"model_id": "Qwen/Qwen3-0.6B"}
 
         with tempfile.TemporaryDirectory() as temp_dir:
-            config = TGIModelConfig(model_name="Qwen/Qwen3-0.6B", cache_dir=temp_dir)
+            config = TGIModelConfig(
+                model_name="Qwen/Qwen3-0.6B", cache_dir=temp_dir, inference_server_address="http://localhost:8080"
+            )
             model = ModelClient(config)
 
             self._test_cache(model)