Add support for chat template parameters

Author: LysandreJik

src/lighteval/logging/info_loggers.py

@@ -87,6 +87,7 @@ class GeneralConfigLogger:
     model_size: str = None
 
     generation_parameters: dict | None = None
+    chat_template_parameters: dict | None = None
 
     # Nanotron config
     config: "Config" = None
@@ -129,7 +130,9 @@ def log_args_info(
         self.job_id = job_id
         self.config = config
 
-    def log_model_info(self, generation_parameters: dict, model_info: ModelInfo) -> None:
+    def log_model_info(
+        self, generation_parameters: dict, model_info: ModelInfo, chat_template_parameters: dict
+    ) -> None:
         """
         Logs the model information.
 
@@ -139,6 +142,7 @@ def log_model_info(self, generation_parameters: dict, model_info: ModelInfo) ->
 
         """
         self.generation_parameters = generation_parameters
+        self.chat_template_parameters = chat_template_parameters
         self.model_name = model_info.model_name
         self.model_sha = model_info.model_sha
         self.model_dtype = model_info.model_dtype

src/lighteval/main_baseline.py

@@ -89,6 +89,7 @@ def baseline(
             model_dtype=None,
             model_size=None,
         ),
+        {},
     )
     evaluation_tracker.task_config_logger.log(tasks_dict)
 

src/lighteval/models/custom/custom_model.py

@@ -70,5 +70,4 @@ def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
     An example of a custom model can be found in `examples/custom_models/google_translate_model.py`.
     """
 
-    model_name: str
     model_definition_file_path: str

src/lighteval/models/endpoints/endpoint_model.py

@@ -95,7 +95,6 @@ class ServerlessEndpointModelConfig(ModelConfig):
         ```
     """
 
-    model_name: str
     add_special_tokens: bool = True
     batch_size: int = 1
 

src/lighteval/models/litellm_model.py

@@ -94,7 +94,6 @@ class LiteLLMModelConfig(ModelConfig):
         ```
     """
 
-    model_name: str
     provider: str | None = None
     base_url: str | None = None
     api_key: str | None = None

src/lighteval/models/model_input.py

@@ -232,3 +232,18 @@ def to_sglang_dict(self) -> dict:
             "min_new_tokens": self.min_new_tokens,
         }
         return {k: v for k, v in args.items() if v is not None}
+
+
+class ChatTemplateParameters(BaseModel):
+    reasoning_effort: str = None
+
+    def to_transformers_dict(self) -> dict:
+        """Selects relevant chat template parameters for transformers models.
+
+        Returns:
+            dict: Valid parameters for the chat template
+        """
+        args = {
+            "reasoning_effort": self.reasoning_effort,
+        }
+        return {k: v for k, v in args.items() if v is not None}

src/lighteval/models/sglang/sglang_model.py

@@ -109,7 +109,6 @@ class SGLangModelConfig(ModelConfig):
         ```
     """
 
-    model_name: str
     load_format: str = "auto"
     dtype: str = "auto"
     tp_size: PositiveInt = 1  # how many GPUs to use for tensor parallelism

src/lighteval/models/transformers/transformers_model.py

@@ -133,7 +133,6 @@ class TransformersModelConfig(ModelConfig):
         (bitsandbytes for 4-bit/8-bit quantization).
     """
 
-    model_name: str
     tokenizer: str | None = None
     subfolder: str | None = None
     revision: str = "main"
@@ -230,7 +229,10 @@ def __init__(
         )
 
         self.prompt_manager = PromptManager(
-            use_chat_template=self.use_chat_template, tokenizer=self.tokenizer, system_prompt=config.system_prompt
+            use_chat_template=self.use_chat_template,
+            tokenizer=self.tokenizer,
+            system_prompt=config.system_prompt,
+            chat_template_parameters=config.chat_template_parameters,
         )
 
     def cleanup(self):

src/lighteval/models/transformers/vlm_transformers_model.py

@@ -104,7 +104,6 @@ class VLMTransformersModelConfig(ModelConfig):
             loading.
     """
 
-    model_name: str
     processor: str | None = None
     use_fast_image_processor: bool | None = None
     subfolder: str | None = None

src/lighteval/models/utils.py

@@ -34,7 +34,7 @@
 from transformers import AutoTokenizer
 from transformers.models.auto.configuration_auto import AutoConfig
 
-from lighteval.models.model_input import GenerationParameters
+from lighteval.models.model_input import ChatTemplateParameters, GenerationParameters
 
 
 logger = logging.getLogger(__name__)
@@ -70,7 +70,7 @@ class ModelConfig(BaseModel, extra="forbid"):
         config = ModelConfig.from_path("model_config.yaml")
 
         # Load from command line arguments
-        config = ModelConfig.from_args("model_name=meta-llama/Llama-3.1-8B-Instruct,system_prompt='You are a helpful assistant.',generation_parameters={temperature=0.7}")
+        config = ModelConfig.from_args("model_name=meta-llama/Llama-3.1-8B-Instruct,system_prompt='You are a helpful assistant.',generation_parameters={temperature:0.7}")
 
         # Direct instantiation
         config = ModelConfig(
@@ -81,7 +81,9 @@ class ModelConfig(BaseModel, extra="forbid"):
         ```
     """
 
+    model_name: str
     generation_parameters: GenerationParameters = GenerationParameters()
+    chat_template_parameters: ChatTemplateParameters = ChatTemplateParameters()
     system_prompt: str | None = None
 
     @classmethod
@@ -131,20 +133,31 @@ def _parse_args(args: str) -> dict:
         """
         # Looking for generation_parameters in the model_args
         generation_parameters_dict = None
-        pattern = re.compile(r"(\w+)=(\{.*\}|[^,]+)")
+        chat_template_parameters_dict = None
+        pattern = re.compile(r"(\w+)\s*=\s*(\{[^{}]*\}|[^,]+?)(?=,|$)")
         matches = pattern.findall(args)
         for key, value in matches:
             key = key.strip()
             if key == "generation_parameters":
                 gen_params = re.sub(r"(\w+):", r'"\1":', value)
                 generation_parameters_dict = json.loads(gen_params)
+            if key == "chat_template_parameters":
+                # Chat template parameters have strings as values that also need to be quoted
+                chat_template_params = re.sub(r"(\w+)\s*:\s*([A-Za-z_][\w.-]*)\s*(?=[,}])", r'"\1":"\2"', value)
+                chat_template_parameters_dict = json.loads(chat_template_params)
 
         args = re.sub(r"generation_parameters=\{.*?\},?", "", args).strip(",")
-        model_config = {k.split("=")[0]: k.split("=")[1] if "=" in k else True for k in args.split(",")}
+        args = re.sub(r"chat_template_parameters=\{.*?\},?", "", args).strip(",")
+        model_config = (
+            {k.split("=")[0]: k.split("=")[1] if "=" in k else True for k in args.split(",")} if args else {}
+        )
 
         if generation_parameters_dict is not None:
             model_config["generation_parameters"] = generation_parameters_dict
 
+        if chat_template_parameters_dict is not None:
+            model_config["chat_template_parameters"] = chat_template_parameters_dict
+
         return model_config