EuroEval 15.6.0__py3-none-any.whl → 15.7.0__py3-none-any.whl

euroeval/benchmark_modules/litellm.py

@@ -33,6 +33,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
@@ -104,6 +105,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 +137,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 +195,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 +209,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 +258,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 +272,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 +323,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 +367,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 +434,7 @@ class LiteLLMModel(BenchmarkModule):
                 Timeout,
                 ServiceUnavailableError,
                 InternalServerError,
+                SystemError,
             ) as e:
                 logger.debug(
                     f"Service temporarily unavailable. The error message was: {e}. "
@@ -359,9 +462,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

euroeval/benchmark_modules/vllm.py

@@ -132,7 +132,7 @@ class VLLMModel(HuggingFaceEncoderModel):
         self._model: LLM = model
         self._tokenizer: PreTrainedTokenizer = tokenizer
         self.end_of_reasoning_token_id = get_end_of_reasoning_token_id(
-            model=self._model, tokenizer=self._tokenizer
+            model=self._model, tokenizer=self._tokenizer, model_id=model_config.model_id
         )
 
         # We specify `HuggingFaceEncoderModel` here instead of `VLLMModel`, as we want
@@ -146,7 +146,10 @@ class VLLMModel(HuggingFaceEncoderModel):
         self.buffer |= dict(
             instruction_model=self._tokenizer.chat_template is not None,
             first_label_token_mapping=get_first_label_token_mapping(
-                dataset_config=self.dataset_config, tokenizer=self._tokenizer
+                dataset_config=self.dataset_config,
+                model_config=self.model_config,
+                tokenizer=self._tokenizer,
+                generative_type=self.generative_type,
             ),
         )
         if self.model_config.adapter_base_model_id is not None:
@@ -332,30 +335,40 @@ class VLLMModel(HuggingFaceEncoderModel):
             if end_of_chat_token:
                 stop_tokens.append(end_of_chat_token)
 
+        logits_processor = None
         if self.dataset_config.task in TASKS_USING_JSON:
-            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)
-            logits_processor = JSONLogitsProcessor(
-                schema=pydantic_class,
-                tokenizer=adapt_tokenizer(tokenizer=self._tokenizer),  # type: ignore
-                whitespace_pattern=r" ?",
-            )
-            log_once(
-                "Using structured generation with the schema "
-                f"{pydantic_class.model_json_schema()}",
-                level=logging.DEBUG,
-            )
-        else:
-            logits_processor = None
+            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,
+                )
+            else:
+                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)
+                logits_processor = JSONLogitsProcessor(
+                    schema=pydantic_class,
+                    tokenizer=adapt_tokenizer(tokenizer=self._tokenizer),  # type: ignore
+                    whitespace_pattern=r" ?",
+                )
+                log_once(
+                    "Using structured generation with the JSON schema "
+                    f"{pydantic_class.model_json_schema()}",
+                    level=logging.DEBUG,
+                )
 
         # 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=self._tokenizer
+            dataset_config=self.dataset_config,
+            model_config=self.model_config,
+            tokenizer=self._tokenizer,
+            generative_type=self.generative_type,
         )
 
         # Define the parameters used for vLLM generation
@@ -391,7 +404,10 @@ class VLLMModel(HuggingFaceEncoderModel):
         ) and should_prompts_be_stripped(
             labels_to_be_generated=labels_to_be_generated, tokenizer=self._tokenizer
         ):
-            log_once(message="Stripping prompts.", level=logging.DEBUG)
+            log_once(
+                f"Stripping prompts for model {self.model_config.model_id!r}.",
+                level=logging.DEBUG,
+            )
             prompts = [prompt.strip() for prompt in prompts]
 
         # Generate sequences using vLLM
@@ -411,18 +427,64 @@ class VLLMModel(HuggingFaceEncoderModel):
                     f"Encountered error during vLLM generation: {str(e)}. Retrying..."
                 )
                 sleep(1)
+            except ValueError as e:
+                # Truncate the prompts if they are too long for the model
+                truncate_error_messages = [
+                    r"prompt \(length [0-9]+\) is longer than the maximum model length"
+                ]
+                if any(
+                    re.search(pattern, str(e), flags=re.IGNORECASE) is not None
+                    for pattern in truncate_error_messages
+                ):
+                    logger.info(
+                        "Prompts are too long, so truncating them and trying again..."
+                    )
+                    tokenized_prompts = self._tokenizer(
+                        text=prompts,
+                        truncation=True,
+                        max_length=max(
+                            self._tokenizer.model_max_length - max_tokens, 0
+                        ),
+                    )
+                    prompts = self._tokenizer.batch_decode(
+                        sequences=tokenized_prompts.input_ids, skip_special_tokens=True
+                    )
+                else:
+                    raise InvalidBenchmark(
+                        f"An error occurred during vLLM generation: {str(e)}"
+                    )
         else:
             raise InvalidBenchmark(
                 f"Could not generate sequences after {num_attempts} attempts."
             )
 
