EuroEval 15.6.1__py3-none-any.whl → 15.7.1__py3-none-any.whl

euroeval/benchmark_modules/litellm.py

@@ -1,11 +1,8 @@
 """Generative models from an inference API, using the LiteLLM framework."""
 
 import collections.abc as c
-import itertools as it
-import json
 import logging
 import os
-import random
 import re
 import typing as t
 from functools import cached_property, partial
@@ -33,6 +30,7 @@ from litellm.exceptions import (
 )
 from litellm.llms.vertex_ai.common_utils import VertexAIError
 from litellm.types.utils import ChoiceLogprobs, ModelResponse
+from pydantic import conlist, create_model
 from requests.exceptions import RequestException
 from tqdm.auto import tqdm
 from transformers.trainer import Trainer
@@ -59,6 +57,7 @@ from ..exceptions import (
     NeedsEnvironmentVariable,
     NeedsExtraInstalled,
 )
+from ..generation_utils import apply_prompt, extract_few_shot_examples
 from ..task_group_utils import (
     question_answering,
     sequence_classification,
@@ -104,6 +103,7 @@ MODEL_MAX_LENGTH_MAPPING = {
     r"o1-(mini|preview)(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 128_000,
     r"o1(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 200_000,
     r"o[2-9](-mini|-preview)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 200_000,
+    r"gpt-4.1.*": 1_047_576,
     # Anthropic models
     r"(anthropic/)?claude-[1-9](-[1-9])?-(opus|sonnet|haiku)-[0-9]{8}": 200_000,
     # Gemini models
@@ -135,20 +135,23 @@ ALLOWED_PARAMS = {
     r"gpt-4.*": [],
     r"o[1-9](-mini|-preview)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": ["low", "high"],
     # Anthropic models
-    r"(anthropic/)?claude-3-.*": [],
-    r"(anthropic/)?claude-3.5-.*": [],
-    r"(anthropic/)?claude-3.7-sonnet.*": ["thinking"],
+    r"(anthropic/)?claude-3-(haiku|sonnet|opus).*": [],
+    r"(anthropic/)?claude-3-5-.*": [],
+    r"(anthropic/)?claude-3-7-sonnet.*": ["thinking"],
     # Gemini models
     r"(gemini/)?gemini-.*": [],
     # xAI models
-    r"(xai/)?grok.*": [],
+    r"(xai/)?grok-2.*": [],
+    r"(xai/)?grok-3(-fast)?(-beta)?": [],
+    r"(xai/)?grok-3-mini(-fast)?(-beta)?": ["low", "high"],
 }
 
 
 REASONING_MODELS = [
     r"o[1-9](-mini|-preview)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?",
     r"(gemini/)?gemini.*thinking.*",
-    r"(gemini/)?gemini-2.5-pro.*",
+    r"(gemini/)?gemini-2.5.*",
+    r"(xai/)?grok-3-mini.*",
 ]
 
 
@@ -190,7 +193,10 @@ class LiteLLMModel(BenchmarkModule):
         )
 
         self.buffer["first_label_token_mapping"] = get_first_label_token_mapping(
-            dataset_config=self.dataset_config, tokenizer=None
+            dataset_config=self.dataset_config,
+            model_config=self.model_config,
+            tokenizer=None,
+            generative_type=self.generative_type,
         )
 
     @property
@@ -201,13 +207,20 @@ class LiteLLMModel(BenchmarkModule):
             The generative type of the model, or None if it has not been set yet.
         """
         if self.model_config.revision == "thinking":
-            return GenerativeType.REASONING
+            type_ = GenerativeType.REASONING
         elif re.fullmatch(
             pattern="|".join(REASONING_MODELS), string=self.model_config.model_id
         ):
-            return GenerativeType.REASONING
+            type_ = GenerativeType.REASONING
         else:
-            return GenerativeType.INSTRUCTION_TUNED
+            type_ = GenerativeType.INSTRUCTION_TUNED
+
+        log_once(
+            f"Detected generative type {type_.name!r} for model "
+            f"{self.model_config.model_id!r}",
+            level=logging.DEBUG,
+        )
+        return type_
 
     def generate(self, inputs: dict) -> GenerativeModelOutput:
         """Generate outputs from the model.
@@ -243,7 +256,10 @@ class LiteLLMModel(BenchmarkModule):
         # Get the mapping from labels to the first token in the label. We call this each
         # time we generate a new dataset since the dataset config can change
         self.buffer["first_label_token_mapping"] = get_first_label_token_mapping(
-            dataset_config=self.dataset_config, tokenizer=None
+            dataset_config=self.dataset_config,
+            model_config=self.model_config,
+            tokenizer=None,
+            generative_type=self.generative_type,
         )
 
         if self.buffer["first_label_token_mapping"]:
@@ -254,16 +270,41 @@ class LiteLLMModel(BenchmarkModule):
             assert "json" in messages[0]["content"].lower(), (
                 "Prompt must contain 'json' for JSON tasks."
             )
-            generation_kwargs["response_format"] = dict(type="json_object")
-            log_once(
-                "Enabling JSON response format for model "
-                f"{self.model_config.model_id!r}",
-                level=logging.DEBUG,
-            )
+            if self.generative_type == GenerativeType.REASONING:
+                log_once(
+                    f"The model {self.model_config.model_id!r} is a reasoning model "
+                    "and thus does not support structured generation, so we do not "
+                    "enable it.",
+                    level=logging.DEBUG,
+                )
+            elif litellm.utils.supports_response_schema(
+                model=self.model_config.model_id
+            ):
+                ner_tag_names = list(self.dataset_config.prompt_label_mapping.values())
+                keys_and_their_types: dict[str, t.Any] = {
+                    tag_name: (conlist(str, max_length=5), ...)
+                    for tag_name in ner_tag_names
+                }
+                pydantic_class = create_model("AnswerFormat", **keys_and_their_types)
+                generation_kwargs["response_format"] = pydantic_class
+                log_once(
+                    "Enabling structured generation for model "
+                    f"{self.model_config.model_id!r} with the JSON schema "
+                    f"{pydantic_class.model_json_schema()}",
+                    level=logging.DEBUG,
+                )
+            else:
+                generation_kwargs["response_format"] = dict(type="json_object")
+                log_once(
+                    "Enabling structured JSON generation for model "
+                    f"{self.model_config.model_id!r} with no custom JSON schema, as "
+                    "the model does not support schemas.",
+                    level=logging.DEBUG,
+                )
 
         if self.model_config.revision == "thinking":
             generation_kwargs["thinking"] = dict(
-                type="enabled", budget_tokens=REASONING_MAX_TOKENS
+                type="enabled", budget_tokens=REASONING_MAX_TOKENS - 1
             )
             log_once(
                 f"Enabling thinking mode for model {self.model_config.model_id!r}",
@@ -280,28 +321,42 @@ class LiteLLMModel(BenchmarkModule):
         # This drops generation kwargs that are not supported by the model
         litellm.drop_params = True
 
+        # Error messages that we want to catch and handle
+        stop_messages = ["stop_sequences", "'stop' is not supported with this model"]
+        logprobs_messages = [
+            "you are not allowed to request logprobs",
+            "you've reached the maximum number of requests with logprobs",
+            "logprobs is not supported",
+            "logprobs is not enabled",
+        ]
+        temperature_messages = [
+            "'temperature' is not supported with this model.",
+            "temperature is not supported with this model",
+        ]
+        temperature_must_be_one_messages = [
+            "`temperature` may only be set to 1",
+            "'temperature' does not support 0.0 with this model. Only the default "
+            "(1) value is supported",
+        ]
+        max_items_messages = ["'maxItems' is not permitted."]
+        no_json_schema_messages = ["Property keys should match pattern"]
+
         # Extract the generated sequences from the model response. Some APIs cannot
         # handle using newlines as stop sequences, so we try both.
         num_attempts = 10
         for _ in range(num_attempts):
-            stop_messages = ["stop_sequences"]
-            logprobs_messages = [
-                "you are not allowed to request logprobs",
-                "you've reached the maximum number of requests with logprobs",
-                "logprobs is not supported",
-                "logprobs is not enabled",
-            ]
-            temperature_messages = [
-                "'temperature' is not supported with this model.",
-                "temperature is not supported with this model",
-            ]
             try:
-                model_response = litellm.completion(
-                    messages=messages, max_retries=3, **generation_kwargs
+                model_response = litellm.completion_with_retries(
+                    messages=messages, **generation_kwargs
                 )
                 break
             except (BadRequestError, RateLimitError) as e:
                 if any(msg.lower() in str(e).lower() for msg in stop_messages):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} does not support "
+                        "stop sequences, so disabling them.",
+                        level=logging.DEBUG,
+                    )
                     generation_kwargs["stop"] = None
                 elif (
                     any(msg.lower() in str(e).lower() for msg in logprobs_messages)
@@ -310,10 +365,55 @@ class LiteLLMModel(BenchmarkModule):
                     # we ignore this since the rate limiting makes it unusable anyway.
                     or (isinstance(e, VertexAIError) and "logprobs" in str(e).lower())
                 ):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} does not support "
