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

pyproject.toml

@@ -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"]

src/lighteval/models/endpoints/endpoint_model.py

@@ -478,6 +478,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
             ]
@@ -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,
             )
             for request in requests
         ]

src/lighteval/models/endpoints/tgi_model.py

@@ -101,7 +101,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,
@@ -127,7 +127,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/model_input.py

@@ -42,6 +42,7 @@ class GenerationParameters:
     min_p: Optional[float] = None  # vllm, transformers
     top_p: Optional[int] = None  # vllm, transformers, tgi
     truncate_prompt: Optional[bool] = None  # vllm, tgi
+    grammar: Optional[str] = None  # tgi
 
     @classmethod
     def from_dict(cls, config_dict: dict):
@@ -117,5 +118,6 @@ 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}

src/lighteval/models/model_loader.py

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