+        # When we shorten the prompts then some residual model outputs persist, so we
+        # need to filter these out
+        num_extra_outputs = len(raw_outputs) - len(prompts)
+        if num_extra_outputs > 0:
+            raw_outputs = raw_outputs[num_extra_outputs:]
+            if not all(
+                raw_output.prompt == prompt
+                for raw_output, prompt in zip(raw_outputs, prompts)
+            ):
+                raise InvalidBenchmark(
+                    f"The prompts and the model outputs do not match. There were "
+                    f"{num_extra_outputs!r} extra outputs."
+                )
+            else:
+                logger.debug(
+                    f"Filtered out {num_extra_outputs:,} extra outputs from the model, "
+                    "which occured as we interupted the generation when we truncated "
+                    "the prompts."
+                )
+
         # Parse the raw model outputs
         completion_ids: list[list[int]] = [
             output.outputs[0].token_ids for output in raw_outputs
         ]
         if self.end_of_reasoning_token_id in completion_ids[0]:
             completion_ids = [
-                token_ids[token_ids.index(self.end_of_reasoning_token_id) + 2 :]
+                token_ids[token_ids.index(self.end_of_reasoning_token_id) + 1 :]
                 if self.end_of_reasoning_token_id in token_ids
                 else token_ids
                 for token_ids in completion_ids
@@ -435,6 +497,13 @@ class VLLMModel(HuggingFaceEncoderModel):
         )
         completions = [completion.strip() for completion in completions]
 