+                        "logprobs, so disabling it.",
+                        level=logging.DEBUG,
+                    )
                     generation_kwargs.pop("logprobs")
                     generation_kwargs.pop("top_logprobs")
                 elif any(msg.lower() in str(e).lower() for msg in temperature_messages):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} does not support "
+                        "temperature, so disabling it.",
+                        level=logging.DEBUG,
+                    )
                     generation_kwargs.pop("temperature")
+                elif any(
+                    msg.lower() in str(e).lower()
+                    for msg in temperature_must_be_one_messages
+                ):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} requires "
+                        "temperature to be set to 1, so setting it.",
+                        level=logging.DEBUG,
+                    )
+                    generation_kwargs["temperature"] = 1.0
+                elif any(msg.lower() in str(e).lower() for msg in max_items_messages):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} does not support "
+                        "maxItems in the JSON schema, so disabling it.",
+                        level=logging.DEBUG,
+                    )
+                    ner_tag_names = list(
+                        self.dataset_config.prompt_label_mapping.values()
+                    )
+                    keys_and_their_types = {
+                        tag_name: (list[str], ...) for tag_name in ner_tag_names
+                    }
+                    pydantic_class = create_model(
+                        "AnswerFormat", **keys_and_their_types
+                    )
+                    generation_kwargs["response_format"] = pydantic_class
+                elif any(
+                    msg.lower() in str(e).lower() for msg in no_json_schema_messages
+                ):
+                    log_once(
+                        f"The model {self.model_config.model_id!r} does not support "
+                        "JSON schemas, so using the vanilla JSON format.",
+                        level=logging.DEBUG,
+                    )
+                    generation_kwargs["response_format"] = dict(type="json_object")
                 elif isinstance(e, RateLimitError):
                     raise InvalidModel(
                         "You have encountered your rate limit for model "
@@ -332,6 +432,7 @@ class LiteLLMModel(BenchmarkModule):
                 Timeout,
                 ServiceUnavailableError,
                 InternalServerError,
+                SystemError,
             ) as e:
                 logger.debug(
                     f"Service temporarily unavailable. The error message was: {e}. "
@@ -359,9 +460,11 @@ class LiteLLMModel(BenchmarkModule):
                 "reasoning. Returning an empty string."
             )
             return GenerativeModelOutput(sequences=[""])
