Initial implementation for chat template parameters

Author: LysandreJik

src/lighteval/models/model_input.py

@@ -232,3 +232,17 @@ 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/transformers/transformers_model.py

@@ -230,7 +230,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/utils.py

@@ -34,8 +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 GenerationParameters, ChatTemplateParameters
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +81,7 @@ class ModelConfig(BaseModel, extra="forbid"):
     """
 
     generation_parameters: GenerationParameters = GenerationParameters()
+    chat_template_parameters: ChatTemplateParameters = ChatTemplateParameters()
     system_prompt: str | None = None
 
     @classmethod
@@ -131,20 +131,29 @@ 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(",")
+        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 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
 
 

src/lighteval/tasks/prompt_manager.py

@@ -28,6 +28,7 @@
 from itertools import cycle
 from typing import TYPE_CHECKING
 
+from lighteval.models.model_input import ChatTemplateParameters
 from lighteval.tasks.requests import Doc
 from lighteval.utils.utils import as_list
 
@@ -40,10 +41,17 @@
 
 
 class PromptManager:
-    def __init__(self, use_chat_template: bool = False, tokenizer=None, system_prompt: str | None = None):
+    def __init__(
+        self,
+        use_chat_template: bool = False,
+        tokenizer=None,
+        system_prompt: str | None = None,
+        chat_template_parameters: ChatTemplateParameters | None = None,
+    ):
         self.use_chat_template = use_chat_template
         self.tokenizer = tokenizer
         self.system_prompt = system_prompt  # System prompt to be used in chat templates
+        self.chat_template_parameters = chat_template_parameters if chat_template_parameters else {}
 
     def prepare_prompt(self, doc: Doc) -> str:
         """Prepare a prompt from a document, either using chat template or plain text format."""
@@ -123,6 +131,7 @@ def _prepare_chat_template(self, doc: Doc, tokenize: bool = True) -> str:
                 messages,
                 tokenize=False,
                 add_generation_prompt=True,
+                **self.chat_template_parameters.to_transformers_dict()
             )
 
         else:  # for apis