+        # Sanity check
+        if len(completions) != len(prompts):
+            breakpoint()
+            raise InvalidBenchmark(
+                f"Expected {len(prompts):,} completions, but got {len(completions):,}."
+            )
+
         # Add logprobs scores to the output
         if self.buffer["first_label_token_mapping"]:
             scores: list[list[list[tuple[str, float]]]] = [
@@ -809,7 +878,8 @@ class VLLMModel(HuggingFaceEncoderModel):
                     if name.lower() in language_codes:
                         chat_template = candidate_template
                         log_once(
-                            f"Using the {name!r} chat template for the tokenizer.",
+                            f"Using the {name!r} chat template for the tokenizer for "
+                            f"model {self.model_config.model_id!r}.",
                             level=logging.DEBUG,
                         )
                         break
@@ -1169,7 +1239,7 @@ def clear_vllm() -> None:
 
 
 def get_end_of_reasoning_token_id(
-    model: "LLM", tokenizer: "PreTrainedTokenizer"
+    model: "LLM", tokenizer: "PreTrainedTokenizer", model_id: str
 ) -> int | None:
     """Get the end of reasoning token ID for a generative model.
 
@@ -1182,6 +1252,8 @@ def get_end_of_reasoning_token_id(
             The vLLM model.
         tokenizer:
             The tokenizer.
+        model_id:
+            The model ID.
 
     Returns:
         The end of reasoning token ID, or None if it could not be found.
@@ -1220,10 +1292,8 @@ def get_end_of_reasoning_token_id(
     completion_match = re.search(pattern=r"<\w+>", string=completion)
     if completion_match is None and prompt_match is None:
         log_once(
-            message=(
-                "Could not find a reasoning token, so assuming the model is not a "
-                "reasoning model."
-            ),
+            f"Could not find a reasoning token for model {model_id!r}, so assuming "
+            "the model is not a reasoning model.",
             level=logging.DEBUG,
         )
         return None
@@ -1249,20 +1319,17 @@ def get_end_of_reasoning_token_id(
         or end_of_reasoning_token not in special_tokens
     ):
         log_once(
-            message=(
-                f"Detected reasoning token {reasoning_token!r} and end-of-reasoning "
-                f"token {end_of_reasoning_token!r}, but one of them is not registered "
-                "as a special token, so assuming it is not a real reasoning token."
-            ),
+            f"Detected reasoning token {reasoning_token!r} and end-of-reasoning "
+            f"token {end_of_reasoning_token!r} for model {model_id!r}, but one of "
+            "them is not registered as a special token, so assuming it is not a "
+            "real reasoning token.",
             level=logging.DEBUG,
         )
         return None
 
     log_once(
-        message=(
-            f"Detected reasoning token {reasoning_token!r} and end-of-reasoning "
-            f"token {end_of_reasoning_token!r}."
-        ),
+        f"Detected reasoning token {reasoning_token!r} and end-of-reasoning "
+        f"token {end_of_reasoning_token!r} for model {model_id!r}.",
         level=logging.DEBUG,
     )
 

euroeval/benchmarker.py

@@ -782,7 +782,11 @@ class Benchmarker:
                     dataset_languages=[
                         language.code for language in dataset_config.languages
                     ],
-                    model=model_config.model_id,
+                    model=(
+                        f"{model_config.model_id}@{model_config.revision}"
+                        if model_config.revision and model_config.revision != "main"
+                        else model_config.model_id
+                    ),
                     results=results,
                     num_model_parameters=model.num_params,
                     max_sequence_length=model.model_max_length,
@@ -1076,6 +1080,10 @@ def initial_logging(
         benchmark_config:
             The general benchmark configuration.
     """
+    model_id = model_config.model_id
+    if model_config.revision and model_config.revision != "main":
+        model_id += f"@{model_config.revision}"
+
     split_type = "validation" if not benchmark_config.evaluate_test_split else "test"
     if model_config.task in GENERATIVE_PIPELINE_TAGS:
         if benchmark_config.few_shot:
@@ -1084,8 +1092,9 @@ def initial_logging(
             eval_type = "Zero-shot benchmarking"
     else:
         eval_type = "Benchmarking"
+
     logger.info(
-        f"{eval_type} {model_config.model_id} on the {split_type} split of "
+        f"{eval_type} {model_id} on the {split_type} split of "
         f"{dataset_config.pretty_name}"
     )
 
@@ -1095,6 +1104,7 @@ def initial_logging(
             "meaning that the resulting evaluation will not be included in the "
             "official leaderboard."
         )
+
     if benchmark_config.debug:
         logger.info(
             "Running in debug mode. This will output additional information, as "

euroeval/constants.py

@@ -16,7 +16,7 @@ MAX_CONTEXT_LENGTH = 5_000
 
 # We need to raise the amount of tokens generated for reasoning models, to give them
 # time to think
-REASONING_MAX_TOKENS = 8_192
+REASONING_MAX_TOKENS = 32_768
 
 
 # The Hugging Face Hub pipeline tags used to classify models as generative

euroeval/data_loading.py

@@ -39,32 +39,9 @@ def load_data(
         HuggingFaceHubDown:
             If the Hugging Face Hub is down.
     """
-    num_attempts = 5
-    for _ in range(num_attempts):
-        try:
-            dataset = load_dataset(
-                path=dataset_config.huggingface_id,
-                cache_dir=benchmark_config.cache_dir,
-                token=unscramble("HjccJFhIozVymqXDVqTUTXKvYhZMTbfIjMxG_"),
-            )
-            break
-        except (FileNotFoundError, DatasetsError, ConnectionError, ReadTimeout):
-            logger.warning(
-                f"Failed to load dataset {dataset_config.huggingface_id!r}. Retrying..."
-            )
-            time.sleep(1)
-            continue
-        except HfHubHTTPError:
-            raise HuggingFaceHubDown()
-    else:
-        raise InvalidBenchmark(
-            f"Failed to load dataset {dataset_config.huggingface_id!r} after "
-            f"{num_attempts} attempts."
-        )
-
-    assert isinstance(dataset, DatasetDict)  # type: ignore[used-before-def]
-
-    dataset = DatasetDict({key: dataset[key] for key in ["train", "val", "test"]})
+    dataset = load_raw_data(
+        dataset_config=dataset_config, cache_dir=benchmark_config.cache_dir
+    )
 
     if not benchmark_config.evaluate_test_split:
         dataset["test"] = dataset["val"]
@@ -101,3 +78,48 @@ def load_data(
         for idx in range(benchmark_config.num_iterations)
     ]
     return datasets
+
+
+def load_raw_data(dataset_config: "DatasetConfig", cache_dir: str) -> DatasetDict:
+    """Load the raw dataset.
+
+    Args:
+        dataset_config:
+            The configuration for the dataset.
+        cache_dir:
+            The directory to cache the dataset.
+
+    Returns:
+        The dataset.
+    """
+    num_attempts = 5
+    for _ in range(num_attempts):
+        try:
+            dataset = load_dataset(
+                path=dataset_config.huggingface_id,
+                cache_dir=cache_dir,
+                token=unscramble("HjccJFhIozVymqXDVqTUTXKvYhZMTbfIjMxG_"),
+            )
+            break
+        except (FileNotFoundError, DatasetsError, ConnectionError, ReadTimeout):
+            logger.warning(
+                f"Failed to load dataset {dataset_config.huggingface_id!r}. Retrying..."
+            )
+            time.sleep(1)
+            continue
+        except HfHubHTTPError:
+            raise HuggingFaceHubDown()
+    else:
+        raise InvalidBenchmark(
+            f"Failed to load dataset {dataset_config.huggingface_id!r} after "
+            f"{num_attempts} attempts."
+        )
+    assert isinstance(dataset, DatasetDict)  # type: ignore[used-before-def]
+    required_keys = ["train", "val", "test"]
+    missing_keys = [key for key in required_keys if key not in dataset]
+    if missing_keys:
+        raise InvalidBenchmark(
+            "The dataset is missing the following required splits: "
+            f"{', '.join(missing_keys)}"
+        )
+    return DatasetDict({key: dataset[key] for key in required_keys})