+
         model_response_choices = model_response.choices[0]
         assert isinstance(model_response_choices, litellm.Choices)
-        generation_output = model_response_choices.message["content"] or ""
+        generated_message: litellm.Message = model_response_choices.message
+        generation_output = generated_message.content or ""
         generation_output = generation_output.strip()
 
         # Structure the model output as a GenerativeModelOutput object
@@ -838,14 +941,22 @@ class LiteLLMModel(BenchmarkModule):
             )
 
         if self.benchmark_config.few_shot:
-            few_shot_examples = self._extract_few_shot_examples(
-                dataset=dataset, task=task, itr_idx=itr_idx
+            few_shot_examples = extract_few_shot_examples(
+                dataset=dataset, dataset_config=self.dataset_config, itr_idx=itr_idx
             )
         else:
             few_shot_examples = list()
 
         dataset["test"] = dataset["test"].map(
-            partial(self._apply_prompt, few_shot_examples=few_shot_examples, task=task),
+            partial(
+                apply_prompt,
+                few_shot_examples=few_shot_examples,
+                model_config=self.model_config,
+                dataset_config=self.dataset_config,
+                instruction_model=True,
+                always_populate_text_field=False,
+                tokenizer=None,
+            ),
             batched=True,
             load_from_cache_file=False,
             keep_in_memory=True,
@@ -853,253 +964,6 @@ class LiteLLMModel(BenchmarkModule):
 
         return dataset
 
-    def _extract_few_shot_examples(
-        self, dataset: DatasetDict, task: Task, itr_idx: int
-    ) -> list[dict[str, t.Any]]:
-        """Extract few-shot examples from a dataset.
-
-        This will always extract the examples from the training split.
-
-        We ensure that the few-shot examples are unique by picking them one at a time.
-
-        Args:
-            dataset:
-                The dataset to extract the few-shot examples from.
-            task:
-                The task that is being benchmarked.
-            itr_idx:
-                The index of the dataset in the iterator.
-
-        Returns:
-            The few-shot examples.
-        """
-        random_seed = 4242 + itr_idx
-        num_few_shots = self.dataset_config.num_few_shot_examples
-        few_shot_examples: list[dict[str, t.Any]] = list()
-        shuffled_train = dataset["train"].shuffle(seed=random_seed)
-
-        match task.task_group:
-            case (
-                TaskGroup.SEQUENCE_CLASSIFICATION
-                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
-            ):
-                labels = it.cycle(self.dataset_config.labels)
-                while (
-                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
-                ):
-                    label = next(labels)
-                    possible_examples = shuffled_train.filter(
-                        lambda x: x["label"].lower() == label.lower()
-                    )
-                    if len(possible_examples) == 0:
-                        continue
-                    example = possible_examples.select(range(1))[0]
-                    few_shot_examples.append(example)
-                    shuffled_train = shuffled_train.filter(
-                        lambda x: x["text"] != example["text"]
-                    )
-
-            case TaskGroup.TEXT_TO_TEXT:
-                while (
-                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
-                ):
-                    example = shuffled_train.select(range(1))[0]
-                    few_shot_examples.append(example)
-                    shuffled_train = shuffled_train.filter(
-                        lambda x: x["text"] != example["text"]
-                    )
-
-            case TaskGroup.TOKEN_CLASSIFICATION:
-                labels = it.cycle(
-                    [
-                        label.lower()
-                        for label in self.dataset_config.labels
-                        if label.lower().startswith("b-")
-                    ]
-                )
-                while (
-                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
-                ):
-                    label = next(labels)
-                    possible_examples = shuffled_train.filter(
-                        lambda x: label in [tag.lower() for tag in x["labels"]]
-                    )
-                    if len(possible_examples) == 0:
-                        continue
-                    example = possible_examples.select(range(1))[0]
-                    few_shot_examples.append(example)
-                    shuffled_train = shuffled_train.filter(
-                        lambda x: x["tokens"] != example["tokens"]
-                    )
-
-            case TaskGroup.QUESTION_ANSWERING:
-                # Locate the maximum number of tokens that constitutes a short example
-                for max_num_tokens in [512, 1024, 2048, 4096, 8192]:
-                    train_with_short_examples = dataset["train"].filter(
-                        lambda example: len(example["context"]) < max_num_tokens
-                    )
-                    num_short_examples = len(train_with_short_examples)
-                    if num_short_examples >= self.dataset_config.num_few_shot_examples:
-                        break
-                else:
-                    raise InvalidBenchmark(
-                        "Could not find enough short examples for few-shot learning."
-                    )
-
-                shuffled_train = train_with_short_examples.shuffle(seed=random_seed)
-                while (
-                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
-                ):
-                    example = shuffled_train.select(range(1))[0]
-                    few_shot_examples.append(example)
-                    shuffled_train = shuffled_train.filter(
-                        lambda x: x["context"] != example["context"]
-                    )
-
-            case _:
-                raise NotImplementedError(f"Unsupported task group: {task.task_group}.")
-
-        random.seed(random_seed)
-        random.shuffle(few_shot_examples)
-        return few_shot_examples
-
-    def _apply_prompt(
-        self,
-        examples: dict[str, t.Any],
-        few_shot_examples: list[dict[str, t.Any]],
-        task: Task,
-    ) -> dict[str, t.Any]:
-        """Apply prompt template to an example, potentially with few-shot examples.
-
-        Args:
-            examples:
-                The examples to apply the few-shot examples to.
-            few_shot_examples:
-                The few-shot examples to apply.
-            task:
-                The task that is being benchmarked.
-
-        Returns:
-            The example with the few-shot examples applied.
-        """
-
-        def create_prompt(**kwargs: str) -> tuple[str, str]:
-            """Create a prompt from the given keyword arguments.
-
-            Args:
-                kwargs:
-                    The keyword arguments to use in the prompt.
-
-            Returns:
-                A pair (prompt, label), where "label" is an empty string if the model is
-                not instruction tuned (as in this case it is included in the prompt).
-            """
-            label_key = "label" if "label" in kwargs else "target_text"
-            label = kwargs.pop(label_key)
-            label_mapping = self.dataset_config.prompt_label_mapping
-            label = label_mapping.get(label, label)
-            prompt = self.dataset_config.instruction_prompt.format(**kwargs)
-            return prompt, label
-
-        match task.task_group:
-            case (
-                TaskGroup.SEQUENCE_CLASSIFICATION
-                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
-            ):
-                few_shot_sections = [
-                    create_prompt(
-                        text=example["text"].replace("\n", " ").strip(),
-                        label=example["label"].replace("\n", " ").strip(),
-                    )
-                    for example in few_shot_examples
-                ]
-                new_sections = [
-                    create_prompt(text=text.replace("\n", " ").strip(), label="")
-                    for text in examples["text"]
-                ]
-
-            case TaskGroup.TEXT_TO_TEXT:
-                few_shot_sections = [
-                    create_prompt(
-                        text=example["text"].replace("\n", " ").strip(),
-                        target_text=example["target_text"].replace("\n", " ").strip(),
-                    )
-                    for example in few_shot_examples
-                ]
-                new_sections = [
-                    create_prompt(text=text.replace("\n", " ").strip(), target_text="")
-                    for text in examples["text"]
-                ]
-
-            case TaskGroup.TOKEN_CLASSIFICATION:
-
-                def create_label(example: dict) -> str:
-                    prompt_labels = self.dataset_config.prompt_label_mapping.values()
-                    labels: dict[str, list[str]] = {
-                        prompt_label: list() for prompt_label in prompt_labels
-                    }
-                    for token, label in zip(example["tokens"], example["labels"]):
-                        label = label.lower()
-                        if label == "o":
-                            continue
-                        prompt_label = self.dataset_config.prompt_label_mapping[label]
-                        if label.startswith("b-"):
-                            labels[prompt_label].append(token)
-                        elif label.startswith("i-"):
-                            labels[prompt_label][-1] += " " + token
-                    return json.dumps(labels, ensure_ascii=False)
-
-                few_shot_sections = [
-                    create_prompt(
-                        text=" ".join(example["tokens"]).replace("\n", " ").strip(),
-                        label=create_label(example=example),
-                    )
-                    for example in few_shot_examples
-                ]
-                new_sections = [
-                    create_prompt(
-                        text=" ".join(tokens).replace("\n", " ").strip(), label=""
-                    )
-                    for tokens in examples["tokens"]
-                ]
-
-            case TaskGroup.QUESTION_ANSWERING:
-                few_shot_sections = [
-                    create_prompt(
-                        text=example["context"].replace("\n", " ").strip(),
-                        question=example["question"].replace("\n", " ").strip(),
-                        label=example["answers"]["text"][0].replace("\n", " "),
-                    )
-                    for example in few_shot_examples
-                ]
-                new_sections = [
-                    create_prompt(
-                        text=context.replace("\n", " ").strip(),
-                        question=question.replace("\n", " ").strip(),
-                        label="",
-                    )
-                    for context, question in zip(
-                        examples["context"], examples["question"]
-                    )
-                ]
-
-            case _:
-                raise NotImplementedError(f"Unsupported task group: {task.task_group}.")
-
-        few_shot_messages = [
-            dict(role=role, content=content)
-            for prompt, label in few_shot_sections
-            for role, content in [("user", prompt), ("assistant", label)]
-        ]
-
-        messages_list = [
-            few_shot_messages + [dict(role="user", content=prompt)]
-            for prompt, _ in new_sections
-        ]
-
-        examples["messages"] = messages_list
-        return examples
-
 
 def raise_if_wrong_params(
     model_config: ModelConfig, allowed_params: dict[str, list[str]]