Repository: ScandEval/ScandEval (language: python)

Issue:
[FEATURE REQUEST] Allow benchmarking Anthropic models
### 🚀 The feature, motivation and pitch

We should be able to benchmark Anthropic models just like OpenAI models. A lot of the OpenAI code can probably be reused.

### Alternatives

_No response_

### Additional context

_No response_
[FEATURE REQUEST] Support llama.cpp backend
### 🚀 The feature, motivation and pitch

Using the llama.cpp backend (via the Python bindings) will enable us to benchmark gguf models, and also make it easier to benchmark larger models with less VRAM available. 

### Alternatives

_No response_

### Additional context

Since the gguf models are saved in separate model repos, there’s full transparency in case the fp16 model is significantly better, say. 


Discussion / hints:

Just adding my two cents: 

It seems it varies a little how people publish the gguf's. NorwAI put them in the same repo, and one then has to download the file manually, or use hf-hub and exclude everything else. Others have a branch for them, which makes it easy when loading from transformers (same repo, but with revision/branch x in the destination).

Is there an easy way to cover both scenarios? 😄 
> It seems it varies a little how people publish the gguf's. NorwAI put them in the same repo, and one then has to download the file manually, or use hf-hub and exclude everything else. Others have a branch for them, which makes it easy when loading from transformers (same repo, but with revision/branch x in the destination).
> Is there an easy way to cover both scenarios? 😄

ScandEval supports loading revisions of models, so in case the ggufs are in a separate branch then one could simply evaluate <model-id>@<branch-name>. If they're in the same repo then I assume that llama.cpp looks for gguf files in the repo, but I'm not sure about this. In any case, it sounds doable!

Fix patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 00e7876fb..bebd48478 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -25,7 +25,7 @@ jobs:
     if: github.event.pull_request.draft == false
     strategy:
         matrix:
-            python-version: ["3.10", "3.11"]
+            python-version: ["3.10", "3.11", "3.12"]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -44,7 +44,9 @@ jobs:
       - name: Test with pytest
         run: uv run pytest
         env:
-          HUGGINGFACE_HUB_TOKEN: ${{ secrets.HUGGINGFACE_HUB_TOKEN }}
+          HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
 
       - name: Delete ScandEval cache
         run: rm -rf .scandeval_cache
@@ -69,7 +71,9 @@ jobs:
       - name: Test with pytest
         run: uv run pytest
         env:
-          HUGGINGFACE_HUB_TOKEN: ${{ secrets.HUGGINGFACE_HUB_TOKEN }}
+          HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
 
       - name: Delete ScandEval cache
         run: rm -rf .scandeval_cache
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0295659c0..d0643398c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -8,7 +8,7 @@ repos:
     hooks:
       - id: end-of-file-fixer
       - id: trailing-whitespace
-      - id: debug-statements
+      #- id: debug-statements
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.7.1
     hooks:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e60265410..5d6b7807c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,30 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 
 ## [Unreleased]
+### Added
+- Added support for [LiteLLM](https://docs.litellm.ai/), meaning that all LLMs on 100+
+  APIs can now be benchmarked! This includes OpenAI, Anthropic, Google, Mistral AI,
+  Cohere, Ollama, LM Studio, vLLM servers, and Hugging Face inference endpoints. Check
+  out the full list of LiteLLM providers [here](https://docs.litellm.ai/docs/providers).
+
+### Changed
+- No more tokenisation for generation tasks, resulting in faster preprocessing times.
+- Now evaluates models on the validation split by default, to avoid overfitting to the
+  test set. The test set can be evaluated on using the `--evaluate-test-split` flag.
 
+### Removed
+- Removed the option to evaluate on the training split, as this is not a common use
+  case and simplified the codebase. If you find that this should be re-added, please
+  open an issue in the GitHub repository.
+- All generative on-premises models are now evaluated with vLLM and thus does not use
+  the `transformers` backend as a backup, as this was not used in practice, and
+  simplified the codebase. If you find that this should be re-added, please open an
+  issue in the GitHub repository.
+- Removed the `--only-validation-split` flag, as this is now the default behaviour. If
+  you find that this should be re-added, please open an issue in the GitHub repository.
+- Removed the option to benchmark local models, as this was not used in practice, and
+  simplified the codebase. If you find that this should be re-added, please open an
+  issue in the GitHub repository.
 
 
 ## [v13.3.0] - 2024-11-29
diff --git a/makefile b/makefile
index 0e5f69ac8..d78448aa5 100644
--- a/makefile
+++ b/makefile
@@ -74,50 +74,7 @@ publish-docs:  ## Publish documentation to GitHub Pages
 	@echo "Updated documentation website: https://scandeval.com/ScandEval/"
 
 test:  ## Run tests
-	@if [ "$(shell which nvidia-smi)" != "" ]; then \
-		$(MAKE) --quiet test-cuda-vllm; \
-	else \
-		$(MAKE) --quiet test-cpu; \
-	fi
-	@$(MAKE) --quiet update-coverage-badge
-	@date "+%H:%M:%S ⋅ All done!"
-
-test-cuda-vllm:
-	@rm tests_with_cuda_and_vllm.log; \
-		date "+%H:%M:%S ⋅ Running tests with CUDA and vLLM..." \
-		&& USE_CUDA=1 USE_VLLM=1 \
-			uv run pytest | tee tests_with_cuda_and_vllm.log \
-		&& date "+%H:%M:%S ⋅ Finished testing with CUDA and vLLM!"
-
-test-cuda-no-vllm:
-	@rm tests_with_cuda_and_no_vllm.log; \
-		date "+%H:%M:%S ⋅ Running tests with CUDA and no vLLM..." \
-		&& USE_CUDA=1 USE_VLLM=0 \
-			uv run pytest | tee tests_with_cuda_and_no_vllm.log \
-		&& date "+%H:%M:%S ⋅ Finished testing with CUDA and no vLLM!"
-
-test-cpu:
-	@rm tests_with_cpu.log; \
-		date "+%H:%M:%S ⋅ Running tests with CPU..." \
-		&& USE_CUDA=0 uv run pytest | tee tests_with_cpu.log \
-		&& date "+%H:%M:%S ⋅ Finished testing with CPU!"
-
-test-fast:  # Run CPU tests without evaluations
-	@rm tests_with_cpu_fast.log; \
-		date "+%H:%M:%S ⋅ Running fast tests with CPU..." \
-		&& USE_CUDA=0 TEST_EVALUATIONS=0 \
-			uv run pytest | tee tests_with_cpu_fast.log \
-		&& date "+%H:%M:%S ⋅ Finished fast testing with CPU!"
-
-test-slow:  # Run all tests
-	@if [ "$(shell which nvidia-smi)" != "" ]; then \
-		TEST_ALL_DATASETS=1 $(MAKE) --quiet test-cuda-vllm; \
-		TEST_ALL_DATASETS=1 $(MAKE) --quiet test-cuda-no-vllm; \
-	else \
-		TEST_ALL_DATASETS=1 $(MAKE) --quiet test-cpu; \
-	fi
-	@$(MAKE) --quiet update-coverage-badge
-	@date "+%H:%M:%S ⋅ All done!"
+	@uv run pytest && uv run readme-c
 
 tree:  ## Print directory tree
 	@tree -a --gitignore -I .git .
diff --git a/pyproject.toml b/pyproject.toml
index e91787fab..b84e5a78a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,6 +29,12 @@ dependencies = [
     "sentencepiece>=0.1.96",
     "protobuf~=3.20.0",
     "sacremoses>=0.1.1",
+    "more-itertools>=10.5.0",
+    "tenacity>=9.0.0",
+    "litellm>=1.53.1",
+    "rouge-score>=0.1.2",
+    "bert-score>=0.3.13",
+    "levenshtein>=0.24.0",
 ]
 
 [project.optional-dependencies]
@@ -42,46 +48,26 @@ generative = [
     "outlines>=0.0.44,<0.1",
     "bitsandbytes>=0.43.1; platform_system == 'Linux'",
     "vllm>=0.6.3; platform_system == 'Linux'",
-    "rouge-score>=0.1.2",
-    "bert-score>=0.3.13",
-]
-openai = [
-    "rouge-score>=0.1.2",
-    "bert-score>=0.3.13",
-    "openai>=1.51.2",
-    "tiktoken>=0.7.0",
-    "levenshtein>=0.24.0",
 ]
 human_evaluation = [
     "gradio>=4.26.0",
-    "levenshtein>=0.24.0",
 ]
 all = [
     "jax>=0.4.24",
     "jaxlib>=0.4.24",
     "flax>=0.8.1",
-    "rouge-score>=0.1.2",
-    "bert-score>=0.3.13",
     "demjson3>=3.0.6",
     "outlines>=0.0.44,<0.1",
     "bitsandbytes>=0.43.1; platform_system == 'Linux'",
     "vllm>=0.6.3; platform_system == 'Linux'",
-    "openai>=1.51.2",
-    "tiktoken>=0.7.0",
-    "levenshtein>=0.24.0",
     "gradio>=4.26.0",
 ]
 cpu_all = [
     "jax>=0.4.24",
     "jaxlib>=0.4.24",
     "flax>=0.8.1",
-    "rouge-score>=0.1.2",
-    "bert-score>=0.3.13",
     "demjson3>=3.0.6",
     "outlines>=0.0.44,<0.1",
-    "openai>=1.51.2",
-    "tiktoken>=0.7.0",
-    "levenshtein>=0.24.0",
     "gradio>=4.26.0",
 ]
 
@@ -161,10 +147,8 @@ convention = "google"
 [tool.pytest.ini_options]
 minversion = "7.0"
 addopts = [
-    '--durations=10',
-    '--color=yes',
-    '--doctest-modules',
-    '--cov=src/scandeval',
+    "--cov=src/scandeval",
+    "--color=yes",
 ]
 xfail_strict = true
 filterwarnings = [
diff --git a/src/scandeval/__init__.py b/src/scandeval/__init__.py
index de3639e03..3b863153d 100644
--- a/src/scandeval/__init__.py
+++ b/src/scandeval/__init__.py
@@ -12,7 +12,7 @@
 from .utils import block_terminal_output
 
 # Fetches the version of the package as defined in pyproject.toml
-__version__ = importlib.metadata.version(__package__)
+__version__ = importlib.metadata.version("scandeval")
 
 
 # Block unwanted terminal outputs
diff --git a/src/scandeval/benchmark_config_factory.py b/src/scandeval/benchmark_config_factory.py
index 9172d619a..004d00b64 100644
--- a/src/scandeval/benchmark_config_factory.py
+++ b/src/scandeval/benchmark_config_factory.py
@@ -2,24 +2,23 @@
 
 import importlib.util
 import logging
-import os
 import sys
-from typing import TYPE_CHECKING
+import typing as t
 
 import torch
 
-from .config import BenchmarkConfig
+from .data_models import BenchmarkConfig
 from .dataset_configs import get_all_dataset_configs
 from .enums import Device, Framework
 from .exceptions import InvalidBenchmark
 from .languages import get_all_languages
 from .tasks import get_all_tasks
 
-if TYPE_CHECKING:
-    from .config import Language, Task
+if t.TYPE_CHECKING:
+    from .data_models import Language, Task
 
 
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 def build_benchmark_config(
@@ -33,22 +32,16 @@ def build_benchmark_config(
     framework: Framework | str | None,
     device: Device | None,
     batch_size: int,
-    evaluate_train: bool,
     raise_errors: bool,
     cache_dir: str,
-    token: bool | str | None,
-    openai_api_key: str | None,
-    prefer_azure: bool,
-    azure_openai_api_key: str | None,
-    azure_openai_endpoint: str | None,
-    azure_openai_api_version: str | None,
+    api_key: str | None,
     force: bool,
     verbose: bool,
     trust_remote_code: bool,
     load_in_4bit: bool | None,
     use_flash_attention: bool | None,
     clear_model_cache: bool,
-    only_validation_split: bool,
+    evaluate_test_split: bool,
     few_shot: bool,
     num_iterations: int,
     debug: bool,
@@ -86,30 +79,18 @@ def build_benchmark_config(
             set automatically.
         batch_size:
             The batch size to use for running the models.
-        evaluate_train:
-            Whether to evaluate the models on the training set.
         raise_errors:
             Whether to raise errors when running the benchmark.
         cache_dir:
             The directory to use for caching the models.
-        token:
-            The token to use for running the models.
-        openai_api_key:
-            The OpenAI API key to use for running the models.
-        prefer_azure:
-            Whether to prefer the Azure OpenAI API for running the models, over the
-            OpenAI API.
-        azure_openai_api_key:
-            The Azure OpenAI API key to use for running the models.
-        azure_openai_endpoint:
-            The Azure OpenAI endpoint to use for running the models.
-        azure_openai_api_version:
-            The Azure OpenAI api version to use for running the models.
+        api_key:
+            The API key to use for a given inference server.
         force:
             Whether to force the benchmark to run even if the results are already
             cached.
         verbose:
-            Whether to print verbose output when running the benchmark.
+            Whether to print verbose output when running the benchmark. This is
+            automatically set if `debug` is True.
         trust_remote_code:
             Whether to trust remote code when running the benchmark.
         load_in_4bit:
@@ -119,8 +100,8 @@ def build_benchmark_config(
             if it is available.
         clear_model_cache:
             Whether to clear the model cache before running the benchmark.
-        only_validation_split:
-            Whether to only use the validation split for the datasets.
+        evaluate_test_split:
+            Whether to use the test split for the datasets.
         few_shot:
             Whether to use few-shot learning for the models.
         num_iterations:
@@ -150,52 +131,8 @@ def build_benchmark_config(
 
     torch_device = prepare_device(device=device)
 
-    if openai_api_key is None:
-        openai_api_key = os.getenv("OPENAI_API_KEY")
-    if azure_openai_api_key is None:
-        azure_openai_api_key = os.getenv("AZURE_OPENAI_API_KEY")
-    if azure_openai_endpoint is None:
-        azure_openai_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
-    if azure_openai_api_version is None:
-        azure_openai_api_version = os.getenv("AZURE_OPENAI_API_VERSION")
-
-    # Ensure that we are not using both OpenAI and Azure OpenAI API keys
-    if all(
-        value is not None
-        for value in (
-            openai_api_key,
-            azure_openai_api_key,
-            azure_openai_endpoint,
-            azure_openai_api_version,
-        )
-    ):
-        if prefer_azure:
-            logger.info(
-                "Both OpenAI and Azure OpenAI API keys are set. Using Azure OpenAI."
-            )
-            openai_api_key = None
-        else:
-            if run_with_cli:
-                logger.info(
-                    "Both OpenAI and Azure OpenAI API keys are set. Using OpenAI since "
-                    "the `--prefer-azure` flag is not set."
-                )
-            else:
-                logger.info(
-                    "Both OpenAI and Azure OpenAI API keys are set. Using OpenAI since "
-                    "the `prefer_azure` argument is not set to True."
-                )
-            azure_openai_api_key = None
-            azure_openai_endpoint = None
-
-    # Sanity check
-    assert not (openai_api_key is not None and azure_openai_api_key is not None)
-
     framework_obj = Framework(framework) if framework is not None else None
 
-    if token is True:
-        token = None
-
     if use_flash_attention is None:
         if torch_device.type != "cuda":
             use_flash_attention = False
@@ -232,23 +169,18 @@ def build_benchmark_config(
         batch_size=batch_size,
         raise_errors=raise_errors,
         cache_dir=cache_dir,
-        evaluate_train=evaluate_train,
-        token=token,
-        openai_api_key=openai_api_key,
-        azure_openai_api_key=azure_openai_api_key,
-        azure_openai_endpoint=azure_openai_endpoint,
-        azure_openai_api_version=azure_openai_api_version,
+        api_key=api_key,
         force=force,
         progress_bar=progress_bar,
         save_results=save_results,
-        verbose=verbose,
+        verbose=verbose or debug,
         framework=framework_obj,
         device=torch_device,
         trust_remote_code=trust_remote_code,
         load_in_4bit=load_in_4bit,
         use_flash_attention=use_flash_attention,
         clear_model_cache=clear_model_cache,
-        only_validation_split=only_validation_split,
+        evaluate_test_split=evaluate_test_split,
         few_shot=few_shot,
         num_iterations=num_iterations,
         debug=debug,
diff --git a/src/scandeval/benchmark_dataset.py b/src/scandeval/benchmark_dataset.py
deleted file mode 100644
index 4a904f751..000000000
--- a/src/scandeval/benchmark_dataset.py
+++ /dev/null
@@ -1,753 +0,0 @@
-"""Abstract benchmarking dataset class."""
-
-import logging
-import sys
-from abc import ABC, abstractmethod
-from functools import partial
-from typing import TYPE_CHECKING, Any, Type
-
-import evaluate
-import torch
-from datasets.dataset_dict import DatasetDict
-from datasets.load import load_dataset
-from huggingface_hub import HfApi
-from huggingface_hub.hf_api import ModelInfo
-from huggingface_hub.utils import HfHubHTTPError, HFValidationError
-from requests import RequestException
-from tqdm.auto import tqdm
-from transformers import Trainer
-from transformers.modeling_utils import PreTrainedModel
-
-from .exceptions import InvalidBenchmark
-from .finetuning import finetune
-from .generation import generate
-from .model_config import get_model_config
-from .model_loading import load_model
-from .scores import log_scores
-from .speed_benchmark import benchmark_speed
-from .tasks import SPEED
-from .utils import (
-    GENERATIVE_MODEL_TASKS,
-    enforce_reproducibility,
-    get_model_max_length,
-    model_is_generative,
-    should_prompts_be_stripped,
-    unscramble,
-)
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-    from numpy.random import Generator
-    from transformers import PretrainedConfig
-    from transformers.modeling_utils import ModelOutput
-
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Labels, Predictions, ScoreDict
-
-
-logger = logging.getLogger(__package__)
-
-
-class BenchmarkDataset(ABC):
-    """Abstract benchmarking dataset class.
-
-    Args:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-
-    Attributes:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-    """
-
-    def __init__(
-        self, dataset_config: "DatasetConfig", benchmark_config: "BenchmarkConfig"
-    ) -> None:
-        """Initialise the dataset.
-
-        Args:
-            dataset_config:
-                The configuration for the dataset.
-            benchmark_config:
-                The configuration for the benchmark.
-        """
-        self.dataset_config = dataset_config
-        self.benchmark_config = benchmark_config
-        self._metrics = {
-            metric_cfg.name: (
-                evaluate.load(
-                    path=metric_cfg.huggingface_id,
-                    cache_dir=self.benchmark_config.cache_dir,
-                )
-                if metric_cfg.huggingface_id != ""
-                else None
-            )
-            for metric_cfg in dataset_config.task.metrics
-        }
-
-        # Set logging level based on verbosity
-        if hasattr(sys, "_called_from_test"):
-            logging_level = logging.CRITICAL
-        elif self.benchmark_config.verbose:
-            logging_level = logging.DEBUG
-        else:
-            logging_level = logging.INFO
-        logger.setLevel(logging_level)
-
-    def benchmark(
-        self,
-        model_id: str,
-        model: "GenerativeModel | None" = None,
-        tokenizer: "Tokenizer | None" = None,
-    ) -> "tuple[ScoreDict, dict[str, bool | int], GenerativeModel | None, Tokenizer | None]":
-        """Benchmark a model.
-
-        Args:
-            model_id:
-                The full Hugging Face Hub path to the pretrained transformer model. The
-                specific model version to use can be added after the suffix '@':
-                "model_id@v1.0.0". It can be a branch name, a tag name, or a commit id,
-                and defaults to the latest version if not specified.
-            model:
-                The model to benchmark. If not provided, the model will be loaded.
-            tokenizer:
-                The tokenizer to use with the model. If not provided, the tokenizer will
-                be loaded.
-
-        Returns:
-            A pair (scores, metadata_dict, model, tokenizer), with `scores` being a
-            dictionary containing the scores, and `metadata_dict` being a dictionary
-            containing various model metadata, such as the number of model parameters,
-            the model's maximum sequence length and the size of the model's vocabulary.
-            The keys in `score_dict` are 'raw' and 'total', with all the raw scores in
-            the first dictionary and the aggregated scores in the second. The tokenizer
-            and model are only not `None` if the model is generative.
-
-        Raises:
-            RuntimeError:
-                If the extracted framework is not recognized.
-        """
-        model_config = get_model_config(
-            model_id=model_id, benchmark_config=self.benchmark_config
-        )
-
-        # Set random seeds to enforce reproducibility of the randomly initialised
-        # weights
-        rng = enforce_reproducibility(framework=model_config.framework)
-
-        if model is None or tokenizer is None:
-            logger.info("Loading model and tokenizer...")
-            model, tokenizer = load_model(
-                model_config=model_config,
-                dataset_config=self.dataset_config,
-                benchmark_config=self.benchmark_config,
-            )
-
-        benchmarking_generative_model = model_is_generative(model=model)
-
-        # This happens when a local model is used, as we cannot fetch the model
-        # metadata. Note that this is only the case if the model type is not any of the
-        # ones hardcoded in `local.py`
-        if model_config.task == "unknown":
-            if benchmarking_generative_model:
-                model_config.task = GENERATIVE_MODEL_TASKS[0]
-            else:
-                model_config.task = "fill-mask"
-
-        metadata_dict = self._get_metadata(
-            model_id=model_id,
-            model=model,
-            tokenizer=tokenizer,
-            benchmarking_generative_model=benchmarking_generative_model,
-        )
-
-        if self.dataset_config.task != SPEED:
-            train, val, tests = self._load_data(rng=rng)
-            prepared_train, prepared_val, prepared_tests = self._load_prepared_data(
-                train=train,
-                val=val,
-                tests=tests,
-                model_config=model_config,
-                hf_model_config=model.config,
-                tokenizer=tokenizer,
-                benchmarking_generative_model=benchmarking_generative_model,
-            )
-
-        # Set up progress bar
-        itr = tqdm(
-            iterable=range(self.benchmark_config.num_iterations),
-            desc="Benchmarking",
-            disable=not self.benchmark_config.progress_bar,
-        )
-
-        data_collator = self._load_data_collator(tokenizer=tokenizer, model=model)
-
-        if self.dataset_config.task == SPEED:
-            scores = benchmark_speed(
-                itr=itr,
-                tokenizer=tokenizer,
-                model=model,
-                model_config=model_config,
-                dataset_config=self.dataset_config,
-                benchmark_config=self.benchmark_config,
-            )
-        elif benchmarking_generative_model:
-            scores = generate(
-                itr=itr,
-                prepared_train=prepared_train,
-                prepared_tests=prepared_tests,
-                model=model,
-                model_config=model_config,
-                tokenizer=tokenizer,
-                data_collator=data_collator,
-                compute_metrics=self._compute_metrics,
-                extract_labels_fn=self._extract_labels_from_generation,
-                benchmark_config=self.benchmark_config,
-                dataset_config=self.dataset_config,
-            )
-        else:
-            scores = finetune(
-                itr=itr,
-                train=train,
-                val=val,
-                tests=tests,
-                prepared_train=prepared_train,
-                prepared_val=prepared_val,
-                prepared_tests=prepared_tests,
-                model=model,
-                tokenizer=tokenizer,
-                model_config=model_config,
-                dataset_config=self.dataset_config,
-                benchmark_config=self.benchmark_config,
-                compute_metrics=self._compute_metrics,
-                data_collator=data_collator,
-                trainer_class=self._get_trainer_class(),
-                evaluate_inputs_fn=self._get_evaluate_inputs,
-                preprocess_logits_for_metrics=self._preprocess_logits_for_metrics,
-            )
-
-        all_scores = log_scores(
-            dataset_name=self.dataset_config.pretty_name,
-            metric_configs=self.dataset_config.task.metrics,
-            scores=scores,
-            model_id=model_config.model_id,
-        )
-
-        if benchmarking_generative_model:
-            return all_scores, metadata_dict, model, tokenizer
-        else:
-            return all_scores, metadata_dict, None, None
-
-    def _get_metadata(
-        self,
-        model_id: str,
-        model: "PreTrainedModel | GenerativeModel",
-        tokenizer: "Tokenizer",
-        benchmarking_generative_model: bool,
-    ) -> dict[str, int]:
-        """Get metadata about the model.
-
-        Args:
-            model_id:
-                The Hugging Face model ID.
-            model:
-                The model to get metadata about.
-            tokenizer:
-                The tokenizer to get metadata about.
-            benchmarking_generative_model:
-                Whether the model is a generative model.
-
-        Returns:
-            A dictionary containing metadata about the model, with the keys being the
-            metadata names and the values being the metadata values.
-        """
-        api = HfApi()
-        try:
-            repo_info = api.repo_info(repo_id=model_id, repo_type="model")
-            assert isinstance(repo_info, ModelInfo)
-        except (RequestException, HFValidationError):
-            repo_info = None
-
-        if (
-            repo_info is not None
-            and hasattr(repo_info, "safetensors")
-            and repo_info.safetensors is not None
-            and "total" in repo_info.safetensors
-        ):
-            num_params = repo_info.safetensors["total"]
-        elif (
-            repo_info is not None
-            and hasattr(repo_info, "card_data")
-            and repo_info.card_data is not None
-            and "base_model" in repo_info.card_data
-            and repo_info.card_data["base_model"] is not None
-            and len(repo_info.card_data["base_model"]) > 0
-            and (
-                base_repo_info := api.repo_info(
-                    repo_info.card_data["base_model"][-1], repo_type="model"
-                )
-            )
-            is not None
-            and hasattr(base_repo_info, "safetensors")
-            and base_repo_info.safetensors is not None
-            and "total" in base_repo_info.safetensors
-        ):
-            num_params = base_repo_info.safetensors["total"]
-        elif (
-            hasattr(model.config, "num_params") and model.config.num_params is not None
-        ):
-            num_params = model.config.num_params
-        elif isinstance(model, PreTrainedModel):
-            num_params = sum(p.numel() for p in model.parameters())
-        else:
-            num_params = -1
-
-        max_seq_length = get_model_max_length(model=model)
-
-        if hasattr(model.config, "vocab_size") and model.config.vocab_size is not None:
-            vocab_size = model.config.vocab_size
-        elif hasattr(tokenizer, "vocab_size") and tokenizer.vocab_size is not None:
-            vocab_size = tokenizer.vocab_size
-        else:
-            vocab_size = -1
-
-        # Store the metadata in a dictionary
-        metadata_dict = dict(
-            num_model_parameters=num_params,
-            max_sequence_length=max_seq_length,
-            vocabulary_size=vocab_size,
-            generative=benchmarking_generative_model,
-            few_shot=self.benchmark_config.few_shot,
-            validation_split=self.benchmark_config.only_validation_split,
-        )
-
-        # Log the metadata
-        logging_msg: str = ""
-        if num_params < 0:
-            logging_msg += "The model has an unknown number of parameters, "
-        else:
-            logging_msg += f"The model has {num_params:,} parameters, "
-        if vocab_size < 0:
-            logging_msg += "an unknown vocabulary size, "
-        else:
-            logging_msg += f"a vocabulary size of {vocab_size:,}, "
-        if max_seq_length < 0:
-            logging_msg += "and an unknown maximum sequence length."
-        else:
-            logging_msg += f"and a maximum sequence length of {max_seq_length:,}."
-        logger.info(logging_msg)
-
-        return metadata_dict
-
-    def _load_data(
-        self, rng: "Generator"
-    ) -> tuple["Dataset", "Dataset", list["Dataset"]]:
-        """Load the raw bootstrapped datasets.
-
-        Args:
-            rng:
-                The random number generator to use.
-
-        Returns:
-            A tuple containing the training, validation and test datasets.
-        """
-        # Download dataset from the HF Hub
-        try:
-            dataset_dict = load_dataset(
-                path=self.dataset_config.huggingface_id,
-                cache_dir=self.benchmark_config.cache_dir,
-                token=unscramble("HjccJFhIozVymqXDVqTUTXKvYhZMTbfIjMxG_"),
-            )
-        except HfHubHTTPError:
-            raise InvalidBenchmark("The Hugging Face Hub seems to be down.")
-
-        # If the dataset turns out not to be a DatasetDict, then we raise an error
-        if not isinstance(dataset_dict, DatasetDict):
-            raise InvalidBenchmark(
-                f"Expected `dataset_dict` to be a `DatasetDict`, but got "
-                f"{type(dataset_dict)}."
-            )
-
-        # Remove all other keys than 'train', 'val' and 'test'
-        dataset_dict = DatasetDict(
-            {key: dataset_dict[key] for key in ["train", "val", "test"]}
-        )
-
-        # Process the datasets
-        dataset_dict = self._process_data(dataset_dict)
-
-        # Extract the dataset splits
-        train = dataset_dict["train"]
-        val = dataset_dict["val"]
-        test = dataset_dict["test"]
-
-        if self.benchmark_config.only_validation_split:
-            test = val
-
-        # Remove empty examples from the datasets
-        for text_feature in ["tokens", "text"]:
-            if text_feature in train.features:
-                train = train.filter(lambda x: len(x[text_feature]) > 0)
-                val = val.filter(lambda x: len(x[text_feature]) > 0)
-                test = test.filter(lambda x: len(x[text_feature]) > 0)
-
-        # If we are testing then truncate the test set
-        if hasattr(sys, "_called_from_test"):
-            test = test.select(range(1))
-
-        # Bootstrap the test set
-        test_bidxs = rng.integers(
-            0, len(test), size=(self.benchmark_config.num_iterations, len(test))
-        )
-        tests = [test.select(test_bidxs[idx]) for idx in range(test_bidxs.shape[0])]
-
-        return train, val, tests
-
-    def _load_prepared_data(
-        self,
-        train: "Dataset",
-        val: "Dataset",
-        tests: list["Dataset"],
-        model_config: "ModelConfig",
-        hf_model_config: "PretrainedConfig",
-        tokenizer: "Tokenizer",
-        benchmarking_generative_model: bool,
-    ) -> tuple["Dataset", "Dataset", list["Dataset"]]:
-        """Load the data and prepare it for training.
-
-        Args:
-            train:
-                The raw training dataset.
-            val:
-                The raw validation dataset.
-            tests:
-                The raw bootstrapped test datasets.
-            model_config:
-                The model configuration.
-            hf_model_config:
-                The Hugging Face model configuration.
-            tokenizer:
-                The tokenizer.
-            benchmarking_generative_model:
-                Whether the model is a generative model.
-
-        Returns:
-            A tuple containing the prepared training, validation and test datasets.
-        """
-        # Set up the preprocessing parameters
-        preprocess_params: dict[str, Any] = dict(
-            hf_model_config=hf_model_config,
-            model_config=model_config,
-            tokenizer=tokenizer,
-            generative_model=benchmarking_generative_model,
-        )
-
-        # Prepare the train and validation datasets
-        with tqdm(
-            total=2 + self.benchmark_config.num_iterations,
-            desc="Preprocessing data splits",
-            disable=hasattr(sys, "_called_from_test"),
-        ) as pbar:
-            # When evaluating generative models we only need the test split, so
-            # there's no need to prepare the train split
-            try:
-                prepared_train = train
-                if not benchmarking_generative_model:
-                    prepared_train = self._preprocess_data(
-                        train, split="train", **preprocess_params
-                    )
-                pbar.update(1)
-            except ValueError:
-                raise InvalidBenchmark(
-                    "Preprocessing of the training dataset could not be done."
-                )
-
-            # When evaluating generative models we only need the test split, so
-            # there's no need to prepare the validation split
-            try:
-                prepared_val = val
-                if not benchmarking_generative_model:
-                    prepared_val = self._preprocess_data(
-                        val, split="val", **preprocess_params
-                    )
-                pbar.update(1)
-            except ValueError:
-                raise InvalidBenchmark(
-                    "Preprocessing of the validation dataset could not be done."
-                )
-
-            try:
-                prepared_tests: list["Dataset"] = list()
-                for itr_idx, test in enumerate(tests):
-                    if benchmarking_generative_model:
-                        if self.benchmark_config.few_shot:
-                            itr_seed = 4242 + itr_idx
-                            few_shot_examples = self._extract_few_shot_examples(
-                                train_dataset=train, random_seed=itr_seed
-                            )
-                            few_shot_fn = partial(
-                                self._apply_few_shot_prompt,
-                                few_shot_examples=few_shot_examples,
-                                tokenizer=tokenizer,
-                            )
-                            test = test.map(
-                                few_shot_fn,
-                                batched=True,
-                                load_from_cache_file=False,
-                                keep_in_memory=True,
-                            )
-                        else:
-                            instruction_fn = partial(
-                                self._apply_instruction_prompt, tokenizer=tokenizer
-                            )
-                            test = test.map(
-                                function=instruction_fn,
-                                batched=True,
-                                load_from_cache_file=False,
-                                keep_in_memory=True,
-                            )
-
-                        # Determine if we should strip the prompts. This is the case if
-                        # the tokenizer needs to include the space as part of the label
-                        # token
-                        labels_to_be_generated = list(
-                            self.dataset_config.prompt_label_mapping.values()
-                        )
-                        strip_prompts = should_prompts_be_stripped(
-                            labels_to_be_generated=labels_to_be_generated,
-                            tokenizer=tokenizer,
-                        )
-                        if strip_prompts:
-                            test = test.map(
-                                lambda x: dict(text=x["text"].strip()),
-                                load_from_cache_file=False,
-                                keep_in_memory=True,
-                            )
-
-                    prepared_test = self._preprocess_data(
-                        test, split="test", **preprocess_params
-                    )
-
-                    if benchmarking_generative_model:
-                        prepared_test = prepared_test.map(
-                            lambda examples: dict(
-                                text=tokenizer.batch_decode(
-                                    sequences=examples["input_ids"],
-                                    skip_special_tokens=True,
-                                )
-                            ),
-                            batched=True,
-                            load_from_cache_file=False,
-                            keep_in_memory=True,
-                        )
-
-                    prepared_tests.append(prepared_test)
-                    pbar.update(1)
-            except ValueError:
-                raise InvalidBenchmark(
-                    "Preprocessing of the test dataset could not be done."
-                )
-
-        return prepared_train, prepared_val, prepared_tests
-
-    def _preprocess_logits_for_metrics(
-        self, model_outputs: torch.Tensor | tuple, labels: torch.Tensor
-    ) -> torch.Tensor | tuple:
-        """Ensure that only the logits are returned from the model.
-
-        This is to avoid memory issues when the model returns hidden states as well.
-
-        Args:
-            model_outputs:
-                The raw model outputs.
-            labels:
-                The ground truth labels.
-
-        Returns:
-            The preprocessed logits.
-        """
-        if isinstance(model_outputs, tuple) and isinstance(
-            model_outputs[0], torch.Tensor
-        ):
-            model_output_tensors = [
-                model_output
-                for model_output in model_outputs
-                if isinstance(model_output, torch.Tensor)
-            ]
-            if len(model_output_tensors) == 1:
-                return model_output_tensors[0]
-            return tuple(model_output_tensors)
-        else:
-            return model_outputs
-
-    def __call__(self, *args, **kwargs):
-        """Call the benchmark method. See `benchmark` for details."""
-        return self.benchmark(*args, **kwargs)
-
-    def _process_data(self, dataset_dict: DatasetDict) -> DatasetDict:
-        """Process the data.
-
-        Args:
-            dataset_dict:
-                The dataset dictionary.
-
-        Returns:
-            The processed dataset dictionary.
-        """
-        return dataset_dict
-
-    def _get_trainer_class(self) -> Type[Trainer]:
-        """Returns the trainer class to use."""
-        return Trainer
-
-    def _get_evaluate_inputs(
-        self, dataset: "Dataset", prepared_dataset: "Dataset", metric_key_prefix: str
-    ) -> dict[str, Any]:
-        """Returns the inputs to the `Trainer.evaluate` method.
-
-        Args:
-            dataset:
-                The raw dataset.
-            prepared_dataset:
-                The prepared dataset.
-            metric_key_prefix:
-                The prefix to use for the metric keys.
-        """
-        return dict(eval_dataset=prepared_dataset, metric_key_prefix=metric_key_prefix)
-
-    @abstractmethod
-    def _preprocess_data(self, dataset: "Dataset", **kwargs) -> "Dataset":
-        """Preprocess a dataset.
-
-        Args:
-            dataset:
-                The dataset to preprocess.
-            kwargs:
-                Extra keyword arguments containing objects used in preprocessing the
-                dataset.
-
-        Returns:
-            The preprocessed dataset.
-        """
-        pass
-
-    @abstractmethod
-    def _load_data_collator(
-        self,
-        tokenizer: "Tokenizer | None" = None,
-        model: "PreTrainedModel | GenerativeModel | None" = None,
-    ):
-        """Load the data collator used to prepare samples during finetuning.
-
-        Args:
-            tokenizer:
-                A pretrained tokenizer. Can be None if the tokenizer is not used in the
-                initialisation of the data collator. Defaults to None.
-            model:
-                A pretrained model. Can be None if the model is not used in the
-                initialisation of the data collator. Defaults to None.
-
-        Returns:
-            The data collator.
-        """
-        pass
-
-    @abstractmethod
-    def _compute_metrics(
-        self,
-        model_outputs_and_labels: tuple["Predictions", "Labels"],
-        id2label: dict[int, str],
-    ) -> dict[str, float]:
-        """Compute the metrics needed for evaluation.
-
-        Args:
-            model_outputs_and_labels:
-                The first sequence contains the model outputs and the second sequence
-                contains the true labels.
-            id2label:
-                Conversion of indices to labels.
-
-        Returns:
-            A dictionary with the names of the metrics as keys and the metric values as
-            values.
-        """
-        pass
-
-    @abstractmethod
-    def _extract_few_shot_examples(
-        self, train_dataset: "Dataset", random_seed: int
-    ) -> list[dict[str, Any]]:
-        """Extract few-shot examples from the training dataset.
-
-        Args:
-            train_dataset:
-                The training dataset.
-            random_seed:
-                The random seed to use when extracting the few-shot examples.
-
-        Returns:
-            The few-shot examples.
-        """
-        pass
-
-    @abstractmethod
-    def _apply_few_shot_prompt(
-        self, examples: dict, few_shot_examples: list[dict], tokenizer: "Tokenizer"
-    ) -> dict:
-        """Apply a few-shot prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            few_shot_examples:
-                The examples to be included in the few-shot prompt.
-            tokenizer:
-                The tokenizer to use to encode the few-shot prompt.
-
-        Returns:
-            The examples with the few-shot prompt applied.
-        """
-        pass
-
-    @abstractmethod
-    def _apply_instruction_prompt(self, examples: dict, tokenizer: "Tokenizer") -> dict:
-        """Apply an instruction prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            tokenizer:
-                The tokenizer to use to encode the instruction prompt.
-
-        Returns:
-            The examples with the instruction prompt applied.
-        """
-        pass
-
-    @abstractmethod
-    def _extract_labels_from_generation(
-        self,
-        input_batch: dict[str, list],
-        model_output: "ModelOutput",
-        tokenizer: "Tokenizer",
-    ) -> list[Any]:
-        """Extract the predicted labels from the generated output.
-
-        Args:
-            input_batch:
-                The input batch, where the keys are the feature names and the values
-                are lists with the feature values.
-            model_output:
-                The raw generated output of the model.
-            tokenizer:
-                The tokenizer used together with the model.
-
-        Returns:
-            The predicted labels.
-        """
-        pass
diff --git a/src/scandeval/benchmark_modules/__init__.py b/src/scandeval/benchmark_modules/__init__.py
new file mode 100644
index 000000000..25fd72cad
--- /dev/null
+++ b/src/scandeval/benchmark_modules/__init__.py
@@ -0,0 +1,7 @@
+"""The different types of modules that can be benchmarked."""
+
+from .base import BenchmarkModule
+from .fresh import FreshEncoderModel
+from .hf import HuggingFaceEncoderModel
+from .litellm import LiteLLMModel
+from .vllm import VLLMModel
diff --git a/src/scandeval/benchmark_modules/base.py b/src/scandeval/benchmark_modules/base.py
new file mode 100644
index 000000000..422c6cc04
--- /dev/null
+++ b/src/scandeval/benchmark_modules/base.py
@@ -0,0 +1,349 @@
+"""Abstract benchmark module class that the model classes inherit from."""
+
+import collections.abc as c
+import logging
+import sys
+import typing as t
+from abc import ABC, abstractmethod
+from functools import cached_property, partial
+
+from datasets import DatasetDict
+from torch import nn
+from tqdm.auto import tqdm
+from transformers import PreTrainedTokenizer, Trainer
+
+from ..data_models import (
+    BenchmarkConfig,
+    DatasetConfig,
+    GenerativeModelOutput,
+    ModelConfig,
+    Task,
+)
+from ..enums import BatchingPreference
+from ..exceptions import NeedsEnvironmentVariable, NeedsExtraInstalled
+from ..task_utils import (
+    question_answering,
+    sequence_classification,
+    text_to_text,
+    token_classification,
+)
+from ..types import ComputeMetricsFunction, ExtractLabelsFunction
+
+logger = logging.getLogger("scandeval")
+
+
+class BenchmarkModule(ABC):
+    """Abstract class for a benchmark module.
+
+    Attributes:
+        model_config:
+            The model configuration.
+        dataset_config:
+            The dataset configuration.
+        benchmark_config:
+            The benchmark configuration.
+        is_generative:
+            Whether the model is generative.
+    """
+
+    _is_generative: bool | None
+    batching_preference: BatchingPreference
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        dataset_config: DatasetConfig,
+        benchmark_config: BenchmarkConfig,
+    ) -> None:
+        """Initialise the benchmark module.
+
+        Args:
+            model_config:
+                The model configuration.
+            dataset_config:
+                The dataset configuration.
+            benchmark_config:
+                The benchmark configuration.
+        """
+        self.model_config = model_config
+        self.dataset_config = dataset_config
+        self.benchmark_config = benchmark_config
+        self._log_metadata()
+
+    def _log_metadata(self) -> None:
+        """Log the metadata of the model."""
+        # Set logging level based on verbosity
+        if hasattr(sys, "_called_from_test"):
+            logging_level = logging.CRITICAL
+        elif self.benchmark_config.verbose:
+            logging_level = logging.DEBUG
+        else:
+            logging_level = logging.INFO
+        logger.setLevel(logging_level)
+
+        logging_msg: str = ""
+        if self.num_params < 0:
+            logging_msg += "The model has an unknown number of parameters, "
+        else:
+            logging_msg += f"The model has {self.num_params:,} parameters, "
+        if self.vocab_size < 0:
+            logging_msg += "an unknown vocabulary size, "
+        else:
+            logging_msg += f"a vocabulary size of {self.vocab_size:,}, "
+        if self.model_max_length < 0:
+            logging_msg += "and an unknown maximum sequence length."
+        else:
+            logging_msg += f"and a maximum context length of {self.model_max_length:,}."
+        logger.info(logging_msg)
+
+    def get_pytorch_module(self) -> "nn.Module":
+        """Get the underlying PyTorch module.
+
+        Returns:
+            The PyTorch module.
+        """
+        if hasattr(self, "_model"):
+            return self._model
+        raise NotImplementedError(
+            "The `get_pytorch_module` method has not been implemented for "
+            f"{self.__class__.__name__}."
+        )
+
+    def get_tokenizer(self) -> "PreTrainedTokenizer":
+        """Get the underlying tokenizer.
+
+        Returns:
+            The tokenizer.
+        """
+        if hasattr(self, "_tokenizer"):
+            return self._tokenizer
+        raise NotImplementedError(
+            "The `get_tokenizer` method has not been implemented for "
+            f"{self.__class__.__name__}."
+        )
+
+    @cached_property
+    def is_generative(self) -> bool:
+        """Whether the model is generative.
+
+        Returns:
+            Whether the model is generative.
+        """
+        if self._is_generative is not None:
+            return self._is_generative
+        raise NotImplementedError(
+            "The model type must define whether it is generative or not."
+        )
+
+    @cached_property
+    @abstractmethod
+    def num_params(self) -> int:
+        """The number of parameters in the model.
+
+        Returns:
+            The number of parameters in the model.
+        """
+        ...
+
+    @cached_property
+    @abstractmethod
+    def vocab_size(self) -> int:
+        """The vocabulary size of the model.
+
+        Returns:
+            The vocabulary size of the model.
+        """
+        ...
+
+    @cached_property
+    @abstractmethod
+    def model_max_length(self) -> int:
+        """The maximum length of the model.
+
+        Returns:
+            The maximum length of the model.
+        """
+        ...
+
+    @cached_property
+    @abstractmethod
+    def data_collator(self) -> c.Callable[[list[t.Any]], dict[str, t.Any]]:
+        """The data collator used to prepare samples during finetuning.
+
+        Returns:
+            The data collator.
+        """
+        ...
+
+    @cached_property
+    def compute_metrics(self) -> ComputeMetricsFunction:
+        """The function used to compute the metrics.
+
+        Returns:
+            The function used to compute the metrics.
+        """
+        match self.dataset_config.task.supertask:
+            case "sequence-classification":
+                return partial(
+                    sequence_classification.compute_metrics,
+                    id2label=self.dataset_config.id2label,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=self.benchmark_config,
+                )
+            case "text-to-text":
+                return partial(
+                    text_to_text.compute_metrics,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=self.benchmark_config,
+                )
+            case "token-classification":
+                return partial(
+                    token_classification.compute_metrics,
+                    has_misc_tags=self._has_misc_tags,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=self.benchmark_config,
+                )
+            case "question-answering":
+                return partial(
+                    question_answering.compute_metrics,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=self.benchmark_config,
+                )
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {self.dataset_config.task.supertask}."
+                )
+
+    @cached_property
+    @abstractmethod
+    def extract_labels_from_generation(self) -> ExtractLabelsFunction:
+        """The function used to extract the labels from the generated output.
+
+        Returns:
+            The function used to extract the labels from the generated output.
+        """
+        ...
+
+    @cached_property
+    @abstractmethod
+    def trainer_class(self) -> t.Type["Trainer"]:
+        """The Trainer class to use for finetuning.
+
+        Returns:
+            The Trainer class.
+        """
+        ...
+
+    def prepare_datasets(
+        self, datasets: list[DatasetDict], task: Task
+    ) -> list[DatasetDict]:
+        """Prepare the datasets for the model.
+
+        This includes things like tokenisation.
+
+        Args:
+            datasets:
+                The datasets to prepare.
+            task:
+                The task to prepare the datasets for.
+
+        Returns:
+            The prepared datasets.
+        """
+        for idx, dataset in enumerate(
+            tqdm(iterable=datasets, desc="Preparing datasets")
+        ):
+            prepared_dataset = self.prepare_dataset(
+                dataset=dataset, task=task, itr_idx=idx
+            )
+            if self.dataset_config.task.supertask == "token-classification":
+                labels_in_train: set[str] = {
+                    tag for tag_list in dataset["train"]["labels"] for tag in tag_list
+                }
+                self._has_misc_tags = (
+                    "B-MISC" in labels_in_train or "I-MISC" in labels_in_train
+                )
+            datasets[idx] = DatasetDict(
+                dict(
+                    train=prepared_dataset["train"],
+                    val=prepared_dataset["val"],
+                    test=prepared_dataset["test"],
+                    original_train=dataset["train"],
+                    original_val=dataset["val"],
+                    original_test=dataset["test"],
+                )
+            )
+        return datasets
+
+    @abstractmethod
+    def prepare_dataset(
+        self, dataset: DatasetDict, task: Task, itr_idx: int
+    ) -> DatasetDict:
+        """Prepare the dataset for the model.
+
+        This includes things like tokenisation.
+
+        Args:
+            dataset:
+                The dataset to prepare.
+            task:
+                The task to prepare the dataset for.
+            itr_idx:
+                The index of the dataset in the iterator.
+
+        Returns:
+            The prepared dataset.
+        """
+        ...
+
+    def generate(self, inputs: dict) -> GenerativeModelOutput:
+        """Generate outputs from the model.
+
+        Args:
+            inputs:
+                A batch of inputs to pass through the model.
+
+        Returns:
+            The generated model outputs.
+        """
+        raise NotImplementedError(
+            "The `generate` method has not been implemented for "
+            f"{self.__class__.__name__}."
+        )
+
+    @classmethod
+    @abstractmethod
+    def model_exists(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
+        """Check if a model exists.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            Whether the model exists, or an error describing why we cannot check
+            whether the model exists.
+        """
+        ...
+
+    @classmethod
+    @abstractmethod
+    def get_model_config(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> ModelConfig:
+        """Fetch the model configuration.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            The model configuration.
+        """
+        ...
diff --git a/src/scandeval/benchmark_modules/fresh.py b/src/scandeval/benchmark_modules/fresh.py
new file mode 100644
index 000000000..6d2dc4aca
--- /dev/null
+++ b/src/scandeval/benchmark_modules/fresh.py
@@ -0,0 +1,216 @@
+"""Freshly initialised encoder models."""
+
+from functools import cached_property
+from json import JSONDecodeError
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    ElectraForQuestionAnswering,
+    ElectraForSequenceClassification,
+    ElectraForTokenClassification,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    XLMRobertaForQuestionAnswering,
+    XLMRobertaForSequenceClassification,
+    XLMRobertaForTokenClassification,
+)
+
+from ..data_models import BenchmarkConfig, ModelConfig
+from ..enums import Framework, ModelType
+from ..exceptions import (
+    InvalidBenchmark,
+    InvalidModel,
+    NeedsEnvironmentVariable,
+    NeedsExtraInstalled,
+)
+from ..utils import block_terminal_output, create_model_cache_dir
+from .hf import (
+    HuggingFaceEncoderModel,
+    align_model_and_tokenizer,
+    setup_model_for_question_answering,
+)
+
+
+class FreshEncoderModel(HuggingFaceEncoderModel):
+    """A freshly initialised encoder model."""
+
+    @cached_property
+    def num_params(self) -> int:
+        """The number of parameters in the model.
+
+        Returns:
+            The number of parameters in the model.
+        """
+        match self.model_config.model_id:
+            case "fresh-xlm-roberta-base":
+                return 278_885_778
+            case "fresh-electra-small":
+                return 13_738_755
+            case _:
+                raise NotImplementedError(
+                    f"Number of parameters for model {self.model_config.model_id} is "
+                    "not implemented."
+                )
+
+    @cached_property
+    def vocab_size(self) -> int:
+        """The vocabulary size of the model.
+
+        Returns:
+            The vocabulary size of the model.
+        """
+        match self.model_config.model_id:
+            case "fresh-xlm-roberta-base":
+                return 250_002
+            case "fresh-electra-small":
+                return 32_000
+            case _:
+                raise NotImplementedError(
+                    f"Vocabulary size for model {self.model_config.model_id} is not "
+                    "implemented."
+                )
+
+    @cached_property
+    def model_max_length(self) -> int:
+        """The maximum context length of the model.
+
+        Returns:
+            The maximum context length of the model.
+        """
+        match self.model_config.model_id:
+            case "fresh-xlm-roberta-base":
+                return 512
+            case "fresh-electra-small":
+                return 128
+            case _:
+                raise NotImplementedError(
+                    f"Maximum context length for model {self.model_config.model_id} is "
+                    "not implemented."
+                )
+
+    @classmethod
+    def model_exists(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
+        """Check if a model exists.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            Whether the model exists, or an error describing why we cannot check
+            whether the model exists.
+        """
+        valid_models = ["fresh-electra-small", "fresh-xlm-roberta-base"]
+        return model_id in valid_models
+
+    @classmethod
+    def get_model_config(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> ModelConfig:
+        """Fetch the model configuration.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            The model configuration.
+        """
+        return ModelConfig(
+            model_id=model_id,
+            framework=Framework.PYTORCH,
+            task="fill-mask",
+            languages=list(),
+            revision="main",
+            model_type=ModelType.FRESH,
+            model_cache_dir=create_model_cache_dir(
+                cache_dir=benchmark_config.cache_dir, model_id=model_id
+            ),
+            adapter_base_model_id=None,
+        )
+
+    def _load_model_and_tokenizer(self) -> tuple[PreTrainedModel, PreTrainedTokenizer]:
+        """Load the model and tokenizer.
+
+        Returns:
+            The loaded model and tokenizer.
+        """
+        config: "PretrainedConfig"
+        block_terminal_output()
+
+        # Get the fresh model ID and the corresponding real model ID
+        model_id = self.model_config.model_id.replace("-", "_")
+        fresh_to_real_model_id_mapping = dict(
+            fresh_xlm_roberta_base="FacebookAI/xlm-roberta-base",
+            fresh_electra_small="google/electra-small-discriminator",
+        )
+        real_model_id = fresh_to_real_model_id_mapping[model_id]
+
+        match self.dataset_config.task.supertask:
+            case "sequence-classification":
+                model_cls_mapping = dict(
+                    fresh_xlm_roberta_base=XLMRobertaForSequenceClassification,
+                    fresh_electra_small=ElectraForSequenceClassification,
+                )
+            case "token-classification":
+                model_cls_mapping = dict(
+                    fresh_xlm_roberta_base=XLMRobertaForTokenClassification,
+                    fresh_electra_small=ElectraForTokenClassification,
+                )
+            case "question-answering":
+                model_cls_mapping = dict(
+                    fresh_xlm_roberta_base=XLMRobertaForQuestionAnswering,
+                    fresh_electra_small=ElectraForQuestionAnswering,
+                )
+            case _:
+                raise InvalidBenchmark(
+                    f"Supertask {self.dataset_config.task.supertask} is not supported "
+                    f"for model {self.model_config.model_id}"
+                )
+        model_cls = model_cls_mapping[model_id]
+
+        config = AutoConfig.from_pretrained(
+            real_model_id,
+            token=self.benchmark_config.api_key or True,
+            num_labels=self.dataset_config.num_labels,
+            id2label=self.dataset_config.id2label,
+            label2id=self.dataset_config.label2id,
+            cache_dir=self.model_config.model_cache_dir,
+        )
+        model = model_cls(config)
+
+        if self.dataset_config.task.supertask == "question-answering":
+            model = setup_model_for_question_answering(model=model)
+
+        # Load the tokenizer. If the model is a subclass of a RoBERTa model then we
+        # have to add a prefix space to the tokens, by the way the model is constructed
+        prefix_models = ["Roberta", "GPT", "Deberta"]
+        prefix = any(model_type in type(model).__name__ for model_type in prefix_models)
+        try:
+            tokenizer: "PreTrainedTokenizer" = AutoTokenizer.from_pretrained(
+                real_model_id,
+                revision=self.model_config.revision,
+                token=self.benchmark_config.api_key or True,
+                add_prefix_space=prefix,
+                cache_dir=self.model_config.model_cache_dir,
+                use_fast=True,
+                verbose=False,
+            )
+        except (JSONDecodeError, OSError):
+            raise InvalidModel(f"Could not load tokenizer for model {real_model_id!r}.")
+
+        model, tokenizer = align_model_and_tokenizer(
+            model=model,
+            tokenizer=tokenizer,
+            raise_errors=self.benchmark_config.raise_errors,
+        )
+
+        return model, tokenizer
diff --git a/src/scandeval/benchmark_modules/hf.py b/src/scandeval/benchmark_modules/hf.py
new file mode 100644
index 000000000..2587d196b
--- /dev/null
+++ b/src/scandeval/benchmark_modules/hf.py
@@ -0,0 +1,1061 @@
+"""Encoder models from the Hugging Face Hub."""
+
+import collections.abc as c
+import logging
+import typing as t
+from functools import cached_property, partial
+from json import JSONDecodeError
+from time import sleep
+
+import torch
+from datasets import DatasetDict
+from huggingface_hub import HfApi
+from huggingface_hub import whoami as hf_whoami
+from huggingface_hub.hf_api import RepositoryNotFoundError, RevisionNotFoundError
+from huggingface_hub.utils import (
+    GatedRepoError,
+    HFValidationError,
+    LocalTokenNotFoundError,
+)
+from requests.exceptions import RequestException
+from torch import nn
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    DataCollatorForTokenClassification,
+    DataCollatorWithPadding,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    Trainer,
+)
+from urllib3.exceptions import RequestError
+
+from ..constants import DUMMY_FILL_VALUE, GENERATIVE_TAGS
+from ..data_models import BenchmarkConfig, DatasetConfig, HFModelInfo, ModelConfig, Task
+from ..enums import BatchingPreference, Framework, ModelType
+from ..exceptions import (
+    HuggingFaceHubDown,
+    InvalidBenchmark,
+    InvalidModel,
+    NeedsAdditionalArgument,
+    NeedsEnvironmentVariable,
+    NeedsExtraInstalled,
+    NoInternetConnection,
+)
+from ..languages import get_all_languages
+from ..task_utils import question_answering, token_classification
+from ..types import ExtractLabelsFunction
+from ..utils import (
+    block_terminal_output,
+    create_model_cache_dir,
+    get_class_by_name,
+    internet_connection_available,
+)
+from .base import BenchmarkModule
+
+logger = logging.getLogger("scandeval")
+
+
+class HuggingFaceEncoderModel(BenchmarkModule):
+    """An encoder model from the Hugging Face Hub."""
+
+    _is_generative = False
+    batching_preference = BatchingPreference.NO_PREFERENCE
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        dataset_config: DatasetConfig,
+        benchmark_config: BenchmarkConfig,
+    ) -> None:
+        """Initialise the model.
+
+        Args:
+            model_config:
+                The model configuration.
+            dataset_config:
+                The dataset configuration.
+            benchmark_config:
+                The benchmark configuration.
+        """
+        self.model_config = model_config
+        self.dataset_config = dataset_config
+        self.benchmark_config = benchmark_config
+        model, tokenizer = self._load_model_and_tokenizer()
+        self._model: PreTrainedModel = model
+        self._tokenizer: PreTrainedTokenizer = tokenizer
+        self._log_metadata()
+
+    @cached_property
+    def num_params(self) -> int:
+        """The number of parameters in the model.
+
+        Returns:
+            The number of parameters in the model.
+        """
+        hf_api = HfApi()
+        try:
+            repo_info = hf_api.model_info(
+                repo_id=self.model_config.model_id,
+                revision=self.model_config.revision,
+                token=self.benchmark_config.api_key or True,
+            )
+        except (
+            RepositoryNotFoundError,
+            RevisionNotFoundError,
+            RequestException,
+            HFValidationError,
+        ):
+            repo_info = None
+
+        if (
+            repo_info is not None
+            and hasattr(repo_info, "safetensors")
+            and repo_info.safetensors is not None
+            and "total" in repo_info.safetensors
+        ):
+            num_params = repo_info.safetensors["total"]
+        elif (
+            repo_info is not None
+            and hasattr(repo_info, "card_data")
+            and repo_info.card_data is not None
+            and "base_model" in repo_info.card_data
+            and repo_info.card_data["base_model"] is not None
+            and len(repo_info.card_data["base_model"]) > 0
+            and (
+                base_repo_info := hf_api.repo_info(
+                    repo_info.card_data["base_model"][-1], repo_type="model"
+                )
+            )
+            is not None
+            and hasattr(base_repo_info, "safetensors")
+            and base_repo_info.safetensors is not None
+            and "total" in base_repo_info.safetensors
+        ):
+            num_params = base_repo_info.safetensors["total"]
+        elif (
+            hasattr(self._model.config, "num_params")
+            and self._model.config.num_params is not None
+        ):
+            num_params = self._model.config.num_params
+        else:
+            num_params = sum(p.numel() for p in self._model.parameters())
+        return num_params
+
+    @cached_property
+    def vocab_size(self) -> int:
+        """The vocabulary size of the model.
+
+        Returns:
+            The vocabulary size of the model.
+        """
+        if (
+            hasattr(self._model.config, "vocab_size")
+            and self._model.config.vocab_size is not None
+        ):
+            vocab_size = self._model.config.vocab_size
+        elif (
+            hasattr(self._tokenizer, "vocab_size")
+            and self._tokenizer.vocab_size is not None
+        ):
+            vocab_size = self._tokenizer.vocab_size
+        else:
+            vocab_size = -1
+        return vocab_size
+
+    @cached_property
+    def model_max_length(self) -> int:
+        """The maximum context length of the model.
+
+        Returns:
+            The maximum context length of the model.
+        """
+        all_max_lengths: list[int] = list()
+
+        # Add the registered max length of the tokenizer
+        if hasattr(
+            self._tokenizer, "model_max_length"
+        ) and self._tokenizer.model_max_length < int(1e30):
+            all_max_lengths.append(self._tokenizer.model_max_length)
+
+        # Add the max length derived from the model's input sizes
+        if hasattr(self._tokenizer, "max_model_input_sizes"):
+            all_max_lengths.extend(
+                [
+                    size
+                    for size in self._tokenizer.max_model_input_sizes.values()
+                    if size is not None
+                ]
+            )
+
+        # Add max length candidates from the model's configuration
+        candidate_config_max_lengths = [
+            "max_position_embeddings",
+            "model_max_length",
+            "max_sequence_length",
+            "sliding_window",
+            "sliding_window_size",
+        ]
+        for candidate_config_max_length in candidate_config_max_lengths:
+            if (
+                hasattr(self._model.config, candidate_config_max_length)
+                and (value := getattr(self._model.config, candidate_config_max_length))
+                is not None
+            ):
+                all_max_lengths.append(value)
+
+        # To avoid models having artificially low max lengths, we remove any max lengths
+        # that are less than 128
+        all_max_lengths = [
+            max_length for max_length in all_max_lengths if max_length >= 128
+        ]
+
+        if len(list(all_max_lengths)) > 0:
+            model_max_length = min(list(all_max_lengths))
+        else:
+            model_max_length = -1
+
+        return model_max_length
+
+    @cached_property
+    def data_collator(self) -> c.Callable[[list[t.Any]], dict[str, t.Any]]:
+        """The data collator used to prepare samples during finetuning.
+
+        Returns:
+            The data collator.
+        """
+        match self.dataset_config.task.supertask:
+            case "sequence-classification" | "text-to-text" | "question-answering":
+                return DataCollatorWithPadding(self._tokenizer, padding="longest")
+            case "token-classification":
+                return DataCollatorForTokenClassification(
+                    tokenizer=self._tokenizer, label_pad_token_id=-100
+                )
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {self.dataset_config.task.supertask}."
+                )
+
+    @cached_property
+    def extract_labels_from_generation(self) -> ExtractLabelsFunction:
+        """The function used to extract the labels from the generated output.
+
+        Returns:
+            The function used to extract the labels from the generated output.
+        """
+        raise NotImplementedError(
+            "The `extract_labels_from_generation` property has not been implemented "
+            "for Hugging Face Encoder models."
+        )
+
+    @cached_property
+    def trainer_class(self) -> t.Type["Trainer"]:
+        """The Trainer class to use for finetuning.
+
+        Returns:
+            The Trainer class.
+        """
+        match self.dataset_config.task.supertask:
+            case "sequence-classification" | "text-to-text" | "token-classification":
+                return Trainer
+            case "question-answering":
+                return question_answering.QuestionAnsweringTrainer
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {self.dataset_config.task.supertask}."
+                )
+
+    def prepare_dataset(
+        self, dataset: DatasetDict, task: Task, itr_idx: int
+    ) -> DatasetDict:
+        """Prepare the dataset for the model.
+
+        This includes things like tokenisation.
+
+        Args:
+            dataset:
+                The dataset to prepare.
+            task:
+                The task to prepare the dataset for.
+            itr_idx:
+                The index of the dataset in the iterator.
+
+        Returns:
+            The prepared dataset.
+        """
+
+        def numericalise_labels(examples: dict):
+            if "label" in examples:
+                try:
+                    examples["label"] = [
+                        self._model.config.label2id[lbl.lower()]
+                        for lbl in examples["label"]
+                    ]
+                except KeyError:
+                    raise InvalidBenchmark(
+                        f"One of the labels in the dataset, "
+                        f"{examples['label'].lower()}, does not occur in the "
+                        f"label2id dictionary {self._model.config.label2id}."
+                    )
+            return examples
+
+        def tokenise(examples: dict):
+            return self._tokenizer(text=examples["text"], truncation=True, padding=True)
+
+        match task.supertask:
+            case "sequence-classification":
+                dataset = dataset.map(
+                    numericalise_labels, batched=True, load_from_cache_file=False
+                ).map(tokenise, batched=True, load_from_cache_file=False)
+
+            case "text-to-text":
+                dataset = dataset.map(
+                    tokenise,
+                    batched=True,
+                    load_from_cache_file=False,
+                    keep_in_memory=True,
+                )
+
+            case "token-classification":
+                dataset = dataset.map(
+                    partial(
+                        token_classification.tokenize_and_align_labels,
+                        tokenizer=self._tokenizer,
+                        label2id=self._model.config.label2id,
+                    ),
+                    batched=True,
+                    load_from_cache_file=False,
+                    keep_in_memory=True,
+                )
+
+            case "question-answering":
+                dataset = DatasetDict(
+                    dict(
+                        train=dataset["train"].map(
+                            partial(
+                                question_answering.prepare_train_examples,
+                                tokenizer=self._tokenizer,
+                            ),
+                            batched=True,
+                            batch_size=10,
+                            remove_columns=dataset["test"].column_names,
+                            load_from_cache_file=False,
+                            keep_in_memory=True,
+                        ),
+                        val=dataset["val"].map(
+                            partial(
+                                question_answering.prepare_train_examples,
+                                tokenizer=self._tokenizer,
+                            ),
+                            batched=True,
+                            batch_size=10,
+                            remove_columns=dataset["test"].column_names,
+                            load_from_cache_file=False,
+                            keep_in_memory=True,
+                        ),
+                        test=dataset["test"].map(
+                            partial(
+                                question_answering.prepare_test_examples,
+                                tokenizer=self._tokenizer,
+                            ),
+                            batched=True,
+                            batch_size=10,
+                            remove_columns=dataset["test"].column_names,
+                            load_from_cache_file=False,
+                            keep_in_memory=True,
+                        ),
+                    )
+                )
+
+                # The Trainer hides the columns that are not used by the model (here
+                # `id` and `offset_mapping` which we will need for our post-processing),
+                # so we put them back
+                for split_name, split in dataset.items():
+                    dataset[split_name].set_format(
+                        type=split.format["type"], columns=list(split.features.keys())
+                    )
+
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {task.supertask}."
+                )
+
+        return dataset
+
+    @classmethod
+    def model_exists(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
+        """Check if a model exists.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            Whether the model exists, or an error describing why we cannot check
+            whether the model exists.
+        """
+        model_id, revision = (
+            model_id.split("@") if "@" in model_id else (model_id, "main")
+        )
+        model_info = get_model_repo_info(
+            model_id=model_id, revision=revision, benchmark_config=benchmark_config
+        )
+        return model_info is not None and model_info.pipeline_tag != "text-generation"
+
+    @classmethod
+    def get_model_config(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> ModelConfig:
+        """Fetch the model configuration.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            The model configuration.
+        """
+        model_id, revision = (
+            model_id.split("@") if "@" in model_id else (model_id, "main")
+        )
+        model_info = get_model_repo_info(
+            model_id=model_id, revision=revision, benchmark_config=benchmark_config
+        )
+        if model_info is None:
+            raise InvalidModel(f"The model {model_id!r} could not be found.")
+
+        framework = Framework.PYTORCH
+        if "pytorch" in model_info.tags:
+            pass
+        elif "jax" in model_info.tags:
+            framework = Framework.JAX
+        elif "spacy" in model_info.tags:
+            raise InvalidModel("SpaCy models are not supported.")
+        elif any(tag in model_info.tags for tag in {"tf", "tensorflow", "keras"}):
+            raise InvalidModel("TensorFlow/Keras models are not supported.")
+
+        language_mapping = get_all_languages()
+        language_codes = list(language_mapping.keys())
+
+        model_config = ModelConfig(
+            model_id=model_id,
+            revision=revision,
+            framework=framework,
+            task=model_info.pipeline_tag,
+            languages=[
+                language_mapping[tag]
+                for tag in model_info.tags
+                if tag in language_codes
+            ],
+            model_type=ModelType.HF_HUB_ENCODER,
+            model_cache_dir=create_model_cache_dir(
+                cache_dir=benchmark_config.cache_dir, model_id=model_id
+            ),
+            adapter_base_model_id=None,
+        )
+
+        return model_config
+
+    def _load_model_and_tokenizer(self) -> tuple[PreTrainedModel, PreTrainedTokenizer]:
+        """Load the model and tokenizer.
+
+        Returns:
+            The loaded model and tokenizer.
+        """
+        config: "PretrainedConfig"
+        block_terminal_output()
+
+        model_id = self.model_config.model_id
+        supertask = self.dataset_config.task.supertask
+        from_flax = self.model_config.framework == Framework.JAX
+        ignore_mismatched_sizes = False
+
+        config = load_hf_model_config(
+            model_id=model_id,
+            num_labels=self.dataset_config.num_labels,
+            id2label=self.dataset_config.id2label,
+            label2id=self.dataset_config.label2id,
+            revision=self.model_config.revision,
+            model_cache_dir=self.model_config.model_cache_dir,
+            api_key=self.benchmark_config.api_key,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+            run_with_cli=self.benchmark_config.run_with_cli,
+        )
+
+        model_kwargs = dict(
+            config=config,
+            from_flax=from_flax,
+            ignore_mismatched_sizes=ignore_mismatched_sizes,
+            revision=self.model_config.revision,
+            token=self.benchmark_config.api_key or True,
+            cache_dir=self.model_config.model_cache_dir,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+            torch_dtype=get_torch_dtype(
+                device=self.benchmark_config.device,
+                torch_dtype_is_set=config.to_dict().get("torch_dtype") is not None,
+                bf16_available=(
+                    torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+                ),
+            ),
+        )
+
+        # These are used when a timeout occurs
+        attempts_left = 5
+
+        model: PreTrainedModel | None = None
+        while True:
+            try:
+                # Get the model class associated with the supertask
+                model_cls_or_none: t.Type["PreTrainedModel"] | None = get_class_by_name(
+                    class_name=f"auto-model-for-{supertask}", module_name="transformers"
+                )
+
+                # If the model class could not be found then raise an error
+                if not model_cls_or_none:
+                    raise InvalidBenchmark(
+                        f"The supertask {supertask!r} does not correspond to a "
+                        "Hugging Face AutoModel type (such as "
+                        "`AutoModelForSequenceClassification`)."
+                    )
+
+                # If the model is a DeBERTaV2 model then we ensure that
+                # `pooler_hidden_size` is the same size as `hidden_size`
+                if config.model_type == "deberta-v2":
+                    config.pooler_hidden_size = config.hidden_size
+
+                try:
+                    model_or_tuple = model_cls_or_none.from_pretrained(
+                        self.model_config.model_id, **model_kwargs
+                    )
+                except (KeyError, RuntimeError) as e:
+                    if not model_kwargs["ignore_mismatched_sizes"]:
+                        logger.debug(
+                            f"{type(e).__name__} occurred during the loading "
+                            f"of the {model_id!r} model. Retrying with "
+                            "`ignore_mismatched_sizes` set to True."
+                        )
+                        model_kwargs["ignore_mismatched_sizes"] = True
+                        continue
+                    else:
+                        raise InvalidModel(str(e))
+                except (TimeoutError, RequestError):
+                    attempts_left -= 1
+                    if attempts_left == 0:
+                        raise InvalidModel(
+                            "The model could not be loaded after 5 attempts."
+                        )
+                    logger.info(f"Couldn't load the model {model_id!r}. Retrying.")
+                    sleep(5)
+                    continue
+
+                if isinstance(model_or_tuple, tuple):
+                    model = model_or_tuple[0]
+                else:
+                    model = model_or_tuple
+                break
+
+            except (OSError, ValueError) as e:
+                # If `from_flax` is False but only Flax models are available then
+                # try again with `from_flax` set to True
+                if not from_flax and "Use `from_flax=True` to load this model" in str(
+                    e
+                ):
+                    from_flax = True
+                    continue
+
+                if "checkpoint seems to be incorrect" in str(e):
+                    raise InvalidModel(
+                        f"The model {model_id!r} has an incorrect checkpoint."
+                    )
+                if "trust_remote_code" in str(e):
+                    raise InvalidModel(
+                        f"Loading the model {model_id!r} needs to trust remote code. "
+                        "If you trust the suppliers of this model, then you can enable "
+                        "this by setting the `--trust-remote-code` flag."
+                    )
+                raise InvalidModel(
+                    f"The model {model_id!r} could not be loaded. The error was {e!r}."
+                )
+
+        assert model is not None
+
+        model.eval()
+        model.to(self.benchmark_config.device)
+
+        if isinstance(model, PreTrainedModel) and supertask == "question-answering":
+            model = setup_model_for_question_answering(model=model)
+
+        tokenizer = load_tokenizer(
+            model=model,
+            model_id=model_id,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+        )
+
+        model, tokenizer = align_model_and_tokenizer(
+            model=model,
+            tokenizer=tokenizer,
+            raise_errors=self.benchmark_config.raise_errors,
+        )
+
+        return model, tokenizer
+
+
+def get_model_repo_info(
+    model_id: str, revision: str, benchmark_config: BenchmarkConfig
+) -> HFModelInfo | None:
+    """Get the information about the model from the Hugging Face Hub.
+
+    Args:
+        model_id:
+            The model ID.
+        revision:
+            The revision of the model.
+        benchmark_config:
+            The benchmark configuration.
+
+    Returns:
+        The information about the model, or None if the model could not be found.
+    """
+    hf_api = HfApi(token=benchmark_config.api_key or True)
+    model_id, revision = model_id.split("@") if "@" in model_id else (model_id, "main")
+
+    try:
+        model_info = hf_api.model_info(repo_id=model_id, revision=revision)
+
+    # Case where the model is gated; note this to the user
+    except (GatedRepoError, LocalTokenNotFoundError) as e:
+        try:
+            hf_whoami()
+            logger.warning(
+                f"Could not access the model {model_id} with the revision "
+                f"{revision}. The error was {str(e)!r}."
+            )
+            return None
+        except LocalTokenNotFoundError:
+            raise NeedsAdditionalArgument(
+                cli_argument="--api-key",
+                script_argument="api_key=<your-api-key>",
+                run_with_cli=benchmark_config.run_with_cli,
+            )
+
+    # Case where the model could not be found
+    except (RepositoryNotFoundError, HFValidationError):
+        return None
+
+    # Other internet-related errors
+    except (OSError, RequestException):
+        if internet_connection_available():
+            raise HuggingFaceHubDown()
+        else:
+            raise NoInternetConnection()
+
+    tags = model_info.tags or list()
+
+    has_base_model_tag = any(
+        tag.startswith("base_model:") and tag.count(":") == 1 for tag in tags
+    )
+    base_model_id: str | None = None
+    if has_base_model_tag:
+        base_model_id = [
+            tag.split(":")[1]
+            for tag in tags
+            if tag.startswith("base_model:") and tag.count(":") == 1
+        ][0]
+        base_model_info = hf_api.model_info(
+            repo_id=base_model_id,
+            revision=revision,
+            token=benchmark_config.api_key or True,
+        )
+        tags += base_model_info.tags or list()
+        tags = list(set(tags))
+
+    pipeline_tag = model_info.pipeline_tag
+    if pipeline_tag is None:
+        if any(tag in GENERATIVE_TAGS for tag in tags):
+            pipeline_tag = "text-generation"
+        else:
+            pipeline_tag = "fill-mask"
+
+    return HFModelInfo(
+        pipeline_tag=pipeline_tag, tags=tags, adapter_base_model_id=base_model_id
+    )
+
+
+def load_tokenizer(
+    model: "PreTrainedModel", model_id: str, trust_remote_code: bool
+) -> "PreTrainedTokenizer":
+    """Load the tokenizer.
+
+    Args:
+        model:
+            The model, which is used to determine whether to add a prefix space to
+            the tokens.
+        model_id:
+            The model identifier. Used for logging.
+        trust_remote_code:
+            Whether to trust remote code.
+
+    Returns:
+        The loaded tokenizer.
+    """
+    loading_kwargs: dict[str, bool | str] = dict(
+        use_fast=True,
+        verbose=False,
+        trust_remote_code=trust_remote_code,
+        padding_side="right",
+        truncation_side="right",
+    )
+
+    # If the model is a subclass of a certain model types then we have to add a
+    # prefix space to the tokens, by the way the model is constructed.
+    prefix_models = ["Roberta", "GPT", "Deberta"]
+    add_prefix = any(model_type in type(model).__name__ for model_type in prefix_models)
+    if add_prefix:
+        loading_kwargs["add_prefix_space"] = True
+
+    while True:
+        try:
+            return AutoTokenizer.from_pretrained(model_id, **loading_kwargs)
+        except (JSONDecodeError, OSError, TypeError):
+            raise InvalidModel(f"Could not load tokenizer for model {model_id!r}.")
+        except (TimeoutError, RequestError):
+            logger.info(f"Couldn't load tokenizer for {model_id!r}. Retrying.")
+            sleep(5)
+            continue
+
+
+def get_torch_dtype(
+    device: torch.device, torch_dtype_is_set: bool, bf16_available: bool
+) -> str | torch.dtype:
+    """Get the torch dtype, used for loading the model.
+
+    Args:
+        device:
+            The device to use.
+        torch_dtype_is_set:
+            Whether the torch data type is set in the model configuration.
+        bf16_available:
+            Whether bfloat16 is available.
+
+    Returns:
+        The torch dtype.
+    """
+    using_cuda = device == torch.device("cuda")
+    if using_cuda and torch_dtype_is_set:
+        return "auto"
+    elif using_cuda and bf16_available:
+        return torch.bfloat16
+    elif using_cuda:
+        return torch.float16
+    return torch.float32
+
+
+def load_hf_model_config(
+    model_id: str,
+    num_labels: int,
+    id2label: dict[int, str],
+    label2id: dict[str, int],
+    revision: str,
+    model_cache_dir: str | None,
+    api_key: str | None,
+    trust_remote_code: bool,
+    run_with_cli: bool,
+) -> "PretrainedConfig":
+    """Load the Hugging Face model configuration.
+
+    Args:
+        model_id:
+            The Hugging Face model ID.
+        num_labels:
+            The number of labels in the dataset.
+        id2label:
+            The mapping from label IDs to labels.
+        label2id:
+            The mapping from labels to label IDs.
+        revision:
+            The revision of the model.
+        model_cache_dir:
+            The directory to cache the model in.
+        api_key:
+            The Hugging Face API key.
+        trust_remote_code:
+            Whether to trust remote code.
+        run_with_cli:
+            Whether the script is being run with the CLI.
+
+    Returns:
+        The Hugging Face model configuration.
+    """
+    while True:
+        try:
+            config = AutoConfig.from_pretrained(
+                model_id,
+                num_labels=num_labels,
+                id2label=id2label,
+                label2id=label2id,
+                revision=revision,
+                token=api_key or True,
+                trust_remote_code=trust_remote_code,
+                cache_dir=model_cache_dir,
+            )
+            if config.eos_token_id is not None and config.pad_token_id is None:
+                if isinstance(config.eos_token_id, list):
+                    config.pad_token_id = config.eos_token_id[0]
+                else:
+                    config.pad_token_id = config.eos_token_id
+            return config
+        except KeyError as e:
+            key = e.args[0]
+            raise InvalidModel(
+                f"The model config for the model {model_id!r} could not be "
+                f"loaded, as the key {key!r} was not found in the config."
+            )
+        except OSError as e:
+            # TEMP: When the model is gated then we cannot set cache dir, for some
+            # reason (transformers==4.38.2). This should be included back in when
+            # this is fixed.
+            if "gated repo" in str(e):
+                model_cache_dir = None
+                continue
+            raise InvalidModel(
+                f"Couldn't load model config for {model_id!r}. The error was "
+                f"{e!r}. Skipping"
+            )
+        except (TimeoutError, RequestError):
+            logger.info(f"Couldn't load model config for {model_id!r}. Retrying.")
+            sleep(5)
+            continue
+        except ValueError as e:
+            requires_trust_remote_code = "trust_remote_code" in str(e)
+            if requires_trust_remote_code:
+                raise NeedsAdditionalArgument(
+                    cli_argument="--trust-remote-code",
+                    script_argument="trust_remote_code=True",
+                    run_with_cli=run_with_cli,
+                )
+            raise e
+
+
+def setup_model_for_question_answering(model: "PreTrainedModel") -> "PreTrainedModel":
+    """Setup a model for question answering.
+
+    Args:
+        model:
+            The model to setup.
+
+    Returns:
+        The setup model.
+    """
+    # Get the models' token type embedding children, if they exist
+    children = get_children_of_module(name="model", module=model)
+
+    # If the model has token type embeddings then get them
+    if children:
+        # Get the list of attributes that are token type embeddings
+        attribute_list = list()
+        done = False
+        while not done:
+            for key, value in children.items():
+                attribute_list.append(key)
+                if isinstance(value, dict):
+                    children = value
+                else:
+                    done = True
+                break
+
+        # Get the token type embeddings
+        token_type_embeddings = model
+        for attribute in attribute_list:
+            token_type_embeddings = getattr(token_type_embeddings, attribute)
+
+        # If the token type embeddings has shape (1, ...) then set the shape to
+        # (2, ...) by randomly initializing the second token type embedding
+        if token_type_embeddings.weight.data.shape[0] == 1:
+            token_type_embeddings.weight.data = torch.cat(
+                (
+                    token_type_embeddings.weight.data,
+                    torch.rand_like(token_type_embeddings.weight.data),
+                ),
+                dim=0,
+            )
+            token_type_embeddings.num_embeddings = 2
+
+        # Set the model config to use the new type vocab size
+        model.config.type_vocab_size = 2
+
+    return model
+
+
+def get_children_of_module(
+    name: str, module: nn.Module
+) -> nn.Module | dict[str, t.Any] | None:
+    """Get the children of a module.
+
+    Args:
+        name:
+            The name of the module.
+        module:
+            The module to get the children of.
+
+    Returns:
+        The children of the module, or None if the module has no children.
+    """
+    if len(list(module.children())) == 0:
+        if name == "token_type_embeddings":
+            return module
+        else:
+            return None
+    else:
+        submodules = dict()
+        for subname, submodule in module.named_children():
+            children = get_children_of_module(name=subname, module=submodule)
+            if children:
+                submodules[subname] = children
+        return submodules
+
+
+def align_model_and_tokenizer(
+    model: "PreTrainedModel",
+    tokenizer: "PreTrainedTokenizer",
+    raise_errors: bool = False,
+) -> tuple["PreTrainedModel", "PreTrainedTokenizer"]:
+    """Aligns the model and the tokenizer.
+
+    Args:
+        model:
+            The model to fix.
+        tokenizer:
+            The tokenizer to fix.
+        raise_errors:
+            Whether to raise errors instead of trying to fix them silently.
+
+    Returns:
+        The fixed model and tokenizer.
+    """
+    model_max_length = get_model_max_length(model=model, tokenizer=tokenizer)
+
+    # Ensure that the model max length is at least 5,000, to avoid OOM errors
+    model_max_length = min(model_max_length, 5_000)
+
+    if model_max_length > 0:
+        tokenizer.model_max_length = model_max_length
+    else:
+        tokenizer.model_max_length = 512
+
+    # If we're not dealing with a generative model then we move it to CPU to avoid OOM
+    # errors
+    device: torch.device = torch.device("cpu")
+    model_device = model.device
+    model.to(device)
+
+    # Manually check that this model max length is valid for the model, and adjust
+    # otherwise
+    initial_max_length = tokenizer.model_max_length
+    for max_length in range(initial_max_length, 0, -1):
+        tokenizer.model_max_length = max_length
+        dummy_inputs = torch.full(
+            size=(1, max_length),
+            fill_value=DUMMY_FILL_VALUE,
+            dtype=torch.long,
+            device=device,
+        )
+
+        with torch.inference_mode():
+            try:
+                model(dummy_inputs)
+                break
+
+            # This happens if `max_length` is too large
+            except IndexError:
+                continue
+
+    # If there is a mismatch between the vocab size according to the tokenizer and
+    # the vocab size according to the model, we raise an error
+    if hasattr(model.config, "vocab_size") and hasattr(tokenizer, "vocab_size"):
+        if model.config.vocab_size < tokenizer.vocab_size:
+            if raise_errors:
+                raise InvalidModel(
+                    "The vocab size of the tokenizer is larger than the vocab size of "
+                    "the model. As the --raise-errors option was specified, the "
+                    "embeddings of the model will not be automatically adjusted."
+                )
+            if hasattr(model, "resize_token_embeddings"):
+                model.resize_token_embeddings(new_num_tokens=tokenizer.vocab_size + 1)
+
+    if tokenizer.bos_token is None and tokenizer.eos_token is not None:
+        tokenizer.bos_token = tokenizer.eos_token
+        tokenizer.bos_token_id = tokenizer.eos_token_id
+
+    model.to(model_device)
+
+    return model, tokenizer
+
+
+def get_model_max_length(
+    model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer"
+) -> int:
+    """Get the maximum context length of a model.
+
+    Args:
+        model:
+            The model.
+        tokenizer:
+            The tokenizer.
+
+    Returns:
+        The maximum context length.
+    """
+    all_max_lengths: list[int] = list()
+
+    if tokenizer is not None:
+        # Add the registered max length of the tokenizer
+        if hasattr(tokenizer, "model_max_length") and tokenizer.model_max_length < int(
+            1e30
+        ):
+            all_max_lengths.append(tokenizer.model_max_length)
+
+        # Add the max length derived from the model's input sizes
+        if hasattr(tokenizer, "max_model_input_sizes"):
+            all_max_lengths.extend(
+                [
+                    size
+                    for size in tokenizer.max_model_input_sizes.values()
+                    if size is not None
+                ]
+            )
+
+    # Add max length candidates from the model's configuration
+    candidate_config_max_lengths = [
+        "max_position_embeddings",
+        "model_max_length",
+        "max_sequence_length",
+        "sliding_window",
+        "sliding_window_size",
+    ]
+    for candidate_config_max_length in candidate_config_max_lengths:
+        if (
+            hasattr(model.config, candidate_config_max_length)
+            and (value := getattr(model.config, candidate_config_max_length))
+            is not None
+        ):
+            all_max_lengths.append(value)
+
+    # To avoid models having artificially low max lengths, we remove any max lengths
+    # that are less than 128
+    all_max_lengths = [
+        max_length for max_length in all_max_lengths if max_length >= 128
+    ]
+
+    if len(list(all_max_lengths)) > 0:
+        model_max_length = min(list(all_max_lengths))
+    else:
+        model_max_length = -1
+
+    return model_max_length
diff --git a/src/scandeval/benchmark_modules/litellm.py b/src/scandeval/benchmark_modules/litellm.py
new file mode 100644
index 000000000..6a02f0689
--- /dev/null
+++ b/src/scandeval/benchmark_modules/litellm.py
@@ -0,0 +1,714 @@
+"""Generative models from an inference API, using the LiteLLM framework."""
+
+import collections.abc as c
+import itertools as it
+import json
+import logging
+import random
+import re
+import typing as t
+from functools import cached_property, partial
+from time import sleep
+
+import litellm
+from datasets import DatasetDict
+from litellm.exceptions import (
+    APIError,
+    AuthenticationError,
+    BadRequestError,
+    InternalServerError,
+    NotFoundError,
+)
+from litellm.types.utils import ModelResponse
+from transformers import Trainer
+
+from ..constants import MAX_LOGPROBS, SUPERTASKS_USING_LOGPROBS, TASKS_USING_JSON
+from ..data_models import BenchmarkConfig, GenerativeModelOutput, ModelConfig, Task
+from ..enums import BatchingPreference, Framework, ModelType
+from ..exceptions import (
+    InvalidBenchmark,
+    NeedsAdditionalArgument,
+    NeedsEnvironmentVariable,
+    NeedsExtraInstalled,
+)
+from ..task_utils import (
+    question_answering,
+    sequence_classification,
+    text_to_text,
+    token_classification,
+)
+from ..types import ExtractLabelsFunction
+from ..utils import create_model_cache_dir
+from .base import BenchmarkModule
+
+logger = logging.getLogger("scandeval")
+
+
+VOCAB_SIZE_MAPPING = {
+    # OpenAI models
+    "(text-)?(ada|babbage|curie|davinci)(-001)?": 50_257,
+    "(code|text)-davinci-00[2-9]": 50_281,
+    "gpt-3.5-turbo(-16k)?(-[0-9]{4})?": 100_256,
+    "gpt-4-(32k)?(-[0-9]{4})?": 100_256,
+    "gpt-4-[0-9]{4}-preview": 100_256,
+    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 100_256,
+    "gpt-4-(vision|turbo)(-preview)?": 100_256,
+    "gpt-3.5-turbo-instruct(-[0-9]{4})?": 100_256,
+    "gpt-4o(-mini)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 200_019,
+    # Anthropic models
+    "claude-[1-9](-[1-9])?-(opus|sonnet|haiku)-[0-9]{8}": -1,
+}
+
+
+MODEL_MAX_LENGTH_MAPPING = {
+    # OpenAI models
+    "(text-)?(ada|babbage|curie|davinci)(-001)?": 2_050,
+    "text-davinci-00[2-9]": 4_098,
+    "code-davinci-00[1-9]": 8_002,
+    "gpt-3.5-turbo-0613": 4_096,
+    "gpt-3.5-turbo(-[0-9]{4})?": 16_385,
+    "gpt-3.5-turbo-16k(-[0-9]{4})?": 16_384,
+    "gpt-4(-[0-9]{4})?": 8_191,
+    "gpt-4-32k(-[0-9]{4})?": 32_767,
+    "gpt-4-[0-9]{4}-preview": 128_000,
+    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 128_000,
+    "gpt-4-(vision|turbo)(-preview)?": 128_000,
+    "gpt-3.5-turbo-instruct(-[0-9]{4})?": 4_095,
+    "gpt-4o(-mini)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 128_000,
+    # Anthropic models
+    "claude-[1-9](-[1-9])?-(opus|sonnet|haiku)-[0-9]{8}": 200_000,
+}
+
+
+NUM_PARAMS_MAPPING = {
+    "(text-)?ada(-001)?": 350_000_000,
+    "(text-)?babbage(-001)?": 3_000_000_000,
+    "(text-)?curie(-001)?": 13_000_000_000,
+    "((text|code)-)?davinci(-00[1-9])?": 175_000_000_000,
+    "gpt-(3.5|4)-turbo-((16|32)k)?(-[0-9]{4})?": -1,
+    "gpt-4-[0-9]{4}-preview": -1,
+    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
+    "gpt-4-(vision|turbo)(-preview)?": -1,
+    "gpt-3.5-turbo-instruct(-[0-9]{4})?": -1,
+    "gpt-4o(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
+    "gpt-4o-mini(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
+    # Anthropic models
+    "claude-[1-9](-[1-9])?-(opus|sonnet|haiku)-[0-9]{8}": -1,
+}
+
+
+class LiteLLMModel(BenchmarkModule):
+    """A generative model from LiteLLM."""
+
+    _is_generative = True
+    batching_preference = BatchingPreference.SINGLE_SAMPLE
+
+    def generate(self, inputs: dict) -> GenerativeModelOutput:
+        """Generate outputs from the model.
+
+        Args:
+            inputs:
+                A batch of inputs to pass through the model.
+
+        Returns:
+            The generated model outputs.
+        """
+        assert "messages" in inputs, "The input must contain a 'messages' key."
+        assert (
+            len(inputs["messages"]) == 1
+        ), "API models only support single-sample batching."
+        messages = inputs["messages"][0]
+
+        generation_kwargs: dict[str, t.Any] = dict(
+            model=self.model_config.model_id,
+            max_tokens=self.dataset_config.max_generated_tokens,
+            stop=["\n\n"],
+            temperature=0.0,
+            seed=4242,
+        )
+
+        if self.dataset_config.task.supertask in SUPERTASKS_USING_LOGPROBS:
+            generation_kwargs["logprobs"] = True
+            generation_kwargs["top_logprobs"] = MAX_LOGPROBS
+
+        if self.dataset_config.task.name in TASKS_USING_JSON:
+            assert (
+                "json" in messages[0]["content"]
+            ), "Prompt must contain 'json' for JSON tasks."
+            generation_kwargs["response_format"] = dict(type="json_object")
+
+        # This drops generation kwargs that are not supported by the model
+        litellm.drop_params = True
+
+        # 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):
+            try:
+                model_response = litellm.completion(
+                    messages=messages, max_retries=3, **generation_kwargs
+                )
+                break
+            except BadRequestError as e:
+                if "stop_sequences" not in str(e).lower():
+                    raise InvalidBenchmark(
+                        f"Failed to generate text. The error message was: {e}"
+                    )
+                generation_kwargs["stop"] = None
+            except APIError as e:
+                raise InvalidBenchmark(
+                    f"Failed to generate text. The error message was: {e}"
+                )
+            except InternalServerError as e:
+                if "overloaded" not in str(e).lower():
+                    raise InvalidBenchmark(
+                        f"Failed to generate text. The error message was: {e}"
+                    )
+                sleep(1)
+                continue
+            except AuthenticationError:
+                raise NeedsAdditionalArgument(
+                    cli_argument="--api-key",
+                    script_argument="api_key=<your-api-key>",
+                    run_with_cli=self.benchmark_config.run_with_cli,
+                )
+        else:
+            raise InvalidBenchmark(
+                message=f"Failed to generate text, after {num_attempts} attempts."
+            )
+
+        assert isinstance(model_response, ModelResponse)
+        model_response_choices = model_response.choices[0]
+        assert isinstance(model_response_choices, litellm.Choices)
+        generation_output = model_response_choices.message["content"].strip()
+
+        # Structure the model output as a GenerativeModelOutput object
+        model_output = GenerativeModelOutput(sequences=[generation_output])
+        if hasattr(model_response_choices, "logprobs"):
+            logprobs_list: list[list[tuple[str, float]]] = [
+                [(dct["token"], dct["logprob"]) for dct in content["top_logprobs"]]
+                for content in model_response_choices.logprobs["content"]
+            ]
+            model_output.scores = [logprobs_list]
+
+        return model_output
+
+    @cached_property
+    def num_params(self) -> int:
+        """The number of parameters in the model.
+
+        Returns:
+            The number of parameters in the model.
+        """
+        for key, value in NUM_PARAMS_MAPPING.items():
+            if re.match(pattern=key, string=self.model_config.model_id) is not None:
+                return value
+        return -1
+
+    @cached_property
+    def vocab_size(self) -> int:
+        """The vocabulary size of the model.
+
+        Returns:
+            The vocabulary size of the model.
+        """
+        for key, value in VOCAB_SIZE_MAPPING.items():
+            if re.match(pattern=key, string=self.model_config.model_id) is not None:
+                return value
+        return -1
+
+    @cached_property
+    def model_max_length(self) -> int:
+        """The maximum length of the model.
+
+        Returns:
+            The maximum length of the model.
+        """
+        for key, value in MODEL_MAX_LENGTH_MAPPING.items():
+            if re.match(pattern=key, string=self.model_config.model_id) is not None:
+                return value
+        return -1
+
+    @cached_property
+    def data_collator(self) -> c.Callable[[list[t.Any]], dict[str, t.Any]]:
+        """The data collator used to prepare samples during finetuning.
+
+        Returns:
+            The data collator.
+        """
+        raise NotImplementedError(
+            "The `data_collator` property has not been implemented for LiteLLM models."
+        )
+
+    @cached_property
+    def extract_labels_from_generation(self) -> ExtractLabelsFunction:
+        """The function used to extract the labels from the generated output.
+
+        Returns:
+            The function used to extract the labels from the generated output.
+        """
+        match self.dataset_config.task.supertask:
+            case "sequence-classification":
+                return partial(
+                    sequence_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "text-to-text":
+                return text_to_text.extract_labels_from_generation
+            case "token-classification":
+                return partial(
+                    token_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "question-answering":
+                return question_answering.extract_labels_from_generation
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {self.dataset_config.task.supertask}."
+                )
+
+    @cached_property
+    def trainer_class(self) -> t.Type["Trainer"]:
+        """The Trainer class to use for finetuning.
+
+        Returns:
+            The Trainer class.
+        """
+        raise NotImplementedError(
+            "The `trainer_class` property has not been implemented for LiteLLM models."
+        )
+
+    @classmethod
+    def model_exists(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
+        """Check if a model exists.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            Whether the model exists, or an error describing why we cannot check
+            whether the model exists.
+        """
+        if model_id in litellm.model_list:
+            return True
+
+        try:
+            litellm.completion(
+                messages=[dict(role="user", content="X")], model=model_id, max_tokens=1
+            )
+            return True
+        except (BadRequestError, NotFoundError):
+            candidate_models = [
+                candidate_model_id
+                for candidate_model_id in litellm.model_list
+                if candidate_model_id.startswith(model_id)
+            ]
+            match len(candidate_models):
+                case 0:
+                    pass
+                case 1:
+                    logger.warning(
+                        f"Could not find the model ID {model_id!r}. Did you mean "
+                        f"{candidate_models[0]!r}?"
+                    )
+                case _:
+                    candidate_models_str = "', '".join(candidate_models)
+                    logger.warning(
+                        f"Could not find the model ID {model_id!r}. Did you mean any of "
+                        f"the following model IDs: '{candidate_models_str}'?"
+                    )
+            return False
+
+    @classmethod
+    def get_model_config(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> ModelConfig:
+        """Fetch the model configuration.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            The model configuration.
+        """
+        return ModelConfig(
+            model_id=model_id,
+            revision="main",
+            framework=Framework.API,
+            task="text-generation",
+            languages=list(),
+            model_cache_dir=create_model_cache_dir(
+                cache_dir=benchmark_config.cache_dir, model_id=model_id
+            ),
+            model_type=ModelType.API,
+            adapter_base_model_id=None,
+        )
+
+    def prepare_dataset(
+        self, dataset: DatasetDict, task: Task, itr_idx: int
+    ) -> DatasetDict:
+        """Prepare the dataset for the model.
+
+        This includes things like tokenisation.
+
+        Args:
+            dataset:
+                The dataset to prepare.
+            task:
+                The task to prepare the dataset for.
+            itr_idx:
+                The index of the dataset in the iterator.
+
+        Returns:
+            The prepared dataset.
+        """
+        if task.supertask == "question-answering":
+            dataset = dataset.map(
+                lambda examples: dict(
+                    label=[
+                        dict(
+                            id=id,
+                            answers=dict(
+                                answer_start=answer_dct["answer_start"],
+                                text=[
+                                    answer_text.lower()
+                                    for answer_text in answer_dct["text"]
+                                ],
+                            ),
+                        )
+                        for id, answer_dct in zip(examples["id"], examples["answers"])
+                    ]
+                ),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+
+        if self.benchmark_config.few_shot:
+            few_shot_examples = self._extract_few_shot_examples(
+                dataset=dataset, task=task, itr_idx=itr_idx
+            )
+            dataset["test"] = dataset["test"].map(
+                partial(
+                    self._apply_few_shot_prompt,
+                    few_shot_examples=few_shot_examples,
+                    task=task,
+                ),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+        else:
+            dataset["test"] = dataset["test"].map(
+                partial(self._apply_instruction_prompt, task=task),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+        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.supertask:
+            case "sequence-classification":
+                labels = it.cycle(self.dataset_config.task.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 "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 "token-classification":
+                labels = it.cycle(
+                    [
+                        label.lower()
+                        for label in self.dataset_config.task.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["text"] != example["text"]
+                    )
+
+            case "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 supertask: {task.supertask}."
+                )
+
+        random.seed(random_seed)
+        random.shuffle(few_shot_examples)
+        return few_shot_examples
+
+    def _apply_few_shot_prompt(
+        self,
+        examples: dict[str, t.Any],
+        few_shot_examples: list[dict[str, t.Any]],
+        task: Task,
+    ) -> dict[str, t.Any]:
+        """Apply few-shot examples to an example.
+
+        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 split_section(section: str) -> tuple[str, str]:
+            """Split a section of the prompt to user and assistant messages."""
+            user_part = "\n".join(section.split("\n")[:-1])
+            assistant_part = section.split("\n")[-1]
+            return user_part, assistant_part
+
+        match task.supertask:
+            case "sequence-classification":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=example["text"].replace("\n", " ").strip(),
+                        label=example["label"].replace("\n", " ").strip(),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=text.replace("\n", " ").strip(), label=""
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "text-to-text":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=example["text"].replace("\n", " ").strip(),
+                        target_text=example["target_text"].replace("\n", " ").strip(),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=text.replace("\n", " ").strip(), target_text=""
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "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 = [
+                    self.dataset_config.prompt_template.format(
+                        text=" ".join(example["tokens"]).replace("\n", " ").strip(),
+                        label=create_label(example=example),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=" ".join(tokens).replace("\n", " ").strip(), label=""
+                    )
+                    for tokens in examples["tokens"]
+                ]
+
+            case "question-answering":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        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 = [
+                    self.dataset_config.prompt_template.format(
+                        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 supertask: {task.supertask}."
+                )
+
+        few_shot_messages = [
+            dict(role=role, content=content.split(":", 1)[1].strip())
+            for section in few_shot_sections
+            for role, content in zip(
+                it.cycle(["user", "assistant"]), split_section(section=section)
+            )
+            if content.split(":", 1)[1].strip() != ""
+        ]
+
+        if self.dataset_config.prompt_prefix:
+            few_shot_messages[0]["content"] = (
+                self.dataset_config.prompt_prefix
+                + "\n\n"
+                + few_shot_messages[0]["content"]
+            )
+
+        examples["messages"] = [
+            few_shot_messages
+            + [
+                dict(
+                    role="user",
+                    content=split_section(section=new_section)[0]
+                    .split(":", 1)[1]
+                    .strip(),
+                )
+            ]
+            for new_section in new_sections
+        ]
+
+        return examples
+
+    def _apply_instruction_prompt(
+        self, examples: dict[str, t.Any], task: Task
+    ) -> dict[str, t.Any]:
+        """Apply instruction prompts to an example.
+
+        Args:
+            examples:
+                The examples to apply the instruction prompts to.
+            task:
+                The task that is being benchmarked.
+
+        Returns:
+            The example with the instruction prompts applied.
+        """
+        match task.supertask:
+            case "sequence-classification" | "text-to-text" | "token-classification":
+                prompts = [
+                    self.dataset_config.instruction_prompt.format(
+                        text=re.sub(r"\n+", "\n", text).strip()
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "question-answering":
+                prompts = [
+                    self.dataset_config.instruction_prompt.format(
+                        text=re.sub(pattern=r"\n+", repl="\n", string=text).strip(),
+                        question=question.strip(),
+                    )
+                    for (text, question) in zip(
+                        examples["context"], examples["question"]
+                    )
+                ]
+
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {task.supertask}."
+                )
+
+        examples["messages"] = [
+            [dict(role="user", content=prompt)] for prompt in prompts
+        ]
+        return examples
diff --git a/src/scandeval/benchmark_modules/vllm.py b/src/scandeval/benchmark_modules/vllm.py
new file mode 100644
index 000000000..ef0b66791
--- /dev/null
+++ b/src/scandeval/benchmark_modules/vllm.py
@@ -0,0 +1,900 @@
+"""Generative models using the vLLM inference framework."""
+
+import collections.abc as c
+import importlib.util
+import itertools as it
+import json
+import logging
+import random
+import re
+import sys
+import typing as t
+from functools import cached_property, partial
+from pathlib import Path
+from time import sleep
+from types import MethodType
+
+import torch
+from datasets import DatasetDict
+from huggingface_hub import snapshot_download
+from tqdm.auto import tqdm
+from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizer, Trainer
+from urllib3.exceptions import RequestError
+
+from ..constants import MAX_LOGPROBS, SUPERTASKS_USING_LOGPROBS, TASKS_USING_JSON
+from ..data_models import (
+    BenchmarkConfig,
+    DatasetConfig,
+    GenerativeModelOutput,
+    ModelConfig,
+    Task,
+)
+from ..enums import BatchingPreference, Framework, ModelType
+from ..exceptions import (
+    InvalidBenchmark,
+    InvalidModel,
+    NeedsEnvironmentVariable,
+    NeedsExtraInstalled,
+)
+from ..languages import get_all_languages
+from ..structured_generation_utils import get_ner_logits_processors
+from ..task_utils import (
+    question_answering,
+    sequence_classification,
+    text_to_text,
+    token_classification,
+)
+from ..types import ExtractLabelsFunction
+from ..utils import clear_memory, create_model_cache_dir, get_end_of_chat_token_ids
+from .hf import HuggingFaceEncoderModel, get_model_repo_info
+
+if t.TYPE_CHECKING or importlib.util.find_spec("vllm") is not None:
+    from vllm import LLM, RequestOutput, SamplingParams
+    from vllm.lora.request import LoRARequest
+
+    try:
+        from vllm.model_executor.parallel_utils.parallel_state import (
+            destroy_model_parallel,
+        )
+    except ImportError:
+        from vllm.distributed.parallel_state import destroy_model_parallel
+
+if t.TYPE_CHECKING or importlib.util.find_spec("ray") is not None:
+    import ray
+
+logger = logging.getLogger("scandeval")
+
+
+class VLLMModel(HuggingFaceEncoderModel):
+    """A generative model using the vLLM inference framework."""
+
+    _is_generative = True
+    batching_preference = BatchingPreference.ALL_AT_ONCE
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        dataset_config: DatasetConfig,
+        benchmark_config: BenchmarkConfig,
+    ) -> None:
+        """Initialise the vLLM model.
+
+        Args:
+            model_config:
+                The model configuration.
+            dataset_config:
+                The dataset configuration.
+            benchmark_config:
+                The benchmark configuration.
+        """
+        if (
+            importlib.util.find_spec("vllm") is None
+            or importlib.util.find_spec("ray") is None
+        ):
+            raise NeedsExtraInstalled(extra="generative")
+
+        self.model_config = model_config
+        self.dataset_config = dataset_config
+        self.benchmark_config = benchmark_config
+        self.output_scores = (
+            self.dataset_config.task.supertask in SUPERTASKS_USING_LOGPROBS
+        )
+
+        model, tokenizer = self._load_model_and_tokenizer()
+        self._model: LLM = model
+        self._tokenizer: PreTrainedTokenizer = tokenizer
+
+        self.lora_request: LoRARequest | None = None
+        if self.model_config.adapter_base_model_id is not None:
+            adapter_path = snapshot_download(
+                repo_id=self.model_config.model_id,
+                cache_dir=Path(self.model_config.model_cache_dir),
+            )
+            self.lora_request = LoRARequest(
+                lora_name="adapter", lora_int_id=1, lora_path=adapter_path
+            )
+
+        self._log_metadata()
+
+    @cached_property
+    def extract_labels_from_generation(self) -> ExtractLabelsFunction:
+        """The function used to extract the labels from the generated output.
+
+        Returns:
+            The function used to extract the labels from the generated output.
+        """
+        match self.dataset_config.task.supertask:
+            case "sequence-classification":
+                return partial(
+                    sequence_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "text-to-text":
+                return text_to_text.extract_labels_from_generation
+            case "token-classification":
+                return partial(
+                    token_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "question-answering":
+                return question_answering.extract_labels_from_generation
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {self.dataset_config.task.supertask}."
+                )
+
+    def prepare_dataset(
+        self, dataset: DatasetDict, task: Task, itr_idx: int
+    ) -> DatasetDict:
+        """Prepare the dataset for the model.
+
+        This includes things like tokenisation.
+
+        Args:
+            dataset:
+                The dataset to prepare.
+            task:
+                The task to prepare the dataset for.
+            itr_idx:
+                The index of the dataset in the iterator.
+
+        Returns:
+            The prepared dataset.
+        """
+        if task.supertask == "question-answering":
+            dataset = dataset.map(
+                lambda examples: dict(
+                    label=[
+                        dict(
+                            id=id,
+                            answers=dict(
+                                answer_start=answer_dct["answer_start"],
+                                text=[
+                                    answer_text.lower()
+                                    for answer_text in answer_dct["text"]
+                                ],
+                            ),
+                        )
+                        for id, answer_dct in zip(examples["id"], examples["answers"])
+                    ]
+                ),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+
+        if self.benchmark_config.few_shot:
+            few_shot_examples = self._extract_few_shot_examples(
+                dataset=dataset, task=task, itr_idx=itr_idx
+            )
+            dataset["test"] = dataset["test"].map(
+                partial(
+                    self._apply_few_shot_prompt,
+                    few_shot_examples=few_shot_examples,
+                    task=task,
+                ),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+        else:
+            dataset["test"] = dataset["test"].map(
+                partial(self._apply_instruction_prompt, task=task),
+                batched=True,
+                load_from_cache_file=False,
+                keep_in_memory=True,
+            )
+        return dataset
+
+    def generate(self, inputs: dict) -> GenerativeModelOutput:
+        """Generate outputs from the model.
+
+        Args:
+            inputs:
+                A batch of inputs to pass through the model.
+
+        Returns:
+            The generated model outputs.
+        """
+        # Define which tokens to use as stopping criteria. We want to use the padding
+        # token, end-of-sentence token, and a double newline (since these separate the
+        # few-shot examples in the input)
+        stop_tokens: list[str] = ["\n\n"]
+        if self._tokenizer.pad_token_id is not None:
+            stop_tokens.append(self._tokenizer.pad_token)
+        if self._tokenizer.eos_token_id is not None:
+            stop_tokens.append(self._tokenizer.eos_token)
+            if self._tokenizer.pad_token_id is None:
+                self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
+                self._tokenizer.pad_token = self._tokenizer.eos_token
+        if (
+            self._tokenizer.bos_token_id is not None
+            and self._tokenizer.pad_token_id is None
+        ):
+            self._tokenizer.pad_token_id = self._tokenizer.bos_token_id
+            self._tokenizer.pad_token = self._tokenizer.bos_token
+        assert self._tokenizer.pad_token_id is not None
+
+        # Add end of chat token as a stopping token, if it exists
+        end_of_chat_token_ids = get_end_of_chat_token_ids(tokenizer=self._tokenizer)
+        if end_of_chat_token_ids is not None:
+            end_of_chat_token = self._tokenizer.decode(end_of_chat_token_ids).strip()
+            if end_of_chat_token:
+                stop_tokens.append(end_of_chat_token)
+
+        if self.dataset_config.task.name in TASKS_USING_JSON:
+            ner_tag_names = list(self.dataset_config.prompt_label_mapping.values())
+            logits_processors = get_ner_logits_processors(
+                ner_tag_names=ner_tag_names, llm=self._model
+            )
+        else:
+            logits_processors = None
+
+        # Define the parameters used for vLLM generation
+        max_tokens: int = self.dataset_config.max_generated_tokens
+        sampling_params = SamplingParams(
+            max_tokens=max_tokens,
+            logprobs=MAX_LOGPROBS if self.output_scores else None,
+            temperature=0.0,
+            stop=[stop_token for stop_token in stop_tokens if stop_token],
+            logits_processors=logits_processors,
+        )
+
+        # If any of the prompts are empty then we need to replace them with a BOS token
+        # so that the vLLM model can generate from them
+        prompts = inputs["text"]
+        if any(len(prompt) == 0 for prompt in prompts):
+            logger.debug("Found empty prompts, replacing with BOS token.")
+            prompts = [
+                prompt if len(prompt) > 0 else self._tokenizer.bos_token
+                for prompt in prompts
+            ]
+
+        # Generate sequences using vLLM
+        input_is_a_test = len(prompts) == 1 and len(set(prompts[0])) == 1
+        raw_outputs = self._model.generate(
+            prompts=prompts,
+            sampling_params=sampling_params,
+            use_tqdm=(not input_is_a_test),
+            lora_request=self.lora_request,
+        )
+        completions = self._tokenizer.batch_decode(
+            sequences=[
+                torch.LongTensor(output.outputs[0].token_ids) for output in raw_outputs
+            ],
+            skip_special_tokens=True,
+        )
+        completions = [completion.strip() for completion in completions]
+
+        # Add logprobs scores to the output
+        if self.output_scores:
+            scores: list[list[list[tuple[str, float]]]] = [
+                [
+                    [
+                        (obj.decoded_token, obj.logprob)
+                        for obj in token_logprobs_dict.values()
+                    ]
+                    for token_logprobs_dict in raw_output.outputs[0].logprobs
+                ]
+                for raw_output in raw_outputs
+            ]
+            output = GenerativeModelOutput(sequences=completions, scores=scores)
+        else:
+            output = GenerativeModelOutput(sequences=completions)
+
+        return output
+
+    @classmethod
+    def model_exists(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
+        """Check if a model exists.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            Whether the model exists, or an error describing why we cannot check
+            whether the model exists.
+        """
+        model_id, revision = (
+            model_id.split("@") if "@" in model_id else (model_id, "main")
+        )
+        model_info = get_model_repo_info(
+            model_id=model_id, revision=revision, benchmark_config=benchmark_config
+        )
+        return model_info is not None and model_info.pipeline_tag == "text-generation"
+
+    @classmethod
+    def get_model_config(
+        cls, model_id: str, benchmark_config: BenchmarkConfig
+    ) -> ModelConfig:
+        """Fetch the model configuration.
+
+        Args:
+            model_id:
+                The model ID.
+            benchmark_config:
+                The benchmark configuration.
+
+        Returns:
+            The model configuration.
+        """
+        model_id, revision = (
+            model_id.split("@") if "@" in model_id else (model_id, "main")
+        )
+        model_info = get_model_repo_info(
+            model_id=model_id, revision=revision, benchmark_config=benchmark_config
+        )
+        if model_info is None:
+            raise InvalidModel(f"The model {model_id!r} could not be found.")
+
+        framework = Framework.PYTORCH
+        if "pytorch" in model_info.tags:
+            pass
+        elif "jax" in model_info.tags:
+            framework = Framework.JAX
+        elif "spacy" in model_info.tags:
+            raise InvalidModel("SpaCy models are not supported.")
+        elif any(tag in model_info.tags for tag in {"tf", "tensorflow", "keras"}):
+            raise InvalidModel("TensorFlow/Keras models are not supported.")
+
+        language_mapping = get_all_languages()
+        language_codes = list(language_mapping.keys())
+
+        model_config = ModelConfig(
+            model_id=model_id,
+            revision=revision,
+            framework=framework,
+            task=model_info.pipeline_tag,
+            languages=[
+                language_mapping[tag]
+                for tag in model_info.tags
+                if tag in language_codes
+            ],
+            model_type=ModelType.HF_HUB_GENERATIVE,
+            model_cache_dir=create_model_cache_dir(
+                cache_dir=benchmark_config.cache_dir, model_id=model_id
+            ),
+            adapter_base_model_id=model_info.adapter_base_model_id,
+        )
+
+        return model_config
+
+    def _load_model_and_tokenizer(self) -> "tuple[LLM, PreTrainedTokenizer]":
+        """Load the model and tokenizer.
+
+        Returns:
+            The loaded model and tokenizer.
+        """
+        hf_model_config = AutoConfig.from_pretrained(
+            self.model_config.model_id,
+            revision=self.model_config.revision,
+            cache_dir=self.model_config.model_cache_dir,
+        )
+
+        quantization = None
+        if hasattr(hf_model_config, "quantization_config"):
+            quantization = hf_model_config.quantization_config.get("quant_method")
+
+        # The quantised models require extra dependencies
+        if quantization == "gptq" and (
+            importlib.util.find_spec("auto_gptq") is None
+            or importlib.util.find_spec("optimum") is None
+        ):
+            raise NeedsExtraInstalled(extra="quantization")
+        if quantization == "awq" and importlib.util.find_spec("awq") is None:
+            raise NeedsExtraInstalled(extra="quantization")
+
+        dtype: str | torch.dtype = "auto"
+        if quantization is not None and hf_model_config.torch_dtype != torch.float16:
+            logger.info(
+                "You are loading a quantized model with dtype "
+                f"{hf_model_config.torch_dtype}, which vLLM does not support. Setting "
+                "dtype to float16 instead."
+            )
+            dtype = torch.float16
+
+        if self.model_config.adapter_base_model_id is not None:
+            download_dir = str(Path(self.model_config.model_cache_dir) / "base_model")
+        else:
+            download_dir = str(self.model_config.model_cache_dir)
+
+        max_model_len = 5_000
+        potential_max_model_length_config_names = [
+            "max_position_embeddings",
+            "max_sequence_length",
+            "model_max_length",
+            "sliding_window",
+            "sliding_window_size",
+            "n_positions",
+        ]
+        for config_name in potential_max_model_length_config_names:
+            if hasattr(hf_model_config, config_name):
+                model_len = getattr(hf_model_config, config_name)
+                if model_len is not None:
+                    max_model_len = min(max_model_len, model_len)
+
+        vllm_kwargs = dict(
+            model=self.model_config.adapter_base_model_id or self.model_config.model_id,
+            tokenizer=self.model_config.adapter_base_model_id
+            or self.model_config.model_id,
+            gpu_memory_utilization=0.95,
+            max_model_len=max_model_len,
+            download_dir=download_dir,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+            revision=self.model_config.revision,
+            seed=4242,
+            distributed_executor_backend="mp",  # "ray",
+            tensor_parallel_size=torch.cuda.device_count(),
+            disable_custom_all_reduce=True,
+            quantization=quantization,
+            dtype=dtype,
+            enforce_eager=True,
+            max_logprobs=MAX_LOGPROBS if self.output_scores else None,
+            # TEMP: Prefix caching isn't supported with sliding window in vLLM yet, so
+            # we disable it for now
+            enable_prefix_caching=False,
+            enable_lora=self.model_config.adapter_base_model_id is not None,
+            max_lora_rank=256,
+        )
+
+        clear_vllm()
+        model = LLM(**vllm_kwargs)
+        model._run_engine = MethodType(_run_engine_with_fixed_progress_bars, model)
+        model.config = hf_model_config
+
+        tokenizer = load_tokenizer(
+            model_id=self.model_config.model_id,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+            model_max_length=max_model_len,
+        )
+
+        return model, tokenizer
+
+    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.supertask:
+            case "sequence-classification":
+                labels = it.cycle(self.dataset_config.task.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 "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 "token-classification":
+                labels = it.cycle(
+                    [
+                        label.lower()
+                        for label in self.dataset_config.task.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["text"] != example["text"]
+                    )
+
+            case "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 supertask: {task.supertask}."
+                )
+
+        random.seed(random_seed)
+        random.shuffle(few_shot_examples)
+        return few_shot_examples
+
+    def _apply_few_shot_prompt(
+        self,
+        examples: dict[str, t.Any],
+        few_shot_examples: list[dict[str, t.Any]],
+        task: Task,
+    ) -> dict[str, t.Any]:
+        """Apply few-shot examples to an example.
+
+        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.
+        """
+        match task.supertask:
+            case "sequence-classification":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=example["text"].replace("\n", " ").strip(),
+                        label=example["label"].replace("\n", " ").strip(),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=text.replace("\n", " ").strip(), label=""
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "text-to-text":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=example["text"].replace("\n", " ").strip(),
+                        target_text=example["target_text"].replace("\n", " ").strip(),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=text.replace("\n", " ").strip(), target_text=""
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "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 = [
+                    self.dataset_config.prompt_template.format(
+                        text=" ".join(example["tokens"]).replace("\n", " ").strip(),
+                        label=create_label(example=example),
+                    )
+                    for example in few_shot_examples
+                ]
+                new_sections = [
+                    self.dataset_config.prompt_template.format(
+                        text=" ".join(tokens).replace("\n", " ").strip(), label=""
+                    )
+                    for tokens in examples["tokens"]
+                ]
+
+            case "question-answering":
+                few_shot_sections = [
+                    self.dataset_config.prompt_template.format(
+                        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 = [
+                    self.dataset_config.prompt_template.format(
+                        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 supertask: {task.supertask}."
+                )
+
+        instruction_model = self._tokenizer.chat_template is not None
+        if instruction_model:
+
+            def split_section(section: str) -> tuple[str, str]:
+                """Split a section of the prompt to user and assistant messages."""
+                user_part = "\n".join(section.split("\n")[:-1])
+                assistant_part = section.split("\n")[-1]
+                return user_part, assistant_part
+
+            few_shot_messages = [
+                dict(role=role, content=content.split(":", 1)[1].strip())
+                for section in few_shot_sections
+                for role, content in zip(
+                    it.cycle(["user", "assistant"]), split_section(section=section)
+                )
+                if content.split(":", 1)[1].strip() != ""
+            ]
+
+            if self.dataset_config.prompt_prefix:
+                few_shot_messages[0]["content"] = (
+                    self.dataset_config.prompt_prefix
+                    + "\n\n"
+                    + few_shot_messages[0]["content"]
+                )
+
+            messages_list = [
+                few_shot_messages
+                + [
+                    dict(
+                        role="user",
+                        content=split_section(section=new_section)[0]
+                        .split(":", 1)[1]
+                        .strip(),
+                    )
+                ]
+                for new_section in new_sections
+            ]
+
+            examples["text"] = [
+                self._tokenizer.apply_chat_template(
+                    conversation=messages, tokenize=False, add_generation_prompt=True
+                )
+                for messages in messages_list
+            ]
+
+        else:
+            prompt_prefix = ""
+            if self.dataset_config.prompt_prefix:
+                prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"
+
+            examples["text"] = [
+                prompt_prefix + "\n\n".join(few_shot_sections) + "\n\n" + new_section
+                for new_section in new_sections
+            ]
+
+        return examples
+
+    def _apply_instruction_prompt(
+        self, examples: dict[str, t.Any], task: Task
+    ) -> dict[str, t.Any]:
+        """Apply instruction prompts to an example.
+
+        Args:
+            examples:
+                The examples to apply the instruction prompts to.
+            task:
+                The task that is being benchmarked.
+
+        Returns:
+            The example with the instruction prompts applied.
+        """
+        match task.supertask:
+            case "sequence-classification" | "text-to-text" | "token-classification":
+                prompts = [
+                    self.dataset_config.instruction_prompt.format(
+                        text=re.sub(r"\n+", "\n", text).strip()
+                    )
+                    for text in examples["text"]
+                ]
+
+            case "question-answering":
+                prompts = [
+                    self.dataset_config.instruction_prompt.format(
+                        text=re.sub(pattern=r"\n+", repl="\n", string=text).strip(),
+                        question=question.strip(),
+                    )
+                    for (text, question) in zip(
+                        examples["context"], examples["question"]
+                    )
+                ]
+
+            case _:
+                raise NotImplementedError(
+                    f"Unsupported task supertask: {task.supertask}."
+                )
+
+        examples["text"] = prompts
+        return examples
+
+    def __del__(self) -> None:
+        """Clear the GPU memory used by the model, and remove the model itself."""
+        if hasattr(self, "_model"):
+            del self._model
+        del self
+        clear_vllm()
+
+    @cached_property
+    def data_collator(self) -> c.Callable[[list[t.Any]], dict[str, t.Any]]:
+        """The data collator used to prepare samples during finetuning.
+
+        Returns:
+            The data collator.
+        """
+        raise NotImplementedError(
+            "The `data_collator` property has not been implemented for vLLM models."
+        )
+
+    @cached_property
+    def trainer_class(self) -> t.Type["Trainer"]:
+        """The Trainer class to use for finetuning.
+
+        Returns:
+            The Trainer class.
+        """
+        raise NotImplementedError(
+            "The `trainer_class` property has not been implemented for vLLM models."
+        )
+
+
+def load_tokenizer(
+    model_id: str, trust_remote_code: bool, model_max_length: int
+) -> "PreTrainedTokenizer":
+    """Load the tokenizer.
+
+    Args:
+        model_id:
+            The model identifier. Used for logging.
+        trust_remote_code:
+            Whether to trust remote code.
+        model_max_length:
+            The maximum length of the model.
+
+    Returns:
+        The loaded tokenizer.
+    """
+    while True:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_id,
+                use_fast=True,
+                verbose=False,
+                trust_remote_code=trust_remote_code,
+                padding_side="left",
+                truncation_side="left",
+                model_max_length=model_max_length,
+            )
+            if tokenizer.pad_token_id is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            return tokenizer
+        except (json.JSONDecodeError, OSError, TypeError):
+            raise InvalidModel(f"Could not load tokenizer for model {model_id!r}.")
+        except (TimeoutError, RequestError):
+            logger.info(f"Couldn't load tokenizer for {model_id!r}. Retrying.")
+            sleep(5)
+            continue
+
+
+def _run_engine_with_fixed_progress_bars(
+    self: "LLM", use_tqdm: bool
+) -> list["RequestOutput"]:
+    if use_tqdm:
+        num_requests = self.llm_engine.get_num_unfinished_requests()
+        pbar = tqdm(
+            total=num_requests, leave=False, disable=hasattr(sys, "_called_from_test")
+        )
+    else:
+        pbar = None
+
+    # Run the engine.
+    outputs: list["RequestOutput"] = list()
+    while self.llm_engine.has_unfinished_requests():
+        step_outputs = self.llm_engine.step()
+        for output in step_outputs:
+            if output.finished:
+                outputs.append(output)
+                if pbar is not None:
+                    pbar.update(1)
+
+    if pbar is not None:
+        pbar.close()
+
+    # Sort the outputs by request ID. This is necessary because some requests may be
+    # finished earlier than its previous requests.
+    outputs = sorted(outputs, key=lambda x: int(x.request_id))
+
+    return outputs
+
+
+def clear_vllm() -> None:
+    """Clear the GPU memory used by the vLLM model, enabling re-initialisation."""
+    try:
+        destroy_model_parallel()
+    except ImportError:
+        pass
+    clear_memory()
+    if ray.is_initialized():
+        ray.shutdown()
diff --git a/src/scandeval/benchmarker.py b/src/scandeval/benchmarker.py
index 819b192a8..5b45fdbb0 100644
--- a/src/scandeval/benchmarker.py
+++ b/src/scandeval/benchmarker.py
@@ -1,129 +1,37 @@
 """Class that benchmarks Scandinavian language models."""
 
-import importlib.metadata
 import json
 import logging
 import re
 import sys
+import typing as t
 from copy import deepcopy
 from pathlib import Path
 from shutil import rmtree
 from time import sleep
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel, ConfigDict
 
 from .benchmark_config_factory import build_benchmark_config
-from .config import BenchmarkConfig
+from .constants import GENERATIVE_MODEL_TASKS
+from .data_loading import load_data
+from .data_models import BenchmarkConfigParams, BenchmarkResult
 from .dataset_configs import get_all_dataset_configs
-from .dataset_factory import DatasetFactory
 from .enums import Device, Framework
 from .exceptions import InvalidBenchmark, InvalidModel
-from .types import ScoreDict
-from .utils import get_huggingface_model_lists
-
-if TYPE_CHECKING:
-    from .config import DatasetConfig, Language
-    from .protocols import GenerativeModel, Tokenizer
-
-
-logger = logging.getLogger(__package__)
-
-
-class BenchmarkConfigParams(BaseModel):
-    """The parameters for the benchmark configuration."""
-
-    model_config = ConfigDict(protected_namespaces=())
-
-    progress_bar: bool
-    save_results: bool
-    task: str | list[str] | None
-    dataset: str | list[str] | None
-    language: str | list[str]
-    model_language: str | list[str] | None
-    dataset_language: str | list[str] | None
-    framework: Framework | str | None
-    device: Device | None
-    batch_size: int
-    evaluate_train: bool
-    raise_errors: bool
-    cache_dir: str
-    token: bool | str
-    openai_api_key: str | None
-    prefer_azure: bool
-    azure_openai_api_key: str | None
-    azure_openai_endpoint: str | None
-    azure_openai_api_version: str | None
-    force: bool
-    verbose: bool
-    trust_remote_code: bool
-    load_in_4bit: bool | None
-    use_flash_attention: bool | None
-    clear_model_cache: bool
-    only_validation_split: bool
-    few_shot: bool
-    num_iterations: int
-    debug: bool
-    run_with_cli: bool
-
-
-class BenchmarkResult(BaseModel):
-    """A benchmark result."""
-
-    dataset: str
-    task: str
-    dataset_languages: list[str]
-    model: str
-    results: ScoreDict
-    num_model_parameters: int
-    max_sequence_length: int
-    vocabulary_size: int
-    generative: bool
-    few_shot: bool
-    validation_split: bool
-    scandeval_version: str = importlib.metadata.version(__package__)
-
-    @classmethod
-    def from_dict(cls, config: dict) -> "BenchmarkResult":
-        """Create a benchmark result from a dictionary.
-
-        Args:
-            config:
-                The configuration dictionary.
-
-        Returns:
-            The benchmark result.
-        """
-        # To be backwards compatible, we accept old results which changed the model
-        # name with parameters rather than adding them as explicit parameters
-        val_matches = re.search(r"\(.*val.*\)$", config["model"])
-        few_shot_matches = re.search(r"\(.*few-shot.*\)$", config["model"])
-        config["model"] = re.sub(
-            r"\(.*(few-shot|val).*\)$", "", config["model"]
-        ).strip()
-
-        # The default value for `few_shot` is True. It won't do anything if the model
-        # is not generative, so this is fine
-        if "generative" not in config:
-            config["generative"] = few_shot_matches is not None
-        if "few_shot" not in config:
-            config["few_shot"] = True
+from .finetuning import finetune
+from .generation import generate
+from .model_config import get_model_config
+from .model_loading import load_model
+from .scores import log_scores
+from .speed_benchmark import benchmark_speed
+from .tasks import SPEED
+from .utils import enforce_reproducibility
 
-        if "validation_split" not in config:
-            config["validation_split"] = val_matches is not None
+if t.TYPE_CHECKING:
+    from .benchmark_modules import BenchmarkModule
+    from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
-        return cls(**config)
 
-    def append_to_results(self, results_path: Path) -> None:
-        """Append the benchmark result to the results file.
-
-        Args:
-            results_path:
-                The path to the results file.
-        """
-        json_str = json.dumps(self.model_dump())
-        with results_path.open("a") as f:
-            f.write("\n" + json_str)
+logger = logging.getLogger("scandeval")
 
 
 class Benchmarker:
@@ -137,8 +45,6 @@ class Benchmarker:
         force:
             Whether to force evaluations of models, even if they have been benchmarked
             already.
-        dataset_factory:
-            The factory for creating datasets.
         results_path:
             The path to the results file.
         benchmark_results:
@@ -157,22 +63,16 @@ def __init__(
         framework: Framework | str | None = None,
         device: Device | None = None,
         batch_size: int = 32,
-        evaluate_train: bool = False,
         raise_errors: bool = False,
         cache_dir: str = ".scandeval_cache",
-        token: bool | str = True,
-        openai_api_key: str | None = None,
-        prefer_azure: bool = False,
-        azure_openai_api_key: str | None = None,
-        azure_openai_endpoint: str | None = None,
-        azure_openai_api_version: str | None = None,
+        api_key: str | None = None,
         force: bool = False,
         verbose: bool = False,
         trust_remote_code: bool = False,
         load_in_4bit: bool | None = None,
         use_flash_attention: bool | None = None,
         clear_model_cache: bool = False,
-        only_validation_split: bool = False,
+        evaluate_test_split: bool = False,
         few_shot: bool = True,
         num_iterations: int = 10,
         debug: bool = False,
@@ -213,43 +113,19 @@ def __init__(
                 The device to use for benchmarking. Defaults to None.
             batch_size:
                 The batch size to use. Defaults to 32.
-            evaluate_train:
-                Whether to evaluate the training set as well. Defaults to False.
             raise_errors:
                 Whether to raise errors instead of skipping the model evaluation.
                 Defaults to False.
             cache_dir:
                 Directory to store cached models. Defaults to '.scandeval_cache'.
-            token:
-                The authentication token for the Hugging Face Hub. If a boolean value
-                is specified then the token will be fetched from the Hugging Face CLI,
-                where the user has logged in through `huggingface-cli login`. If a
-                string is specified then it will be used as the token. Defaults to
-                True.
-            openai_api_key:
-                The OpenAI API key to use for authentication. If None, then this will
-                be loaded from the environment variable `OPENAI_API_KEY`. Defaults to
-                None.
-            prefer_azure:
-                In the case where both OpenAI and Azure OpenAI models are available,
-                whether to use the Azure OpenAI models. Defaults to False.
-            azure_openai_api_key:
-                The Azure OpenAI API key to use for authentication. If None, then this
-                will be loaded from the environment variable `AZURE_OPENAI_API_KEY`.
-                Defaults to None.
-            azure_openai_endpoint:
-                The Azure OpenAI endpoint to use for authentication. If None, then this
-                will be loaded from the environment variable `AZURE_OPENAI_ENDPOINT`.
-                Defaults to None.
-            azure_openai_api_version:
-                The Azure OpenAI API version to use for authentication. If None, then this
-                will be loaded from the environment variable `AZURE_OPENAI_API_VERSION`.
-                Defaults to None.
+            api_key:
+                The API key to use for a given inference API.
             force:
                 Whether to force evaluations of models, even if they have been
                 benchmarked already. Defaults to False.
             verbose:
-                Whether to output additional output. Defaults to False.
+                Whether to output additional output. This is automatically set if
+                `debug` is True. Defaults to False.
             trust_remote_code:
                 Whether to trust remote code when loading models. Defaults to False.
             load_in_4bit:
@@ -262,9 +138,8 @@ def __init__(
             clear_model_cache:
                 Whether to clear the model cache after benchmarking each model.
                 Defaults to False.
-            only_validation_split:
-                Whether to only evaluate the validation split of the datasets. Defaults
-                to False.
+            evaluate_test_split:
+                Whether to evaluate the test split of the datasets. Defaults to False.
             few_shot:
                 Whether to only evaluate the model using few-shot evaluation. Only
                 relevant if the model is generative. Defaults to True.
@@ -296,22 +171,16 @@ def __init__(
             framework=framework,
             device=device,
             batch_size=batch_size,
-            evaluate_train=evaluate_train,
             raise_errors=raise_errors,
             cache_dir=cache_dir,
-            token=token,
-            openai_api_key=openai_api_key,
-            prefer_azure=prefer_azure,
-            azure_openai_api_key=azure_openai_api_key,
-            azure_openai_endpoint=azure_openai_endpoint,
-            azure_openai_api_version=azure_openai_api_version,
+            api_key=api_key,
             force=force,
             verbose=verbose,
             trust_remote_code=trust_remote_code,
             load_in_4bit=load_in_4bit,
             use_flash_attention=use_flash_attention,
             clear_model_cache=clear_model_cache,
-            only_validation_split=only_validation_split,
+            evaluate_test_split=evaluate_test_split,
             few_shot=few_shot,
             num_iterations=num_iterations,
             debug=debug,
@@ -325,9 +194,7 @@ def __init__(
         # Initialise variable storing model lists, so we only have to fetch it once
         self._model_lists: dict[str, list[str]] | None = None
 
-        # Set up the results path
         self.results_path = Path.cwd() / "scandeval_benchmark_results.jsonl"
-
         adjust_logging_level(verbose=self.benchmark_config.verbose)
 
     @property
@@ -345,7 +212,7 @@ def benchmark_results(self) -> list[BenchmarkResult]:
 
     def benchmark(
         self,
-        model: list[str] | str | None = None,
+        model: list[str] | str,
         task: str | list[str] | None = None,
         dataset: list[str] | str | None = None,
         progress_bar: bool | None = None,
@@ -356,21 +223,16 @@ def benchmark(
         framework: Framework | str | None = None,
         device: Device | None = None,
         batch_size: int | None = None,
-        evaluate_train: bool | None = None,
         raise_errors: bool | None = None,
         cache_dir: str | None = None,
-        token: bool | str | None = None,
-        openai_api_key: str | None = None,
-        azure_openai_api_key: str | None = None,
-        azure_openai_endpoint: str | None = None,
-        azure_openai_api_version: str | None = None,
+        api_key: str | None = None,
         force: bool | None = None,
         verbose: bool | None = None,
         trust_remote_code: bool | None = None,
         load_in_4bit: bool | None = None,
         use_flash_attention: bool | None = None,
         clear_model_cache: bool | None = None,
-        only_validation_split: bool | None = None,
+        evaluate_test_split: bool | None = None,
         few_shot: bool | None = None,
         num_iterations: int | None = None,
     ) -> list[BenchmarkResult]:
@@ -381,8 +243,7 @@ def benchmark(
                 The full Hugging Face Hub path(s) to the pretrained transformer model.
                 The specific model version to use can be added after the suffix '@':
                 "model@v1.0.0". It can be a branch name, a tag name, or a commit id,
-                and defaults to the latest version if not specified. If None then all
-                relevant model IDs will be benchmarked. Defaults to None.
+                and defaults to the latest version if not specified.
             task:
                 The tasks benchmark the model(s) on. Mutually exclusive with `dataset`.
                 If both `task` and `dataset` are None then all datasets will be
@@ -422,36 +283,14 @@ def benchmark(
             batch_size:
                 The batch size to use. Defaults to the value specified when initialising
                 the benchmarker.
-            evaluate_train:
-                Whether to evaluate the training set as well. Defaults to the value
-                specified when initialising the benchmarker.
             raise_errors:
                 Whether to raise errors instead of skipping the model evaluation.
             cache_dir:
                 Directory to store cached models. Defaults to the value specified when
                 initialising the benchmarker.
-            token:
-                The authentication token for the Hugging Face Hub. If a boolean value is
-                specified then the token will be fetched from the Hugging Face CLI, where
-                the user has logged in through `huggingface-cli login`. If a string is
-                specified then it will be used as the token. Defaults to the value
+            api_key:
+                The API key to use for a given inference server. Defaults to the value
                 specified when initialising the benchmarker.
-            openai_api_key:
-                The OpenAI API key to use for authentication. If None, then this will be
-                loaded from the environment variable `OPENAI_API_KEY`. Defaults to the
-                value specified when initialising the benchmarker.
-            azure_openai_api_key:
-                The Azure OpenAI API key to use for authentication. If None, then this
-                will be loaded from the environment variable `AZURE_OPENAI_API_KEY`.
-                Defaults to the value specified when initialising the benchmarker.
-            azure_openai_endpoint:
-                The Azure OpenAI endpoint to use for authentication. If None, then this
-                will be loaded from the environment variable `AZURE_OPENAI_ENDPOINT`.
-                Defaults to the value specified when initialising the benchmarker.
-            azure_openai_api_version:
-                The api version for the Azure OpenAI API, e.g. "2023-12-01-preview". If
-                None then the environment varaible `AZURE_OPENAI_API_VERSION` will be used.
-                Defaults to the value specified when initialising the benchmarker.
             force:
                 Whether to force evaluations of models, even if they have been
                 benchmarked already. Defaults to the value specified when initialising
@@ -472,9 +311,9 @@ def benchmark(
             clear_model_cache:
                 Whether to clear the model cache after benchmarking each model. Defaults
                 to the value specified when initialising the benchmarker.
-            only_validation_split:
-                Whether to only evaluate the validation split of the datasets. Defaults
-                to the value specified when initialising the benchmarker.
+            evaluate_test_split:
+                Whether to evaluate the test split of the datasets. Defaults to the
+                value specified when initialising the benchmarker.
             few_shot:
                 Whether to only evaluate the model using few-shot evaluation. Only
                 relevant if the model is generative. Defaults to the value specified
@@ -495,86 +334,50 @@ def benchmark(
         if task is not None and dataset is not None:
             raise ValueError("Only one of `task` and `dataset` can be specified.")
 
-        benchmark_config_params = deepcopy(self.benchmark_config_default_params)
-        if task is not None:
-            benchmark_config_params.task = task
-            benchmark_config_params.dataset = None
-        if dataset is not None:
-            benchmark_config_params.dataset = dataset
-            benchmark_config_params.task = None
-        if progress_bar is not None:
-            benchmark_config_params.progress_bar = progress_bar
-        if save_results is not None:
-            benchmark_config_params.save_results = save_results
-        if language is not None:
-            benchmark_config_params.language = language
-        if model_language is not None:
-            benchmark_config_params.model_language = model_language
-        if dataset_language is not None:
-            benchmark_config_params.dataset_language = dataset_language
-        if framework is not None:
-            benchmark_config_params.framework = framework
-        if device is not None:
-            benchmark_config_params.device = device
-        if batch_size is not None:
-            benchmark_config_params.batch_size = batch_size
-        if evaluate_train is not None:
-            benchmark_config_params.evaluate_train = evaluate_train
-        if raise_errors is not None:
-            benchmark_config_params.raise_errors = raise_errors
-        if cache_dir is not None:
-            benchmark_config_params.cache_dir = cache_dir
-        if token is not None:
-            benchmark_config_params.token = token
-        if openai_api_key is not None:
-            benchmark_config_params.openai_api_key = openai_api_key
-        if azure_openai_api_key is not None:
-            benchmark_config_params.azure_openai_api_key = azure_openai_api_key
-        if azure_openai_endpoint is not None:
-            benchmark_config_params.azure_openai_endpoint = azure_openai_endpoint
-        if azure_openai_api_version is not None:
-            benchmark_config_params.azure_openai_api_version = azure_openai_api_version
-        if force is not None:
-            benchmark_config_params.force = force
-        if verbose is not None:
-            benchmark_config_params.verbose = verbose
-        if trust_remote_code is not None:
-            benchmark_config_params.trust_remote_code = trust_remote_code
-        if load_in_4bit is not None:
-            benchmark_config_params.load_in_4bit = load_in_4bit
-        if use_flash_attention is not None:
-            benchmark_config_params.use_flash_attention = use_flash_attention
-        if clear_model_cache is not None:
-            benchmark_config_params.clear_model_cache = clear_model_cache
-        if only_validation_split is not None:
-            benchmark_config_params.only_validation_split = only_validation_split
-        if few_shot is not None:
-            benchmark_config_params.few_shot = few_shot
-        if num_iterations is not None:
-            benchmark_config_params.num_iterations = num_iterations
-
-        benchmark_config = build_benchmark_config(
-            **benchmark_config_params.model_dump()
+        benchmark_config = self._get_updated_benchmark_config(
+            task=task,
+            dataset=dataset,
+            progress_bar=progress_bar,
+            save_results=save_results,
+            language=language,
+            model_language=model_language,
+            dataset_language=dataset_language,
+            framework=framework,
+            device=device,
+            batch_size=batch_size,
+            raise_errors=raise_errors,
+            cache_dir=cache_dir,
+            api_key=api_key,
+            force=force,
+            verbose=verbose,
+            trust_remote_code=trust_remote_code,
+            load_in_4bit=load_in_4bit,
+            use_flash_attention=use_flash_attention,
+            clear_model_cache=clear_model_cache,
+            evaluate_test_split=evaluate_test_split,
+            few_shot=few_shot,
+            num_iterations=num_iterations,
         )
+
         adjust_logging_level(verbose=benchmark_config.verbose)
 
         if benchmark_config.clear_model_cache:
             clear_model_cache_fn(cache_dir=benchmark_config.cache_dir)
 
-        model_ids = self._prepare_model_ids(
-            model=model,
-            model_languages=benchmark_config.model_languages,
-            token=benchmark_config.token,
-        )
+        model_ids = self._prepare_model_ids(model_id=model)
         dataset_configs = prepare_dataset_configs(
             dataset_names=benchmark_config.datasets
         )
 
         current_benchmark_results: list[BenchmarkResult] = list()
         for m_id in model_ids:
-            m_id = m_id.rstrip(" /")
-            loaded_model = None
-            loaded_tokenizer = None
+            try:
+                model_config = get_model_config(
+                    model_id=m_id, benchmark_config=benchmark_config
+                )
+            except InvalidModel as e:
+                logger.info(e.message)
+                continue
 
             for dataset_config in dataset_configs:
                 # Skip if we have already benchmarked this model on this dataset and
@@ -583,7 +386,7 @@ def benchmark(
                     model_id=m_id,
                     dataset=dataset_config.name,
                     few_shot=benchmark_config.few_shot,
-                    validation_split=benchmark_config.only_validation_split,
+                    validation_split=not benchmark_config.evaluate_test_split,
                     benchmark_results=self.benchmark_results,
                 ):
                     logger.debug(
@@ -593,78 +396,72 @@ def benchmark(
                     continue
 
                 # Benchmark a single model on a single dataset
-                try:
-                    benchmark_output = self._benchmark_single(
-                        dataset_config=dataset_config,
-                        model_id=m_id,
-                        raise_errors=benchmark_config.raise_errors,
-                        model=loaded_model,
-                        tokenizer=loaded_tokenizer,
-                        benchmark_config=benchmark_config,
-                    )
-                except InvalidModel as e:
-                    if benchmark_config.raise_errors:
-                        raise e
-                    logger.info(e)
-                    break
+                benchmark_output_or_err = self._benchmark_single(
+                    model_config=model_config,
+                    dataset_config=dataset_config,
+                    benchmark_config=benchmark_config,
+                )
+
+                if (
+                    isinstance(benchmark_output_or_err, Exception)
+                    and benchmark_config.raise_errors
+                ):
+                    raise benchmark_output_or_err
 
-                # If the benchmark was unsuccessful then skip
-                if isinstance(benchmark_output, dict) and "error" in benchmark_output:
-                    error_msg = benchmark_output["error"]
+                elif isinstance(benchmark_output_or_err, InvalidBenchmark):
                     logger.info(
                         f"{m_id} could not be benchmarked on "
                         f"{dataset_config.pretty_name}. Skipping. The error message "
-                        f"raised was {error_msg!r}."
+                        f"raised was {benchmark_output_or_err.message!r}."
                     )
                     continue
 
-                assert isinstance(benchmark_output, tuple)
-                record, loaded_model, loaded_tokenizer = benchmark_output
+                elif isinstance(benchmark_output_or_err, InvalidModel):
+                    logger.info(benchmark_output_or_err.message)
+                    break
 
-                # Save the benchmark results
-                assert isinstance(record, BenchmarkResult)
-                current_benchmark_results.append(record)
-                if benchmark_config.save_results:
-                    record.append_to_results(results_path=self.results_path)
+                else:
+                    record = benchmark_output_or_err
+                    current_benchmark_results.append(record)
+                    if benchmark_config.save_results:
+                        record.append_to_results(results_path=self.results_path)
 
             if benchmark_config.clear_model_cache:
                 clear_model_cache_fn(cache_dir=benchmark_config.cache_dir)
 
         return current_benchmark_results
 
-    def _prepare_model_ids(
-        self,
-        model: list[str] | str | None,
-        model_languages: list["Language"],
-        token: bool | str | None,
-    ) -> list[str]:
-        """Prepare the model ID(s) to be benchmarked.
+    def _get_updated_benchmark_config(self, **kwargs) -> "BenchmarkConfig":
+        """Get an updated benchmark configuration.
 
         Args:
-            model:
-                The model ID(s) of the models to benchmark. If None then all model IDs
-                will be retrieved.
-            model_languages:
-                The languages of the models to fetch.
-            token:
-                The authentication token for the Hugging Face Hub.
+            **kwargs:
+                The new parameters for the benchmark configuration.
 
         Returns:
-            The prepared list of model IDs.
+            The updated benchmark configuration.
         """
-        model_ids: list[str]
-
-        # If `model_id` is not specified, then fetch all the relevant model IDs
-        if model is None:
-            model_ids = self._get_model_ids(languages=model_languages, token=token)
+        benchmark_config_params = deepcopy(self.benchmark_config_default_params)
+        for key, value in kwargs.items():
+            if value is not None and hasattr(benchmark_config_params, key):
+                setattr(benchmark_config_params, key, value)
+                if key == "task":
+                    benchmark_config_params.dataset = None
+                elif key == "dataset":
+                    benchmark_config_params.task = None
+        return build_benchmark_config(**benchmark_config_params.model_dump())
+
+    def _prepare_model_ids(self, model_id: list[str] | str) -> list[str]:
+        """Prepare the model ID(s) to be benchmarked.
 
-        # Otherwise, if `model_id` is a string, ensure that it is a list
-        elif isinstance(model, str):
-            model_ids = [model]
+        Args:
+            model_id:
+                The model ID(s) of the models to benchmark.
 
-        # Otherwise `model_id` is already a list, so we do nothing
-        else:
-            model_ids = model
+        Returns:
+            The prepared list of model IDs.
+        """
+        model_ids = [model_id] if isinstance(model_id, str) else model_id
 
         # Reorder the `model_ids` list to include the ones present in the benchmark
         # results first
@@ -677,40 +474,31 @@ def _prepare_model_ids(
             m_id for m_id in model_ids if m_id not in benchmarked_model_ids
         ]
 
-        return model_ids_sorted
+        return [m_id.rstrip(" /") for m_id in model_ids_sorted]
 
     def _benchmark_single(
         self,
+        model_config: "ModelConfig",
         dataset_config: "DatasetConfig",
-        model_id: str,
-        raise_errors: bool,
-        model: "GenerativeModel | None",
-        tokenizer: "Tokenizer | None",
-        benchmark_config: BenchmarkConfig,
-    ) -> (
-        tuple[BenchmarkResult, "GenerativeModel | None", "Tokenizer | None"]
-        | dict[str, str]
-    ):
+        benchmark_config: "BenchmarkConfig",
+    ) -> BenchmarkResult | InvalidBenchmark | InvalidModel:
         """Benchmark a single model on a single dataset.
 
         Args:
+            model_config:
+                The configuration of the model we are evaluating.
             dataset_config:
-                The dataset configuration to use.
-            model_id:
-                The model ID to use.
-            raise_errors:
-                Whether to raise errors instead of skipping the model evaluation.
-            model:
-                The pre-loaded model, if available.
-            tokenizer:
-                The pre-loaded tokenizer, if available.
+                The configuration of the dataset we are evaluating on.
             benchmark_config:
-                The benchmark configuration.
+                The general benchmark configuration.
 
         Returns:
-            The benchmark result, or a dictionary containing an error message.
+            The benchmark result, or an error if the benchmark was unsuccessful.
         """
-        logger.info(f"Benchmarking {model_id} on {dataset_config.pretty_name}")
+        # Initial logging
+        logger.info(
+            f"Benchmarking {model_config.model_id} on {dataset_config.pretty_name}"
+        )
         if dataset_config.unofficial:
             logger.info(
                 f"Note that the {dataset_config.name!r} dataset is unofficial, "
@@ -724,35 +512,101 @@ def _benchmark_single(
                 "batch. For this reason, evaluation will be slower."
             )
 
+        model: BenchmarkModule | None = None
         while True:
             try:
-                dataset_factory = DatasetFactory(benchmark_config=benchmark_config)
-                dataset = dataset_factory.build_dataset(dataset_config)
-                results, metadata_dict, model, tokenizer = dataset(
-                    model_id=model_id, model=model, tokenizer=tokenizer
+                model_config = get_model_config(
+                    model_id=model_config.model_id,
+                    benchmark_config=self.benchmark_config,
                 )
+
+                # Set random seeds to enforce reproducibility of the randomly
+                # initialised weights
+                rng = enforce_reproducibility(framework=model_config.framework)
+
+                if model is None or not model.is_generative:
+                    logger.info("Loading model...")
+                    model = load_model(
+                        model_config=model_config,
+                        dataset_config=dataset_config,
+                        benchmark_config=benchmark_config,
+                    )
+                assert model is not None
+
+                # This happens when a local model is used, as we cannot fetch the model
+                # metadata. Note that this is only the case if the model type is not any
+                # of the ones hardcoded in `local.py`
+                if model_config.task == "unknown":
+                    if model.is_generative:
+                        model_config.task = GENERATIVE_MODEL_TASKS[0]
+                    else:
+                        model_config.task = "fill-mask"
+
+                if dataset_config.task == SPEED:
+                    scores = benchmark_speed(
+                        model=model, benchmark_config=self.benchmark_config
+                    )
+
+                else:
+                    bootstrapped_datasets = load_data(
+                        rng=rng,
+                        dataset_config=dataset_config,
+                        benchmark_config=benchmark_config,
+                    )
+                    prepared_datasets = model.prepare_datasets(
+                        datasets=bootstrapped_datasets, task=dataset_config.task
+                    )
+                    if model.is_generative:
+                        scores = generate(
+                            model=model,
+                            datasets=prepared_datasets,
+                            model_config=model_config,
+                            dataset_config=dataset_config,
+                            benchmark_config=self.benchmark_config,
+                        )
+                    else:
+                        scores = finetune(
+                            model=model,
+                            datasets=prepared_datasets,
+                            model_config=model_config,
+                            dataset_config=dataset_config,
+                            benchmark_config=benchmark_config,
+                        )
+
+                results = log_scores(
+                    dataset_name=dataset_config.pretty_name,
+                    metric_configs=dataset_config.task.metrics,
+                    scores=scores,
+                    model_id=model_config.model_id,
+                )
+
                 record = BenchmarkResult(
                     dataset=dataset_config.name,
                     task=dataset_config.task.name,
                     dataset_languages=[
                         language.code for language in dataset_config.languages
                     ],
-                    model=model_id,
+                    model=model_config.model_id,
                     results=results,
-                    **metadata_dict,
+                    num_model_parameters=model.num_params,
+                    max_sequence_length=model.model_max_length,
+                    vocabulary_size=model.vocab_size,
+                    generative=model.is_generative,
+                    few_shot=benchmark_config.few_shot,
+                    validation_split=not benchmark_config.evaluate_test_split,
                 )
                 logger.debug(f"Results:\n{results}")
-                return record, model, tokenizer
+                return record
 
-            except InvalidBenchmark as e:
-                # If the model ID is not valid then raise an error, if specified
+            except (InvalidBenchmark, InvalidModel) as e:
+                # If the model ID is not valid then raise an error
                 model_err_msg = "does not exist on the Hugging Face Hub"
-                if raise_errors and model_err_msg in str(e):
+                if benchmark_config.raise_errors and model_err_msg in str(e):
                     raise e
 
                 # Otherwise, if the error is due to Hugging Face Hub being down, then
                 # wait a bit and try again
-                if "The Hugging Face Hub seems to be down." in str(e):
+                elif "The Hugging Face Hub seems to be down." in str(e):
                     wait_time = 30
                     logger.debug(
                         "The Hugging Face Hub seems to be down. Retrying in "
@@ -770,57 +624,14 @@ def _benchmark_single(
                         "environment variable to `1` and try again."
                     )
 
-                # Otherwise, raise the error or return the error message
-                else:
-                    if raise_errors:
-                        raise e
-                    return dict(error=str(e))
+                elif benchmark_config.raise_errors:
+                    raise e
+                return e
 
     def __call__(self, *args, **kwargs) -> list[BenchmarkResult]:
         """Call the benchmarker. See `Benchmarker.benchmark`."""
         return self.benchmark(*args, **kwargs)
 
-    def _get_model_ids(
-        self, languages: list["Language"], token: bool | str | None
-    ) -> list[str]:
-        """Get list of model IDs from the Hugging Face Hub.
-
-        Args:
-            languages:
-                The languages of the models to fetch.
-            token:
-                The authentication token for the Hugging Face Hub.
-
-        Returns:
-            List of model IDs.
-        """
-        # Specify boolean variables determining whether the input variables are new
-        new_languages = self._model_lists is not None and any(
-            lang.code not in self._model_lists for lang in languages
-        )
-
-        # If the model lists have not been fetched already, then do it
-        if self._model_lists is None or new_languages:
-            self._model_lists = get_huggingface_model_lists(
-                languages=languages, token=token
-            )
-
-        # Extract all the model IDs from the model lists, for the chosen languages
-        model_ids: list[str] = list()
-        for language in languages:
-            model_ids.extend(self._model_lists[language.code])
-
-        # Add the multilingual models
-        model_ids.extend(self._model_lists["multilingual"])
-
-        # Add the fresh models
-        model_ids.extend(self._model_lists["fresh"])
-
-        # Remove duplicate model IDs
-        model_ids = list(set(model_ids))
-
-        return model_ids
-
 
 def model_has_been_benchmarked(
     model_id: str,
@@ -855,20 +666,26 @@ def model_has_been_benchmarked(
     return False
 
 
-def adjust_logging_level(verbose: bool) -> None:
+def adjust_logging_level(verbose: bool, ignore_testing: bool = False) -> int:
     """Adjust the logging level based on verbosity.
 
     Args:
         verbose:
             Whether to output additional output.
+        ignore_testing:
+            Whether to ignore the testing flag.
+
+    Returns:
+        The logging level that was set.
     """
-    if hasattr(sys, "_called_from_test"):
+    if hasattr(sys, "_called_from_test") and not ignore_testing:
         logging_level = logging.CRITICAL
     elif verbose:
         logging_level = logging.DEBUG
     else:
         logging_level = logging.INFO
     logger.setLevel(logging_level)
+    return logging_level
 
 
 def clear_model_cache_fn(cache_dir: str) -> None:
diff --git a/src/scandeval/cli.py b/src/scandeval/cli.py
index 020d0d23d..dd322d304 100644
--- a/src/scandeval/cli.py
+++ b/src/scandeval/cli.py
@@ -13,14 +13,9 @@
 @click.option(
     "--model",
     "-m",
-    default=None,
-    show_default=True,
+    required=True,
     multiple=True,
-    help="""The Hugging Face model ID of the model(s) to be benchmarked. If not
-    specified then all models will be benchmarked, filtered by `model_language` or
-    `language`. The specific model revision to use can be added after the suffix "@":
-    "<model>@v1.0.0". It can be a branch name, a tag name, or a commit id, and
-    defaults to the latest version if not specified.""",
+    help="The ID of the model to benchmark.",
 )
 @click.option(
     "--task",
@@ -80,12 +75,6 @@
     type=click.Choice(["1", "2", "4", "8", "16", "32"]),
     help="The batch size to use.",
 )
-@click.option(
-    "--evaluate-train/--no-evaluate-train",
-    default=False,
-    show_default=True,
-    help="Whether to evaluate the model on the training set.",
-)
 @click.option(
     "--progress-bar/--no-progress-bar",
     default=True,
@@ -120,59 +109,14 @@
     help="The directory where models are datasets are cached.",
 )
 @click.option(
-    "--token",
-    type=str,
-    default="",
-    show_default=True,
-    help="""The authentication token for the Hugging Face Hub. If specified then the
-    `--token` flag will be set to True.""",
-)
-@click.option(
-    "--use-token/--no-use-token",
-    default=True,
-    show_default=True,
-    help="""Whether an authentication token should be used, enabling evaluation of
-    private models. Requires that you are logged in via the `huggingface-cli login`
-    command.""",
-)
-@click.option(
-    "--openai-api-key",
-    type=str,
-    default=None,
-    show_default=True,
-    help="""The API key for the OpenAI API. If not specified then the environment
-    variable `OPENAI_API_KEY` will be used.""",
-)
-@click.option(
-    "--prefer-azure/--no-prefer-azure",
-    default=False,
-    show_default=True,
-    help="""Whether to prefer the Azure OpenAI API over the OpenAI API, if both are
-    available.""",
-)
-@click.option(
-    "--azure-openai-api-key",
-    type=str,
-    default=None,
-    show_default=True,
-    help="""The API key for the Azure OpenAI API. If not specified then the environment
-    variable `AZURE_OPENAI_API_KEY` will be used.""",
-)
-@click.option(
-    "--azure-openai-endpoint",
-    type=str,
-    default=None,
-    show_default=True,
-    help="""The endpoint for the Azure OpenAI API. If not specified then the environment
-    variable `AZURE_OPENAI_ENDPOINT` will be used.""",
-)
-@click.option(
-    "--azure-openai-api-version",
+    "--api-key",
     type=str,
     default=None,
     show_default=True,
-    help="""The api version for the Azure OpenAI API, e.g. "2023-12-01-preview". If
-            None then the environment varaible `AZURE_OPENAI_API_VERSION` will be used.""",
+    help="""The API key to use for a given inference API. If you are benchmarking an "
+    "OpenAI model then this would be the OpenAI API key, if you are benchmarking a "
+    "model on the Hugging Face inference API then this would be the Hugging Face API "
+    "key, and so on.""",
 )
 @click.option(
     "--force/--no-force",
@@ -230,11 +174,11 @@
     to avoid running out of disk space.""",
 )
 @click.option(
-    "--only-validation-split/--no-only-validation-split",
+    "--evaluate-test-split/--evaluate-val-split",
     default=False,
     show_default=True,
-    help="""Whether to only evaluate on the validation split. This is useful if you're
-    optimising hyperparameters, to avoid overfitting to the test sets.""",
+    help="""Whether to only evaluate on the test split. Only use this for your final
+    evaluation, as the test split should not be used for development.""",
 )
 @click.option(
     "--few-shot/--zero-shot",
@@ -268,17 +212,10 @@ def benchmark(
     raise_errors: bool,
     task: tuple[str],
     batch_size: str,
-    evaluate_train: bool,
     progress_bar: bool,
     save_results: bool,
     cache_dir: str,
-    use_token: bool,
-    token: str,
-    openai_api_key: str | None,
-    prefer_azure: bool,
-    azure_openai_api_key: str | None,
-    azure_openai_endpoint: str | None,
-    azure_openai_api_version: str | None,
+    api_key: str | None,
     force: bool,
     verbose: bool,
     framework: str | None,
@@ -287,14 +224,13 @@ def benchmark(
     load_in_4bit: bool | None,
     use_flash_attention: bool | None,
     clear_model_cache: bool,
-    only_validation_split: bool,
+    evaluate_test_split: bool,
     few_shot: bool,
     num_iterations: int,
     debug: bool,
 ) -> None:
     """Benchmark pretrained language models on language tasks."""
-    # Set up language variables
-    models = None if len(model) == 0 else list(model)
+    models = list(model)
     datasets = None if len(dataset) == 0 else list(dataset)
     languages: list[str] = list(language)
     model_languages = None if len(model_language) == 0 else list(model_language)
@@ -302,9 +238,7 @@ def benchmark(
     tasks = None if len(task) == 0 else list(task)
     batch_size_int = int(batch_size)
     device = Device[device.upper()] if device is not None else None
-    token_str_bool: str | bool = token if token != "" else use_token
 
-    # Initialise the benchmarker class
     benchmarker = Benchmarker(
         language=languages,
         model_language=model_languages,
@@ -314,15 +248,9 @@ def benchmark(
         batch_size=batch_size_int,
         progress_bar=progress_bar,
         save_results=save_results,
-        evaluate_train=evaluate_train,
         raise_errors=raise_errors,
         verbose=verbose,
-        token=token_str_bool,
-        openai_api_key=openai_api_key,
-        prefer_azure=prefer_azure,
-        azure_openai_api_key=azure_openai_api_key,
-        azure_openai_endpoint=azure_openai_endpoint,
-        azure_openai_api_version=azure_openai_api_version,
+        api_key=api_key,
         force=force,
         cache_dir=cache_dir,
         framework=framework,
@@ -331,7 +259,7 @@ def benchmark(
         load_in_4bit=load_in_4bit,
         use_flash_attention=use_flash_attention,
         clear_model_cache=clear_model_cache,
-        only_validation_split=only_validation_split,
+        evaluate_test_split=evaluate_test_split,
         few_shot=few_shot,
         num_iterations=num_iterations,
         debug=debug,
diff --git a/src/scandeval/constants.py b/src/scandeval/constants.py
new file mode 100644
index 000000000..5597fe109
--- /dev/null
+++ b/src/scandeval/constants.py
@@ -0,0 +1,43 @@
+"""Constants used throughout the project."""
+
+# This is used as input to generative models; it cannot be a special token
+DUMMY_FILL_VALUE = 100
+
+
+GENERATIVE_MODEL_TASKS = [
+    "text-generation",
+    "conversational",
+    # "text2text-generation",  # TODO: Add this when we support it
+]
+
+
+GENERATIVE_DATASET_TASKS = [
+    "knowledge",
+    "common-sense-reasoning",
+    "multiple-choice-reading-comprehension",
+]
+
+
+GENERATIVE_DATASET_SUPERTASKS = ["text-to-text", "text-modelling"]
+
+
+TASKS_USING_JSON = ["ner"]
+
+
+SUPERTASKS_USING_LOGPROBS = ["sequence-classification"]
+
+
+METRIC_ATTRIBUTES_TAKING_UP_MEMORY = ["cached_bertscorer"]
+
+
+GENERATIVE_TAGS = [
+    "trl",
+    "mistral",
+    "text-generation-inference",
+    "unsloth",
+    "text-generation",
+    "llama",
+]
+
+
+MAX_LOGPROBS = 10
diff --git a/src/scandeval/data_loading.py b/src/scandeval/data_loading.py
new file mode 100644
index 000000000..6e95a47b7
--- /dev/null
+++ b/src/scandeval/data_loading.py
@@ -0,0 +1,96 @@
+"""Functions related to the loading of the data."""
+
+import logging
+import sys
+import time
+
+from datasets import Dataset, DatasetDict, load_dataset
+from datasets.exceptions import DatasetsError
+from huggingface_hub.utils import HfHubHTTPError
+from numpy.random import Generator
+
+from .data_models import BenchmarkConfig, DatasetConfig
+from .exceptions import InvalidBenchmark
+from .utils import unscramble
+
+logger = logging.getLogger("scandeval")
+
+
+def load_data(
+    rng: Generator, dataset_config: "DatasetConfig", benchmark_config: "BenchmarkConfig"
+) -> list[DatasetDict]:
+    """Load the raw bootstrapped datasets.
+
+    Args:
+        rng:
+            The random number generator to use.
+        dataset_config:
+            The configuration for the dataset.
+        benchmark_config:
+            The configuration for the benchmark.
+
+    Returns:
+        A list of bootstrapped datasets, one for each iteration.
+    """
+    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):
+            logger.warning(
+                f"Failed to load dataset {dataset_config.huggingface_id!r}. Retrying..."
+            )
+            time.sleep(1)
+            continue
+        except HfHubHTTPError:
+            raise InvalidBenchmark("The Hugging Face Hub seems to be down.")
+    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"]})
+
+    if not benchmark_config.evaluate_test_split:
+        dataset["test"] = dataset["val"]
+
+    # Remove empty examples from the datasets
+    for text_feature in ["tokens", "text"]:
+        if text_feature in dataset["train"].features:
+            dataset = dataset.filter(lambda x: len(x[text_feature]) > 0)
+
+    # If we are testing then truncate the test set
+    if hasattr(sys, "_called_from_test"):
+        dataset["test"] = dataset["test"].select(range(1))
+
+    # Bootstrap the splits
+    bootstrapped_splits: dict[str, list[Dataset]] = dict()
+    for split in ["train", "val", "test"]:
+        bootstrap_indices = rng.integers(
+            0,
+            len(dataset[split]),
+            size=(benchmark_config.num_iterations, len(dataset[split])),
+        )
+        bootstrapped_splits[split] = [
+            dataset[split].select(bootstrap_indices[idx])
+            for idx in range(benchmark_config.num_iterations)
+        ]
+
+    datasets = [
+        DatasetDict(
+            {
+                split: bootstrapped_splits[split][idx]
+                for split in ["train", "val", "test"]
+            }
+        )
+        for idx in range(benchmark_config.num_iterations)
+    ]
+    return datasets
diff --git a/src/scandeval/config.py b/src/scandeval/data_models.py
similarity index 60%
rename from src/scandeval/config.py
rename to src/scandeval/data_models.py
index 21250c9d9..fc30ea9d0 100644
--- a/src/scandeval/config.py
+++ b/src/scandeval/data_models.py
@@ -1,12 +1,18 @@
-"""Configuration classes used throughout the project."""
-
+"""Data models used in ScandEval."""
+
+import collections.abc as c
+import importlib.metadata
+import json
+import pathlib
+import re
+import typing as t
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable
 
+import pydantic
 import torch
 
-if TYPE_CHECKING:
-    from .enums import Framework, ModelType
+from .enums import Device, Framework, ModelType
+from .types import ScoreDict
 
 
 @dataclass
@@ -37,8 +43,8 @@ class MetricConfig:
     pretty_name: str
     huggingface_id: str
     results_key: str
-    compute_kwargs: dict[str, Any] = field(default_factory=dict)
-    postprocessing_fn: Callable[[float], tuple[float, str]] = field(
+    compute_kwargs: dict[str, t.Any] = field(default_factory=dict)
+    postprocessing_fn: c.Callable[[float], tuple[float, str]] = field(
         default_factory=lambda: lambda raw_score: (100 * raw_score, f"{raw_score:.2%}")
     )
 
@@ -113,25 +119,8 @@ class BenchmarkConfig:
             Whether to raise errors instead of skipping them.
         cache_dir:
             Directory to store cached models and datasets.
-        evaluate_train:
-            Whether to evaluate on the training set.
-        token:
-            The authentication token for the Hugging Face Hub. If a boolean value is
-            specified then the token will be fetched from the Hugging Face CLI, where
-            the user has logged in through `huggingface-cli login`. If a string is
-            specified then it will be used as the token.
-        openai_api_key:
-            The API key for the OpenAI API. If None then OpenAI models will not be
-            benchmarked.
-        azure_openai_api_key:
-            The API key for the Azure OpenAI API. If None then Azure OpenAI models will
-            not be benchmarked.
-        azure_openai_endpoint:
-            The endpoint for the Azure OpenAI API. If None then Azure OpenAI models will
-            not be benchmarked.
-        azure_openai_api_version:
-            The api version for the Azure OpenAI API, e.g. "2023-12-01-preview". If
-            None then Azure OpenAI models will not be benchmarked.
+        api_key:
+            The API key to use for a given inference API.
         force:
             Whether to force the benchmark to run even if the results are already
             cached.
@@ -153,8 +142,8 @@ class BenchmarkConfig:
             generative models.
         clear_model_cache:
             Whether to clear the model cache after benchmarking each model.
-        only_validation_split:
-            Whether to only evaluate on the validation split.
+        evaluate_test_split:
+            Whether to evaluate on the test split.
         few_shot:
             Whether to only evaluate the model using few-shot evaluation. Only relevant
             if the model is generative.
@@ -170,16 +159,11 @@ class BenchmarkConfig:
     dataset_languages: list[Language]
     tasks: list[Task]
     datasets: list[str]
-    framework: "Framework | None"
+    framework: Framework | None
     batch_size: int
     raise_errors: bool
     cache_dir: str
-    evaluate_train: bool
-    token: bool | str | None
-    openai_api_key: str | None
-    azure_openai_api_key: str | None
-    azure_openai_endpoint: str | None
-    azure_openai_api_version: str | None
+    api_key: str | None
     force: bool
     progress_bar: bool
     save_results: bool
@@ -189,13 +173,103 @@ class BenchmarkConfig:
     load_in_4bit: bool | None
     use_flash_attention: bool | None
     clear_model_cache: bool
-    only_validation_split: bool
+    evaluate_test_split: bool
     few_shot: bool
     num_iterations: int
     debug: bool
     run_with_cli: bool
 
 
+class BenchmarkConfigParams(pydantic.BaseModel):
+    """The parameters for the benchmark configuration."""
+
+    model_config = pydantic.ConfigDict(protected_namespaces=())
+
+    progress_bar: bool
+    save_results: bool
+    task: str | list[str] | None
+    dataset: str | list[str] | None
+    language: str | list[str]
+    model_language: str | list[str] | None
+    dataset_language: str | list[str] | None
+    framework: Framework | str | None
+    device: Device | None
+    batch_size: int
+    raise_errors: bool
+    cache_dir: str
+    api_key: str | None
+    force: bool
+    verbose: bool
+    trust_remote_code: bool
+    load_in_4bit: bool | None
+    use_flash_attention: bool | None
+    clear_model_cache: bool
+    evaluate_test_split: bool
+    few_shot: bool
+    num_iterations: int
+    debug: bool
+    run_with_cli: bool
+
+
+class BenchmarkResult(pydantic.BaseModel):
+    """A benchmark result."""
+
+    dataset: str
+    task: str
+    dataset_languages: list[str]
+    model: str
+    results: ScoreDict
+    num_model_parameters: int
+    max_sequence_length: int
+    vocabulary_size: int
+    generative: bool
+    few_shot: bool
+    validation_split: bool
+    scandeval_version: str = importlib.metadata.version("scandeval")
+
+    @classmethod
+    def from_dict(cls, config: dict) -> "BenchmarkResult":
+        """Create a benchmark result from a dictionary.
+
+        Args:
+            config:
+                The configuration dictionary.
+
+        Returns:
+            The benchmark result.
+        """
+        # To be backwards compatible, we accept old results which changed the model
+        # name with parameters rather than adding them as explicit parameters
+        val_matches = re.search(r"\(.*val.*\)$", config["model"])
+        few_shot_matches = re.search(r"\(.*few-shot.*\)$", config["model"])
+        config["model"] = re.sub(
+            r"\(.*(few-shot|val).*\)$", "", config["model"]
+        ).strip()
+
+        # The default value for `few_shot` is True. It won't do anything if the model
+        # is not generative, so this is fine
+        if "generative" not in config:
+            config["generative"] = few_shot_matches is not None
+        if "few_shot" not in config:
+            config["few_shot"] = True
+
+        if "validation_split" not in config:
+            config["validation_split"] = val_matches is not None
+
+        return cls(**config)
+
+    def append_to_results(self, results_path: pathlib.Path) -> None:
+        """Append the benchmark result to the results file.
+
+        Args:
+            results_path:
+                The path to the results file.
+        """
+        json_str = json.dumps(self.model_dump())
+        with results_path.open("a") as f:
+            f.write("\n" + json_str)
+
+
 @dataclass
 class DatasetConfig:
     """Configuration for a dataset.
@@ -299,13 +373,86 @@ class ModelConfig:
 
     model_id: str
     revision: str
-    framework: "Framework"
+    framework: Framework
     task: str
     languages: list[Language]
-    model_type: "ModelType | str"
+    model_type: ModelType
     model_cache_dir: str
     adapter_base_model_id: str | None
 
     def __hash__(self) -> int:
         """Return a hash of the model configuration."""
         return hash(self.model_id)
+
+
+@dataclass
+class PreparedModelInputs:
+    """The inputs to a model.
+
+    Attributes:
+        texts:
+            The texts to input to the model. Can be None if the input IDs and attention
+            mask are provided instead.
+        input_ids:
+            The input IDs of the texts. Can be None if the texts are provided instead.
+        attention_mask:
+            The attention mask of the texts. Can be None if the texts are provided
+            instead.
+    """
+
+    texts: list[str] | None = None
+    input_ids: torch.Tensor | None = None
+    attention_mask: torch.Tensor | None = None
+
+
+@dataclass
+class GenerativeModelOutput:
+    """The output of a generative model.
+
+    Attributes:
+        sequences:
+            The generated sequences.
+        scores:
+            The scores of the sequences. This is an array of shape (batch_size,
+            num_tokens, num_logprobs, 2), where the last dimension contains the
+            token and its logprob. Can be None if the scores are not available.
+    """
+
+    sequences: list[str]
+    scores: list[list[list[tuple[str, float]]]] | None = None
+
+
+@dataclass
+class SingleGenerativeModelOutput:
+    """A single output of a generative model.
+
+    Attributes:
+        sequence:
+            The generated sequence.
+        scores:
+            The scores of the sequence. This is an array of shape (num_tokens,
+            num_logprobs, 2), where the last dimension contains the token and its
+            logprob. Can be None if the scores are not available.
+    """
+
+    sequence: str
+    scores: list[list[tuple[str, float]]] | None = None
+
+
+@dataclass
+class HFModelInfo:
+    """Information about a Hugging Face model.
+
+    Attributes:
+        pipeline_tag:
+            The pipeline tag of the model.
+        tags:
+            The other tags of the model.
+        adapter_base_model_id:
+            The model ID of the base model if the model is an adapter model. Can be None
+            if the model is not an adapter model.
+    """
+
+    pipeline_tag: str
+    tags: list[str]
+    adapter_base_model_id: str | None
diff --git a/src/scandeval/dataset_configs.py b/src/scandeval/dataset_configs.py
index bbc90cdb5..b4516199c 100644
--- a/src/scandeval/dataset_configs.py
+++ b/src/scandeval/dataset_configs.py
@@ -1,6 +1,6 @@
 """All dataset configurations used in ScandEval."""
 
-from .config import DatasetConfig
+from .data_models import DatasetConfig
 from .languages import DA, DE, EN, FO, IS, NB, NL, NN, NO, SV, get_all_languages
 from .tasks import COMMON_SENSE, KNOW, LA, MCRC, NER, RC, SENT, SPEED, SUMM
 
@@ -64,7 +64,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Recension: {text}\n\nKlassificera sentimentet i recensionen. "
     "Svara med 'positiv', 'neutral' eller 'negativ'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 ANGRY_TWEETS_CONFIG = DatasetConfig(
@@ -83,7 +83,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Tweet: {text}\n\nKlassificer sentimentet i tweetet. Svar kun "
     "med 'positiv', 'neutral' eller 'negativ'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 NOREC_CONFIG = DatasetConfig(
@@ -102,7 +102,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Anmeldelse: {text}\n\nKlassifiser sentimentet i anmeldelsen. "
     "Svar med 'positiv', 'nøytral' eller 'negativ'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HOTTER_AND_COLDER_SENTIMENT_CONFIG = DatasetConfig(
@@ -121,7 +121,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Texti: {text}\n\nFlokkaðu tilfinninguna í textanum. "
     "Svaraðu með 'jákvætt', 'hlutlaust' eða 'neikvætt'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SB10K_CONFIG = DatasetConfig(
@@ -140,7 +140,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Tweet: {text}\n\nKlassifizieren Sie die Stimmung im Tweet. "
     "Antworten Sie mit 'positiv', 'neutral' oder 'negativ'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 DUTCH_SOCIAL_CONFIG = DatasetConfig(
@@ -159,7 +159,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Tweet: {text}\n\nClassificeer het sentiment in de tweet. "
     "Antwoord met 'positief', 'neutraal' of 'negatief'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SST5_CONFIG = DatasetConfig(
@@ -178,7 +178,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Text: {text}\n\nClassify the sentiment in the text. Answer "
     "with 'positive', 'neutral' or 'negative'.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 FOSENT_CONFIG = DatasetConfig(
@@ -196,7 +196,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Tekstur: {text}\n\nFlokka lyndið í tekstinum. Svara við "
     "'positivt', 'neutralt' ella 'negativt'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 
@@ -527,7 +527,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Mening: {text}\n\nBestäm om meningen är grammatiskt korrekt "
     "eller inte. Svara med 'ja' om meningen är korrekt och 'nej' om den inte är.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_DA_CONFIG = DatasetConfig(
@@ -542,7 +542,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Sætning: {text}\n\nBestem om sætningen er grammatisk korrekt "
     "eller ej. Svar med 'ja', hvis sætningen er korrekt, og 'nej', hvis den ikke er.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_NB_CONFIG = DatasetConfig(
@@ -557,7 +557,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "eller ikke. Svar med 'ja' hvis setningen er korrekt og 'nei' hvis den ikke er.",
     prompt_label_mapping=dict(correct="ja", incorrect="nei"),
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_NN_CONFIG = DatasetConfig(
@@ -572,7 +572,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Setning: {text}\n\nBestem om setningen er grammatisk korrekt "
     "eller ikke. Svar med 'ja' hvis setningen er korrekt og 'nei' hvis den ikke er.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_IS_CONFIG = DatasetConfig(
@@ -588,7 +588,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "rétt eða ekki. Svarið skal vera 'já' ef setningin er rétt og 'nei' ef hún er "
     "ekki.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_FO_CONFIG = DatasetConfig(
@@ -604,7 +604,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "rættur ella ikki. Svarið skal vera 'ja' um setningurin er rættur og 'nei' um "
     "hann ikki er.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_DE_CONFIG = DatasetConfig(
@@ -620,7 +620,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "korrekt ist oder nicht. Antworten Sie mit 'ja', wenn der Satz korrekt ist und "
     "'nein', wenn er es nicht ist.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_NL_CONFIG = DatasetConfig(
@@ -635,7 +635,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Zin: {text}\n\nBepaal of de zin grammaticaal correct is of "
     "niet. Antwoord met 'ja' als de zin correct is en 'nee' als dat niet het geval is.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 SCALA_EN_CONFIG = DatasetConfig(
@@ -652,7 +652,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "grammatically correct or not. Reply with 'yes' if the sentence is correct and "
     "'no' if it is not.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 DUTCH_COLA_CONFIG = DatasetConfig(
@@ -703,7 +703,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "rétt eða ekki. Svarið skal vera 'já' ef setningin er rétt og 'nei' ef hún er "
     "ekki.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -719,7 +719,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Setning: {text}\n\nGreinið hvort setningin er málfræðilega "
     "rétt eða ekki. Svarið skal vera 'já' ef setningin er rétt og 'nei' ef hún er ekki.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -735,7 +735,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Setning: {text}\n\nGreinið hvort setningin er málfræðilega "
     "rétt eða ekki. Svarið skal vera 'já' ef setningin er rétt og 'nei' ef hún er ekki.",
     num_few_shot_examples=12,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1084,7 +1084,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 DANISH_CITIZEN_TESTS_CONFIG = DatasetConfig(
@@ -1099,7 +1099,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_NO_CONFIG = DatasetConfig(
@@ -1115,7 +1115,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørsmål: {text}\n\nBesvar følgende spørsmål med 'a', 'b', "
     "'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_SV_CONFIG = DatasetConfig(
@@ -1131,7 +1131,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Fråga: {text}\n\nBesvara följande fråga med 'a', 'b', 'c' "
     "eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_IS_CONFIG = DatasetConfig(
@@ -1147,7 +1147,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spurningar: {text}\n\nSvaraðu eftirfarandi spurningum með 'a', "
     "'b', 'c' eða 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1164,7 +1164,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Frage: {text}\n\nBeantworten Sie die obige Frage mit 'a', 'b', "
     "'c' oder 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_NL_CONFIG = DatasetConfig(
@@ -1180,7 +1180,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Vraag: {text}\n\nBeantwoord de bovenstaande vraag met 'a', 'b', "
     "'c' of 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_CONFIG = DatasetConfig(
@@ -1195,7 +1195,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Question: {text}\n\nAnswer the above question by "
     "replying with 'a', 'b', 'c' or 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 MMLU_DA_CONFIG = DatasetConfig(
@@ -1211,7 +1211,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1228,7 +1228,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1245,7 +1245,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørsmål: {text}\n\nBesvar følgende spørsmål med 'a', 'b', "
     "'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1262,7 +1262,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Fråga: {text}\n\nBesvara följande fråga med 'a', 'b', 'c' "
     "eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1279,7 +1279,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spurningar: {text}\n\nSvaraðu eftirfarandi spurningum með 'a', "
     "'b', 'c' eða 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 ARC_DE_CONFIG = DatasetConfig(
@@ -1295,7 +1295,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Frage: {text}\n\nBeantworten Sie die obige Frage mit 'a', 'b', "
     "'c' oder 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1312,7 +1312,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Vraag: {text}\n\nBeantwoord de bovenstaande vraag met 'a', "
     "'b', 'c' of 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1328,7 +1328,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Question: {text}\n\nAnswer the above question by "
     "replying with 'a', 'b', 'c' or 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1350,7 +1350,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_NO_CONFIG = DatasetConfig(
@@ -1366,7 +1366,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørsmål: {text}\n\nBesvar følgende spørsmål med 'a', 'b', "
     "'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_SV_CONFIG = DatasetConfig(
@@ -1382,7 +1382,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Fråga: {text}\n\nBesvara följande fråga med 'a', 'b', 'c' "
     "eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_IS_CONFIG = DatasetConfig(
@@ -1398,7 +1398,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spurningar: {text}\n\nSvaraðu eftirfarandi spurningum með 'a', "
     "'b', 'c' eða 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1415,7 +1415,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spurningar: {text}\n\nSvaraðu eftirfarandi spurningum með 'a', "
     "'b', 'c' eða 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_DE_CONFIG = DatasetConfig(
@@ -1431,7 +1431,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Frage: {text}\n\nBeantworten Sie die obige Frage mit 'a', 'b', "
     "'c' oder 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_NL_CONFIG = DatasetConfig(
@@ -1447,7 +1447,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Vraag: {text}\n\nBeantwoord de bovenstaande vraag met 'a', 'b', "
     "'c' of 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 HELLASWAG_CONFIG = DatasetConfig(
@@ -1463,7 +1463,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Question: {text}\n\nAnswer the above question by "
     "replying with 'a', 'b', 'c' or 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
 
 # TODO: Faroese common sense reasoning
@@ -1485,7 +1485,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørgsmål: {text}\n\nBesvar ovenstående spørgsmål ved at "
     "svare med 'a', 'b', 'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1503,7 +1503,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Fråga: {text}\n\nBesvara följande fråga med 'a', 'b', 'c' "
     "eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1520,7 +1520,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Spørsmål: {text}\n\nBesvar følgende spørsmål med 'a', 'b', "
     "'c' eller 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1538,7 +1538,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     "'b', 'c' eða 'd'.",
     prompt_label_mapping=dict(a="a", b="b", c="c", d="d"),
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1556,7 +1556,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Frage: {text}\n\nBeantworten Sie die obige Frage mit 'a', 'b', "
     "'c' oder 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1574,7 +1574,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Vraag: {text}\n\nBeantwoord de bovenstaande vraag met 'a', "
     "'b', 'c' of 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1591,7 +1591,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     instruction_prompt="Question: {text}\n\nAnswer the above question by "
     "replying with 'a', 'b', 'c' or 'd'.",
     num_few_shot_examples=5,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
     unofficial=True,
 )
 
@@ -1608,5 +1608,5 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     prompt_template="",
     instruction_prompt="",
     num_few_shot_examples=0,
-    max_generated_tokens=1,
+    max_generated_tokens=5,
 )
diff --git a/src/scandeval/dataset_factory.py b/src/scandeval/dataset_factory.py
deleted file mode 100644
index a5d85edbb..000000000
--- a/src/scandeval/dataset_factory.py
+++ /dev/null
@@ -1,67 +0,0 @@
-"""Factory which produces datasets from a configuration."""
-
-from typing import TYPE_CHECKING, Type
-
-from .dataset_configs import get_dataset_config
-from .utils import get_class_by_name
-
-if TYPE_CHECKING:
-    from .benchmark_dataset import BenchmarkDataset
-    from .config import BenchmarkConfig, DatasetConfig
-
-
-class DatasetFactory:
-    """Factory which produces datasets from a configuration.
-
-    Attributes:
-        benchmark_config:
-            The benchmark configuration to be used in all datasets constructed.
-    """
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the dataset factory.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration to be used in all datasets constructed.
-        """
-        self.benchmark_config = benchmark_config
-
-    def build_dataset(self, dataset: "str | DatasetConfig") -> "BenchmarkDataset":
-        """Build a dataset from a configuration or a name.
-
-        Args:
-            dataset:
-                The name of the dataset, or the dataset configuration.
-
-        Returns:
-            BenchmarkDataset:
-                The benchmark dataset.
-        """
-        # Get the dataset configuration
-        if isinstance(dataset, str):
-            dataset_config = get_dataset_config(dataset)
-        else:
-            dataset_config = dataset
-
-        # Get the benchmark class based on the task
-        potential_class_names = [
-            dataset_config.name,
-            dataset_config.task.name,
-            dataset_config.task.supertask,
-        ]
-        benchmark_cls: Type["BenchmarkDataset"] | None = get_class_by_name(
-            class_name=potential_class_names
-        )
-        if not benchmark_cls:
-            raise ValueError(
-                "Could not find a benchmark class for any of the following potential "
-                f"names: {', '.join(potential_class_names)}."
-            )
-
-        # Create the dataset
-        dataset_obj = benchmark_cls(
-            dataset_config=dataset_config, benchmark_config=self.benchmark_config
-        )
-
-        return dataset_obj
diff --git a/src/scandeval/enums.py b/src/scandeval/enums.py
index 34ed48787..53db47532 100644
--- a/src/scandeval/enums.py
+++ b/src/scandeval/enums.py
@@ -56,20 +56,23 @@ class ModelType(AutoStrEnum):
     Attributes:
         FRESH:
             Randomly initialised Hugging Face model.
-        HF:
-            Model from the Hugging Face Hub.
+        HF_HUB_ENCODER:
+            Hugging Face encoder model from the Hub.
+        HF_HUB_GENERATIVE:
+            Hugging Face generative model from the Hub.
         LOCAL:
-            Locally stored Hugging Face model.
-        OPENAI:
-            Model from OpenAI.
+            Locally stored model.
+        API:
+            Model accessed through an API.
         HUMAN:
             Human evaluator.
     """
 
     FRESH = auto()
-    HF = auto()
+    HF_HUB_ENCODER = auto()
+    HF_HUB_GENERATIVE = auto()
     LOCAL = auto()
-    OPENAI = auto()
+    API = auto()
     HUMAN = auto()
 
 
@@ -88,3 +91,20 @@ class DataType(AutoStrEnum):
     FP32 = auto()
     FP16 = auto()
     BF16 = auto()
+
+
+class BatchingPreference(AutoStrEnum):
+    """The preference for batching.
+
+    Attributes:
+        NO_PREFERENCE:
+            No preference for batching.
+        SINGLE_SAMPLE:
+            Single sample batching.
+        ALL_AT_ONCE:
+            All samples at once batching.
+    """
+
+    NO_PREFERENCE = auto()
+    SINGLE_SAMPLE = auto()
+    ALL_AT_ONCE = auto()
diff --git a/src/scandeval/exceptions.py b/src/scandeval/exceptions.py
index 216f83a3f..473f17591 100644
--- a/src/scandeval/exceptions.py
+++ b/src/scandeval/exceptions.py
@@ -107,6 +107,7 @@ def __init__(self, extra: str):
             extra:
                 The extra that needs to be installed.
         """
+        self.extra = extra
         self.message = (
             f"The model you are trying to load requires the `{extra}` extra to be "
             f"installed. To install the `{extra}` extra, please run `pip install "
@@ -125,6 +126,7 @@ def __init__(self, package: str):
             package:
                 The package that needs to be manually installed.
         """
+        self.package = package
         self.message = (
             f"The model you are trying to load requires the `{package}` package to be "
             f"installed - please run `pip install {package}` and try again."
@@ -146,6 +148,8 @@ def __init__(self, cli_argument: str, script_argument: str, run_with_cli: bool):
             run_with_cli:
                 Whether the benchmark is being run with the CLI.
         """
+        self.cli_argument = cli_argument
+        self.script_argument = script_argument
         if run_with_cli:
             self.message = (
                 f"The model you are trying to load requires the `{cli_argument}` "
@@ -161,34 +165,19 @@ def __init__(self, cli_argument: str, script_argument: str, run_with_cli: bool):
         super().__init__(self.message)
 
 
-class MissingHuggingFaceToken(InvalidModel):
-    """The evaluation requires a Hugging Face token."""
+class NeedsEnvironmentVariable(InvalidModel):
+    """The evaluation requires an environment variable to be set."""
 
-    def __init__(self, run_with_cli: bool):
+    def __init__(self, env_var: str):
         """Initialize the exception.
 
         Args:
-            run_with_cli:
-                Whether the benchmark is being run with the CLI.
+            env_var:
+                The environment variable that needs to be set.
         """
+        self.env_var = env_var
         self.message = (
-            "The model you are trying to load requires a Hugging Face token. "
+            f"The model you are trying to load requires the `{env_var}` environment "
+            "variable to be set. Please set the environment variable and try again."
         )
-        if run_with_cli:
-            self.message += (
-                "Please run `huggingface-cli login` to login to the Hugging Face "
-                "Hub and try again. Alternatively, you can pass your Hugging Face Hub "
-                "token directly to the `scandeval` command with the `--token <token>` "
-                "argument."
-            )
-        else:
-            self.message += (
-                "Please pass your Hugging Face Hub token to the `Benchmarker` class "
-                "with the `token=<token>` argument  to the `Benchmarker` class and try "
-                "again. Alternatively, you can also simply pass `token=True` (which is "
-                "the default) to the `Benchmarker` class, which assumes that you have "
-                "logged into the Hugging Face Hub. This can be done by running "
-                "`huggingface-cli login` in the terminal or `from huggingface_hub "
-                "import login; login()` in a Python script."
-            )
         super().__init__(self.message)
diff --git a/src/scandeval/finetuning.py b/src/scandeval/finetuning.py
index 410e44f9c..ce9a718be 100644
--- a/src/scandeval/finetuning.py
+++ b/src/scandeval/finetuning.py
@@ -3,107 +3,59 @@
 import importlib.util
 import logging
 import sys
-from collections import defaultdict
+import typing as t
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Type
 
 import torch
-from datasets import Dataset
+from datasets import DatasetDict
 from tqdm.auto import tqdm
 from transformers import (
     EarlyStoppingCallback,
     IntervalStrategy,
-    PreTrainedModel,
     PrinterCallback,
     ProgressCallback,
     TrainingArguments,
 )
 from transformers.trainer import OptimizerNames
 
+from .benchmark_modules import BenchmarkModule
 from .callbacks import NeverLeaveProgressCallback
 from .enums import DataType
 from .exceptions import InvalidBenchmark, NaNValueInModelOutput
 from .model_loading import load_model
 from .utils import block_terminal_output, clear_memory, enforce_reproducibility
 
-if TYPE_CHECKING:
-    from transformers import DataCollator, Trainer
+if t.TYPE_CHECKING:
+    from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-    from .protocols import Tokenizer
-
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 def finetune(
-    itr: tqdm,
-    train: Dataset,
-    val: Dataset,
-    tests: list[Dataset],
-    prepared_train: Dataset,
-    prepared_val: Dataset,
-    prepared_tests: list[Dataset],
-    model: PreTrainedModel,
-    tokenizer: "Tokenizer",
+    model: BenchmarkModule,
+    datasets: list[DatasetDict],
     model_config: "ModelConfig",
-    benchmark_config: "BenchmarkConfig",
     dataset_config: "DatasetConfig",
-    compute_metrics: Callable,
-    data_collator: "DataCollator",
-    trainer_class: Type["Trainer"],
-    evaluate_inputs_fn: Callable[..., dict[str, Any]],
-    preprocess_logits_for_metrics: Callable[
-        [torch.Tensor | tuple, torch.Tensor], torch.Tensor | tuple
-    ],
-) -> dict[str, list[dict[str, float]]]:
+    benchmark_config: "BenchmarkConfig",
+) -> list[dict[str, float]]:
     """Evaluate a model on a dataset through finetuning.
 
     Args:
-        itr:
-            The progress bar iterator.
-        train:
-            The training dataset.
-        val:
-            The validation dataset.
-        tests:
-            The bootstrapped test datasets.
-        prepared_train:
-            The prepared training dataset.
-        prepared_val:
-            The prepared validation dataset.
-        prepared_tests:
-            The prepared bootstrapped test datasets.
         model:
             The model to evaluate.
-        tokenizer:
-            The tokenizer to use.
+        datasets:
+            The datasets to use for training and evaluation.
         model_config:
             The configuration of the model.
-        benchmark_config:
-            The benchmark configuration.
         dataset_config:
             The dataset configuration.
-        compute_metrics:
-            The function used to compute the metrics.
-        data_collator:
-            The data collator to use.
-        trainer_class:
-            The Trainer class to use.
-        evaluate_inputs_fn:
-            A function that generates the appropriate inputs for the `Trainer.evaluate`
-            method.
-        preprocess_logits_for_metrics:
-            A function that preprocesses the logits before they are passed to the
-            `compute_metrics` function. This helps prevent memory issues during
-            evaluation.
+        benchmark_config:
+            The benchmark configuration.
 
     Returns:
-        A dictionary containing the scores, with keys "test" and maybe "train", with
-        values being lists of dicts containing the scores for each metric for each
-        iteration.
+        A list of dicts containing the scores for each metric for each iteration.
     """
-    scores: dict[str, list[dict[str, float]]] = defaultdict(list)
-
+    # Set the data type to use for the model weights
     using_cuda = benchmark_config.device == torch.device("cuda")
     if using_cuda and torch.cuda.is_bf16_supported():
         dtype = DataType.BF16
@@ -113,7 +65,12 @@ def finetune(
         dtype = DataType.FP32
 
     bs: int = benchmark_config.batch_size
-    for idx in itr:
+    scores: list[dict[str, float]] = list()
+    for idx in tqdm(
+        iterable=range(benchmark_config.num_iterations),
+        desc="Benchmarking",
+        disable=not benchmark_config.progress_bar,
+    ):
         # Set variable that tracks whether we need to initialize new models in
         # the single iteration call
         model_already_initialized = idx == 0
@@ -126,18 +83,9 @@ def finetune(
                     del model
                 except UnboundLocalError:
                     pass
-                try:
-                    del tokenizer
-                except UnboundLocalError:
-                    pass
                 clear_memory()
 
             try:
-                test = tests[idx]
-                prepared_test = prepared_tests[idx]
-                assert isinstance(test, Dataset)
-                assert isinstance(prepared_test, Dataset)
-
                 # Re-block terminal output, as it gets unblocked by the `transformers`
                 # package before training
                 block_terminal_output()
@@ -151,32 +99,17 @@ def finetune(
                 )
 
                 itr_scores = finetune_single_iteration(
+                    model=model if model_already_initialized else None,
+                    dataset=datasets[idx],
                     iteration_idx=idx,
-                    model_config=model_config,
-                    train=train,
-                    prepared_train=prepared_train,
-                    prepared_val=prepared_val,
-                    test=test,
-                    prepared_test=prepared_test,
                     training_args=training_args,
-                    benchmark_config=benchmark_config,
+                    model_config=model_config,
                     dataset_config=dataset_config,
-                    data_collator=data_collator,
-                    compute_metrics=compute_metrics,
-                    tokenizer=tokenizer if model_already_initialized else None,
-                    model=model if model_already_initialized else None,
-                    trainer_class=trainer_class,
-                    evaluate_inputs_fn=evaluate_inputs_fn,
-                    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+                    benchmark_config=benchmark_config,
                 )
 
-                if "train" in itr_scores:
-                    logger.debug(
-                        f"Train scores for iteration {idx}: {itr_scores['train']}"
-                    )
-                    scores["train"].append(itr_scores["train"])
-                scores["test"].append(itr_scores["test"])
-                logger.debug(f"Test scores for iteration {idx}: {itr_scores['test']}")
+                scores.append(itr_scores)
+                logger.debug(f"Test scores for iteration {idx}: {itr_scores}")
 
                 break
 
@@ -221,101 +154,58 @@ def finetune(
 
 
 def finetune_single_iteration(
+    model: BenchmarkModule | None,
+    dataset: DatasetDict,
     iteration_idx: int,
-    model_config: "ModelConfig",
-    train: Dataset,
-    test: Dataset,
-    prepared_train: Dataset,
-    prepared_val: Dataset,
-    prepared_test: Dataset,
     training_args: TrainingArguments,
-    benchmark_config: "BenchmarkConfig",
+    model_config: "ModelConfig",
     dataset_config: "DatasetConfig",
-    data_collator: "DataCollator",
-    compute_metrics: Callable,
-    tokenizer: "Tokenizer | None",
-    model: PreTrainedModel | None,
-    trainer_class: Type["Trainer"],
-    evaluate_inputs_fn: Callable[..., dict[str, Any]],
-    preprocess_logits_for_metrics: Callable[
-        [torch.Tensor | tuple, torch.Tensor], torch.Tensor | tuple
-    ],
-) -> dict[str, dict[str, float]]:
+    benchmark_config: "BenchmarkConfig",
+) -> dict[str, float]:
     """Run a single iteration of a benchmark.
 
     Args:
+        model:
+            The model to use in the benchmark. If None then a new model will be loaded.
+        dataset:
+            The dataset to use for training and evaluation.
         iteration_idx:
             The index of the iteration.
-        model_config:
-            The model configuration.
-        train:
-            The original training dataset.
-        test:
-            The original test dataset.
-        prepared_train:
-            The prepared training dataset.
-        prepared_val:
-            The prepared validation dataset.
-        prepared_test:
-            The prepared test dataset.
         training_args:
             The training arguments.
-        benchmark_config:
-            The benchmark configuration.
+        model_config:
+            The model configuration.
         dataset_config:
             The dataset configuration.
-        data_collator:
-            The data collator.
-        compute_metrics:
-            The function to compute the metrics.
-        tokenizer:
-            The tokenizer to use in the benchmark. If None then a new tokenizer
-            will be loaded.
-        model:
-            The model to use in the benchmark. If None then a new model will be
-            loaded.
-        trainer_class:
-            The trainer class to use.
-        evaluate_inputs_fn:
-            A function that generates the appropriate inputs for the `Trainer.evaluate`
-            method.
-        preprocess_logits_for_metrics:
-            A function that preprocesses the logits before they are passed to the
-            `compute_metrics` function. This helps prevent memory issues during
-            evaluation.
+        benchmark_config:
+            The benchmark configuration.
 
     Returns:
-        A dictionary containing the scores for the current iteration, with keys `train`
-        and `test`.
+        The scores for the test dataset.
     """
-    scores: dict[str, dict[str, float]] = dict()
-
     # Set random seeds to enforce reproducibility of the randomly initialised
     # weights
     seed = 4242 + iteration_idx
     enforce_reproducibility(framework=model_config.framework, seed=seed)
 
-    if tokenizer is None or model is None:
-        model_or_generative_model, tokenizer = load_model(
+    if model is None:
+        model = load_model(
             model_config=model_config,
             dataset_config=dataset_config,
             benchmark_config=benchmark_config,
         )
-        assert isinstance(model_or_generative_model, PreTrainedModel)
-        model = model_or_generative_model
 
-    compute_metrics = partial(compute_metrics, id2label=dataset_config.id2label)
-    early_stopping = EarlyStoppingCallback(early_stopping_patience=2)
-    trainer = trainer_class(
-        model=model,
+    trainer = model.trainer_class(
+        model=model.get_pytorch_module(),
+        tokenizer=model.get_tokenizer(),
         args=training_args,
-        train_dataset=prepared_train,
-        eval_dataset=prepared_val,
-        tokenizer=tokenizer,
-        compute_metrics=compute_metrics,
-        callbacks=[early_stopping],
-        data_collator=data_collator,
-        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["val"],
+        compute_metrics=partial(
+            model.compute_metrics, id2label=dataset_config.id2label
+        ),
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
+        data_collator=model.data_collator,
     )
 
     if not benchmark_config.verbose:
@@ -339,30 +229,22 @@ def no_logging(logs: dict[str, float]) -> None:
 
     try:
         trainer.train()
-
-        if benchmark_config.evaluate_train:
-            with torch.inference_mode():
-                evaluate_inputs = evaluate_inputs_fn(
-                    dataset=train,
-                    prepared_dataset=prepared_train,
-                    metric_key_prefix="train",
-                )
-                train_scores = trainer.evaluate(**evaluate_inputs)
-            scores["train"] = train_scores
-
         with torch.inference_mode():
-            evaluate_inputs = evaluate_inputs_fn(
-                dataset=test, prepared_dataset=prepared_test, metric_key_prefix="test"
-            )
-            test_scores = trainer.evaluate(**evaluate_inputs)
-        scores["test"] = test_scores
-
-        return scores
+            try:
+                test_scores = trainer.evaluate(
+                    eval_dataset=dataset["test"],
+                    orig_eval_dataset=dataset["original_test"],
+                    metric_key_prefix="test",
+                )
+            except TypeError:
+                test_scores = trainer.evaluate(
+                    eval_dataset=dataset["test"], metric_key_prefix="test"
+                )
+        return test_scores
 
     except NaNValueInModelOutput as e:
         del trainer
         del model
-        del tokenizer
         clear_memory()
         raise e
 
diff --git a/src/scandeval/generation.py b/src/scandeval/generation.py
index 2bddbd451..5d155ec2f 100644
--- a/src/scandeval/generation.py
+++ b/src/scandeval/generation.py
@@ -2,98 +2,53 @@
 
 import logging
 import sys
-import warnings
-from collections import defaultdict
+import typing as t
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
 
-import torch
-from datasets import Dataset
-from torch.utils.data import DataLoader
+import more_itertools as mit
+from datasets import Dataset, DatasetDict
 from tqdm.auto import tqdm
-from transformers import (
-    GenerationConfig,
-    PreTrainedTokenizerBase,
-    StoppingCriteria,
-    StoppingCriteriaList,
-)
-from transformers.modeling_utils import ModelOutput
 
+from .benchmark_modules import BenchmarkModule
+from .enums import BatchingPreference
 from .exceptions import InvalidBenchmark
 from .model_cache import (
     ModelCache,
     load_cached_model_outputs,
     split_dataset_into_cached_and_non_cached,
 )
-from .openai_models import OpenAIModel
-from .structured_generation_utils import (
-    get_ner_logits_processors,
-    get_ner_prefix_allowed_tokens_fn,
-)
-from .tasks import LA, NER, RC, SENT, SUMM
-from .utils import SUPERTASKS_USING_LOGPROBS, clear_memory, get_end_of_chat_token_ids
-from .vllm_models import VLLMModel
-
-if TYPE_CHECKING:
-    from transformers import DataCollator
+from .utils import clear_memory
 
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-    from .protocols import GenerativeModel, Tokenizer
+if t.TYPE_CHECKING:
+    from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 def generate(
-    itr: tqdm,
-    prepared_train: Dataset,
-    prepared_tests: list[Dataset],
-    model: "GenerativeModel",
+    model: "BenchmarkModule",
+    datasets: list[DatasetDict],
     model_config: "ModelConfig",
-    tokenizer: "Tokenizer",
-    data_collator: "DataCollator",
-    compute_metrics: Callable,
-    extract_labels_fn: Callable[..., list[Any]],
-    benchmark_config: "BenchmarkConfig",
     dataset_config: "DatasetConfig",
-) -> dict[str, list[dict[str, float]]]:
+    benchmark_config: "BenchmarkConfig",
+) -> list[dict[str, float]]:
     """Evaluate a model on a dataset through generation.
 
     Args:
-        itr:
-            The progress bar iterator.
-        prepared_train:
-            The prepared training dataset.
-        prepared_tests:
-            The prepared bootstrapped test datasets.
-        num_iter:
-            The number of iterations to run.
-        rng:
-            The random number generator.
         model:
             The model to evaluate.
+        datasets:
+            The datasets to evaluate on.
         model_config:
             The configuration of the model.
-        tokenizer:
-            The tokenizer to use for the model. If `None` then the model's
-            tokenizer will be used.
-        data_collator:
-            The data collator to use for the model.
-        compute_metrics:
-            The function to use to compute the metrics.
-        extract_labels_fn:
-            The function to use to extract the labels from the model output.
         benchmark_config:
             The configuration of the benchmark.
         dataset_config:
             The configuration of the dataset.
 
     Returns:
-        A dictionary containing the scores, with keys "test" and maybe "train", with
-        values being lists of dicts containing the scores for each metric for each
-        iteration.
+        A list of dictionaries containing the test scores.
     """
-    scores: dict[str, list[dict[str, float]]] = defaultdict(list)
-
     # Set up the name of the model output cache. If we are testing then we save the
     # model outputs to a different cache and ensure that that cache is deleted before
     # the next test, to ensure that the tests are independent of each other
@@ -115,81 +70,23 @@ def generate(
         max_generated_tokens=dataset_config.max_generated_tokens,
     )
 
-    for idx in itr:
-        prepared_test = prepared_tests[idx]
-        assert isinstance(prepared_test, Dataset)
-
-        generation_kwargs = dict(
+    scores: list[dict[str, float]] = list()
+    for idx in tqdm(
+        iterable=range(benchmark_config.num_iterations),
+        desc="Benchmarking",
+        disable=not benchmark_config.progress_bar,
+    ):
+        test_scores = generate_single_iteration(
             model=model,
-            tokenizer=tokenizer,
-            data_collator=data_collator,
-            compute_metrics=compute_metrics,
-            extract_labels_fn=extract_labels_fn,
-            dataset_config=dataset_config,
+            dataset=datasets[idx]["test"],
             cache=cache,
+            dataset_config=dataset_config,
+            benchmark_config=benchmark_config,
         )
 
-        def update_scores(
-            scores: dict[str, list[dict[str, float]]],
-            benchmark_config: "BenchmarkConfig",
-        ) -> dict[str, list[dict[str, float]]]:
-            """Perform a single iteration of generation and update the scores.
-
-            Args:
-                scores:
-                    The scores so far.
-                benchmark_config:
-                    The configuration of the benchmark.
-
-            Returns:
-                The updated scores.
-            """
-            test_scores = generate_single_iteration(
-                prepared_dataset=prepared_test,
-                benchmark_config=benchmark_config,
-                **generation_kwargs,
-            )
-            logger.debug(f"Test scores for iteration {idx}: {test_scores}")
-            scores["test"].append(test_scores)
-
-            if benchmark_config.evaluate_train:
-                train_scores = generate_single_iteration(
-                    prepared_dataset=prepared_train,
-                    benchmark_config=benchmark_config,
-                    **generation_kwargs,
-                )
-                logger.debug(f"Train scores for iteration {idx}: {train_scores}")
-                scores["train"].append(train_scores)
-
-            clear_memory()
-            return scores
-
-        if isinstance(model, VLLMModel):
-            scores = update_scores(scores=scores, benchmark_config=benchmark_config)
-        else:
-            while True:
-                try:
-                    scores = update_scores(
-                        scores=scores, benchmark_config=benchmark_config
-                    )
-                    break
-                except Exception as e:
-                    oom_error = [
-                        "CUDA out of memory",
-                        "CUDA error",
-                        "MPS backend out of memory",
-                        "Too many parallel completions requested.",  # OpenAI specific
-                    ]
-                    if isinstance(model, VLLMModel) or all(
-                        error not in str(e) for error in oom_error
-                    ):
-                        raise InvalidBenchmark(str(e))
-                    clear_memory()
-                    benchmark_config.batch_size //= 2
-                    if benchmark_config.batch_size < 1:
-                        raise InvalidBenchmark(
-                            "GPU out of memory, even with a batch size of 1!"
-                        )
+        logger.debug(f"Test scores for iteration {idx}: {test_scores}")
+        scores.append(test_scores)
+        clear_memory()
 
     if not benchmark_config.debug:
         cache.remove()
@@ -198,12 +95,8 @@ def update_scores(
 
 
 def generate_single_iteration(
-    prepared_dataset: Dataset,
-    model: "GenerativeModel",
-    tokenizer: "Tokenizer",
-    data_collator: "DataCollator",
-    compute_metrics: Callable,
-    extract_labels_fn: Callable[..., list[Any]],
+    dataset: Dataset,
+    model: "BenchmarkModule",
     dataset_config: "DatasetConfig",
     benchmark_config: "BenchmarkConfig",
     cache: ModelCache,
@@ -211,18 +104,10 @@ def generate_single_iteration(
     """Evaluate a model on a dataset in a single iteration through generation.
 
     Args:
-        prepared_dataset:
+        dataset:
             The dataset to evaluate on.
         model:
             The model to evaluate.
-        tokenizer:
-            The tokenizer to use for the model.
-        data_collator:
-            The data collator to use for the model.
-        compute_metrics:
-            The function to use to compute the metrics.
-        extract_labels_fn:
-            The function to use to extract the labels from the dataset.
         dataset_config:
             The configuration of the dataset.
         benchmark_config:
@@ -235,112 +120,63 @@ def generate_single_iteration(
     """
     cache.load()
 
-    # Split up the prepared dataset into a cached and non-cached part
+    # Split up the dataset into a cached and non-cached part
     cached_dataset, non_cached_dataset = split_dataset_into_cached_and_non_cached(
-        dataset=prepared_dataset, cache=cache
+        dataset=dataset, cache=cache
     )
 
-    all_preds: list[dict | str | list[str]] = list()
+    all_preds: list[str] = list()
 
     if len(non_cached_dataset) > 0:
-        # Tokens used in generation to know when generation is finished
-        stopping_criteria = get_generation_stopping_criteria(
-            tokenizer=tokenizer, model=model
-        )
-
-        generation_config = GenerationConfig(
-            # What to output
-            max_new_tokens=dataset_config.max_generated_tokens,
-            output_scores=dataset_config.task.supertask in SUPERTASKS_USING_LOGPROBS,
-            return_dict_in_generate=True,
-            # How to sample
-            do_sample=False,  # Equivalent to greedy decoding (temperature=0)
-            # Special tokens
-            bos_token_id=tokenizer.bos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-        )
-
-        # Sort the non_cached dataset by the length of the text, to minimise the amount
-        # of padding that needs to be added, speeding up generation
-        non_cached_dataset = non_cached_dataset.add_column(
-            name="length", column=[len(x) for x in non_cached_dataset["text"]]
-        )
-        non_cached_dataset = non_cached_dataset.sort("length", reverse=True)
-
-        # Enable batching by building a dataloader. The dataloader cannot deal with
-        # text columns, so we create a copy of the dataset without these
-        torch_dataset = non_cached_dataset.with_format("torch").remove_columns(
-            [
-                column
-                for column in non_cached_dataset.column_names
-                if column != "input_ids"
-            ]
-        )
+        match model.batching_preference:
+            case BatchingPreference.SINGLE_SAMPLE:
+                itr = tqdm(iterable=non_cached_dataset, leave=False)
+            case BatchingPreference.ALL_AT_ONCE:
+                itr = [non_cached_dataset[:]]
+            case _:
+                num_batches = len(non_cached_dataset) // benchmark_config.batch_size
+                if len(non_cached_dataset) % benchmark_config.batch_size != 0:
+                    num_batches += 1
+                itr = tqdm(
+                    iterable=mit.batched(
+                        iterable=non_cached_dataset, n=benchmark_config.batch_size
+                    ),
+                    total=len(non_cached_dataset) // benchmark_config.batch_size,
+                )
 
-        if isinstance(model, OpenAIModel):
-            batch_size = 1
-        elif isinstance(model, VLLMModel):
-            batch_size = len(torch_dataset)
-        else:
-            batch_size = benchmark_config.batch_size
-
-        dataloader = DataLoader(
-            dataset=torch_dataset,
-            batch_size=batch_size,
-            shuffle=False,
-            num_workers=4,
-            collate_fn=data_collator,
-        )
+        # Generate the completions for the non-cached examples
+        for batch in itr:
+            assert isinstance(batch, dict)
 
-        with warnings.catch_warnings():
-            # This ignores the following warning, which is out of our control:
-            #   "os.fork() was called. os.fork() is incompatible with multithreaded
-            #   code, and JAX is multithreaded, so this will likely lead to a deadlock."
-            warnings.simplefilter("ignore", category=RuntimeWarning)
+            single_sample_batch = (
+                "text" in batch
+                and isinstance(batch["text"], str)
+                or "messages" in batch
+                and isinstance(batch["messages"][0], dict)
+            )
+            if single_sample_batch:
+                batch = {key: [value] for key, value in batch.items()}
 
-            itr = (
-                dataloader
-                if batch_size == len(torch_dataset) or hasattr(sys, "_called_from_test")
-                else tqdm(iterable=dataloader, leave=False)
+            model_output = model.generate(inputs=batch)
+            extracted_labels = model.extract_labels_from_generation(
+                input_batch=batch, model_output=model_output
             )
 
-            # Generate the completions for the non-cached examples
-            for batch_idx, batch in enumerate(itr):
-                model_output, extracted_labels = generate_batch(
+            # Extended logging if we are running in debug mode
+            if benchmark_config.debug:
+                debug_log(
                     batch=batch,
-                    batch_idx=batch_idx,
-                    batch_size=batch_size,
-                    non_cached_dataset=non_cached_dataset,
-                    model=model,
-                    tokenizer=tokenizer,
-                    stopping_criteria=stopping_criteria,
-                    generation_config=generation_config,
-                    extract_labels_fn=extract_labels_fn,
+                    extracted_labels=extracted_labels,  # type: ignore[arg-type]
                     dataset_config=dataset_config,
                 )
 
-                # Extended logging if we are running in debug mode
-                if benchmark_config.debug:
-                    debug_log(
-                        batch_idx=batch_idx,
-                        batch_size=batch_size,
-                        non_cached_dataset=non_cached_dataset,
-                        extracted_labels=extracted_labels,
-                        dataset_config=dataset_config,
-                    )
-
-                cache.add_to_cache(
-                    model_input=batch["input_ids"],
-                    model_output=model_output,
-                    tokenizer=tokenizer,
-                )
-                all_preds.extend(extracted_labels)
+            cache.add_to_cache(model_inputs=batch, model_output=model_output)
+            all_preds.extend(extracted_labels)
 
-                # If we are debugging then we save the cache often, but since this makes
-                # evaluation slower, we do not do this by default
-                if benchmark_config.debug:
-                    cache.save()
+            # If we are debugging then we save the cache often, but since this makes
+            # evaluation slower, we do not do this by default
+            if benchmark_config.debug:
+                cache.save()
 
         if isinstance(itr, tqdm):
             itr.close()
@@ -351,10 +187,10 @@ def generate_single_iteration(
     # Fetch the cached predictions for the cached examples
     if len(cached_dataset) > 0:
         model_output = load_cached_model_outputs(
-            cached_dataset=cached_dataset, cache=cache, tokenizer=tokenizer
+            cached_dataset=cached_dataset, cache=cache
         )
-        extracted_labels = extract_labels_fn(
-            input_batch=cached_dataset, model_output=model_output, tokenizer=tokenizer
+        extracted_labels = model.extract_labels_from_generation(
+            input_batch=cached_dataset[:], model_output=model_output
         )
         all_preds.extend(extracted_labels)
 
@@ -375,7 +211,7 @@ def generate_single_iteration(
             "The dataset must have either a 'label', 'labels', or 'target_text' column"
         )
 
-    itr_scores: dict[str, float] = compute_metrics(
+    itr_scores: dict[str, float] = model.compute_metrics(
         model_outputs_and_labels=(all_preds, ground_truth),
         id2label=dataset_config.id2label,
     )
@@ -383,330 +219,114 @@ def generate_single_iteration(
     return itr_scores
 
 
-class StopWordCriteria(StoppingCriteria):
-    """Stopping criteria for generation based on stop words.
-
-    Attributes:
-        stop_word_id_lists:
-            A list of lists of token IDs that are used to determine whether generation
-            should stop.
-        indices_done:
-            A list of indices of the examples for which generation has already stopped.
-            Resets every batch.
-    """
-
-    def __init__(self, stop_word_id_lists: list[list[int]]):
-        """Initialize the stopping criteria.
-
-        Args:
-            stop_word_id_lists:
-                A list of lists of token IDs that are used to determine whether
-                generation should stop.
-        """
-        super().__init__()
-        self.stop_word_id_lists = stop_word_id_lists
-        self.indices_done: list[int] = list()
-
-    def clear(self) -> None:
-        """Clear the example indices for which generation has already stopped."""
-        self.indices_done = list()
-
-    def __call__(
-        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
-    ) -> bool:
-        """Determine whether generation should stop.
-
-        Args:
-            input_ids:
-                The input IDs of the generated sequences.
-            scores:
-                The scores of the generated sequences. Not used.
-            **kwargs:
-                Additional keyword arguments. Not used.
-
-        Returns:
-            Whether generation should stop.
-        """
-        for stop_word_id_list in self.stop_word_id_lists:
-            for batch_idx in range(input_ids.shape[0]):
-                inputs = input_ids[batch_idx].tolist()
-                sample_ends_with_stop_word = (
-                    inputs[-len(stop_word_id_list) :] == stop_word_id_list
-                )
-                if sample_ends_with_stop_word:
-                    self.indices_done.append(batch_idx)
-                if all(idx in self.indices_done for idx in range(input_ids.shape[0])):
-                    return True
-        return False
-
-
-def generate_batch(
-    batch: dict[str, torch.Tensor],
-    batch_idx: int,
-    batch_size: int,
-    non_cached_dataset: Dataset,
-    model: "GenerativeModel",
-    tokenizer: "Tokenizer",
-    stopping_criteria: StopWordCriteria,
-    generation_config: GenerationConfig,
-    extract_labels_fn: Callable[..., list[str]],
-    dataset_config: "DatasetConfig",
-) -> tuple[ModelOutput, list[dict | str | list[str]]]:
-    """Evaluate a model on a single batch of examples through generation.
-
-    Args:
-        batch:
-            The batch of examples to evaluate on.
-        batch_idx:
-            The index of the batch.
-        batch_size:
-            The size of the batch.
-        non_cached_dataset:
-            The dataset to evaluate on.
-        model:
-            The model to evaluate.
-        tokenizer:
-            The tokenizer used to encode the examples.
-        stopping_criteria:
-            The stopping criteria to use to stop generation.
-        generation_config:
-            The generation configuration to use.
-        extract_labels_fn:
-            The function to use to extract the labels from the model output.
-        dataset_config:
-            The configuration of the dataset.
-
-    Returns:
-        The predictions generated so far, with the predictions for the current batch
-        appended.
-    """
-    # Generate the completions of the documents in the batch
-    with torch.inference_mode():
-        inputs = batch["input_ids"].to(model.device)
-        stopping_criteria.clear()
-
-        prefix_allowed_tokens_fn = None
-        logits_processors = None
-        if dataset_config.task == NER and isinstance(
-            tokenizer, PreTrainedTokenizerBase
-        ):
-            ner_tag_names = list(dataset_config.prompt_label_mapping.values())
-            prefix_allowed_tokens_fn = get_ner_prefix_allowed_tokens_fn(
-                ner_tag_names=ner_tag_names, tokenizer=tokenizer
-            )
-            if isinstance(model, VLLMModel):
-                logits_processors = get_ner_logits_processors(
-                    ner_tag_names=ner_tag_names, llm=model
-                )
-
-        model_output = model.generate(
-            inputs=inputs,
-            attention_mask=batch["attention_mask"].to(model.device)
-            if "attention_mask" in batch
-            else None,
-            generation_config=generation_config,
-            stopping_criteria=StoppingCriteriaList([stopping_criteria]),
-            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
-            logits_processors=logits_processors,
-        )
-        assert isinstance(model_output, ModelOutput)
-
-    # Some models include the input in the generated sequence, so we need to remove the
-    # input if it is present
-    inputs = inputs.detach().cpu()
-    model_output.sequences = model_output.sequences.detach().cpu()
-    if torch.equal(model_output.sequences[:, : inputs.shape[1]], inputs):
-        model_output.sequences = model_output.sequences[:, inputs.shape[1] :]
-
-    # Extract the labels from the model output and store them for metric computation
-    # later
-    batch_start = batch_idx * batch_size
-    batch_end = (batch_idx + 1) * batch_size
-    input_batch = non_cached_dataset[batch_start:batch_end]
-    extracted_labels: list = extract_labels_fn(
-        input_batch=input_batch, model_output=model_output, tokenizer=tokenizer
-    )
-
-    return model_output, extracted_labels
-
-
-def extract_raw_predictions(
-    generated_sequences: torch.Tensor, tokenizer: "Tokenizer"
-) -> list[str]:
-    """Get the raw predictions from the generated sequences.
-
-    Args:
-        generated_sequences:
-            The generated sequences from the model. The outer-most list is the
-            batch dimension, the inner-most list is the sequence dimension,
-            consisting of token IDs.
-        tokenizer:
-            The tokenizer used to generate the tokens.
-
-    Returns:
-        The candidate labels with the smallest edit distance to the predicted labels.
-    """
-    raw_predictions: list[str] = [
-        tokenizer.decode(completion_ids.tolist(), skip_special_tokens=True).split(
-            "\n\n"
-        )[0]
-        for completion_ids in generated_sequences.long()
-    ]
-
-    end_of_chat_token_ids = get_end_of_chat_token_ids(tokenizer=tokenizer)
-    if end_of_chat_token_ids is not None:
-        end_of_chat_token = tokenizer.decode(end_of_chat_token_ids).strip()
-        if end_of_chat_token:
-            raw_predictions = [
-                raw_prediction.split(end_of_chat_token)[0]
-                for raw_prediction in raw_predictions
-            ]
-
-    raw_predictions = [raw_prediction.strip() for raw_prediction in raw_predictions]
-
-    return raw_predictions
-
-
-def get_generation_stopping_criteria(
-    tokenizer: "Tokenizer", model: "GenerativeModel"
-) -> StopWordCriteria:
-    """Get the stopping criteria for generation.
-
-    Args:
-        tokenizer:
-            The tokenizer used to tokenize the stop words.
-        model:
-            The generative model, which we use to ensure the tensors are on the
-            same device, and also determine whether stop words are needed, based on
-            the model type.
-
-    Returns:
-        The stopping criteria for generation.
-    """
-    if isinstance(model, OpenAIModel):
-        return StopWordCriteria(stop_word_id_lists=[])
-
-    double_newline_ids: list[int] = tokenizer(
-        text=["\n\n"], add_special_tokens=False
-    ).input_ids[0]
-    single_newline_ids: list[int] = tokenizer(
-        text=["\n"], add_special_tokens=False
-    ).input_ids[0]
-
-    stop_word_id_lists = [double_newline_ids, single_newline_ids + single_newline_ids]
-
-    end_chat_token_ids = get_end_of_chat_token_ids(tokenizer=tokenizer)
-    if end_chat_token_ids is not None:
-        stop_word_id_lists.append(end_chat_token_ids)
-
-    if tokenizer.bos_token is not None:
-        bos_token_ids: list[int] = tokenizer(
-            text=[tokenizer.bos_token], add_special_tokens=False
-        ).input_ids[0]
-        stop_word_id_lists.append(bos_token_ids)
-
-    if tokenizer.eos_token is not None:
-        eos_token_ids: list[int] = tokenizer(
-            text=[tokenizer.eos_token], add_special_tokens=False
-        ).input_ids[0]
-        stop_word_id_lists.append(eos_token_ids)
-
-    def remove_empty_tokens(token_id_list: list[int]) -> list[int]:
-        return [
-            token_id for token_id in token_id_list if tokenizer.decode([token_id]) != ""
-        ]
-
-    stop_word_id_lists = [
-        remove_empty_tokens(token_id_list) for token_id_list in stop_word_id_lists
-    ]
-
-    return StopWordCriteria(stop_word_id_lists=stop_word_id_lists)
+# TODO: Consider removing
+# def extract_raw_predictions(generated_sequences: torch.Tensor) -> list[str]:
+#     """Get the raw predictions from the generated sequences.
+#
+#     Args:
+#         generated_sequences:
+#             The generated sequences from the model. The outer-most list is the
+#             batch dimension, the inner-most list is the sequence dimension,
+#             consisting of token IDs.
+#
+#     Returns:
+#         The candidate labels with the smallest edit distance to the predicted labels.
+#     """
+#     raw_predictions: list[str] = [
+#         tokenizer.decode(completion_ids.tolist(), skip_special_tokens=True).split(
+#             "\n\n"
+#         )[0]
+#         for completion_ids in generated_sequences.long()
+#     ]
+#
+#     end_of_chat_token_ids = get_end_of_chat_token_ids(tokenizer=tokenizer)
+#     if end_of_chat_token_ids is not None:
+#         end_of_chat_token = tokenizer.decode(end_of_chat_token_ids).strip()
+#         if end_of_chat_token:
+#             raw_predictions = [
+#                 raw_prediction.split(end_of_chat_token)[0]
+#                 for raw_prediction in raw_predictions
+#             ]
+#
+#     raw_predictions = [raw_prediction.strip() for raw_prediction in raw_predictions]
+#
+#     return raw_predictions
 
 
 def debug_log(
-    batch_idx: int,
-    batch_size: int,
-    non_cached_dataset: Dataset,
+    batch: dict[str, t.Any],
     extracted_labels: list[dict | str | list[str]],
     dataset_config: "DatasetConfig",
 ) -> None:
     """Log inputs and outputs for debugging purposes.
 
     Args:
-        batch_idx:
-            The index of the batch.
-        batch_size:
-            The size of the batch.
-        non_cached_dataset:
-            The dataset to evaluate on.
+        batch:
+            The batch of examples to evaluate on.
         extracted_labels:
             The extracted labels from the model output.
         dataset_config:
             The configuration of the dataset.
     """
-    sample_idxs = range(
-        batch_idx * batch_size,
-        min((batch_idx + 1) * batch_size, len(non_cached_dataset)),
-    )
-    samples = non_cached_dataset.select(sample_idxs)
-
-    if dataset_config.task == NER:
-        log_msgs = [""]
-        for tokens, predictions, labels in zip(
-            samples["tokens"], extracted_labels, samples["labels"]
-        ):
-            predictions = [tag.upper() for tag in predictions]
-            sample = list(zip(tokens, predictions, labels))
-            log_batches = [
-                [("Tokens: ", "Predictions: ", "Labels: ")] + sample[i : i + 10]
-                for i in range(0, len(sample), 10)
-            ]
-            for log_batch in log_batches:
-                lengths = [len(max(triple, key=len)) for triple in log_batch]
-                log_batch = [
-                    [f"{x:<{length}}" for x in triple]
-                    for triple, length in zip(log_batch, lengths)
+    match dataset_config.task.supertask:
+        case "token-classification":
+            log_msgs = [""]
+            for tokens, predictions, labels in zip(
+                batch["tokens"], extracted_labels, batch["labels"]
+            ):
+                predictions = [tag.upper() for tag in predictions]
+                sample = list(zip(tokens, predictions, labels))
+                log_batches = [
+                    [("Tokens: ", "Predictions: ", "Labels: ")] + sample[i : i + 10]
+                    for i in range(0, len(sample), 10)
                 ]
-                tokens = [triple[0] for triple in log_batch]
-                predictions = [triple[1] for triple in log_batch]
-                labels = [triple[2] for triple in log_batch]
-                log_msgs.append(
-                    "\t".join(tokens)
-                    + "\n"
-                    + "\t".join(predictions)
-                    + "\n"
-                    + "\t".join(labels)
-                )
-        logger.info("\n\n".join(log_msgs))
+                for log_batch in log_batches:
+                    lengths = [len(max(triple, key=len)) for triple in log_batch]
+                    log_batch = [
+                        [f"{x:<{length}}" for x in triple]
+                        for triple, length in zip(log_batch, lengths)
+                    ]
+                    tokens = [triple[0] for triple in log_batch]
+                    predictions = [triple[1] for triple in log_batch]
+                    labels = [triple[2] for triple in log_batch]
+                    log_msgs.append(
+                        "\t".join(tokens)
+                        + "\n"
+                        + "\t".join(predictions)
+                        + "\n"
+                        + "\t".join(labels)
+                    )
+            logger.info("\n\n".join(log_msgs))
+            return
 
-    else:
-        # Define predictions
-        if dataset_config.task == RC:
+        case "sequence-classification":
+            labels = [
+                dataset_config.prompt_label_mapping.get(label, label).lower()
+                for label in batch["label"]
+            ]
+
+        case "question-answering":
             extracted_labels = [
                 prediction["prediction_text"]
                 for prediction in extracted_labels
                 if isinstance(prediction, dict)
             ]
+            labels = [label["answers"]["text"][0] for label in batch["label"]]
 
-        # Define labels
-        if dataset_config.task in {SENT, LA}:
-            labels = [
-                dataset_config.prompt_label_mapping.get(label, label).lower()
-                for label in samples["label"]
-            ]
-        elif dataset_config.task == RC:
-            labels = [label["answers"]["text"][0] for label in samples["label"]]
-        elif dataset_config.task == SUMM:
-            labels = samples["target_text"]
-        else:
-            labels = samples["label"]
-
-        # Log inputs and outputs
-        for input_text, prediction, label in zip(
-            samples["text"], extracted_labels, labels
-        ):
-            logger.info(
-                f"Input: '{input_text}'\n"
-                f"Prediction: '{prediction}'\n"
-                f"Label: '{label}'"
+        case "text-to-text":
+            labels = batch["target_text"]
+
+        case _:
+            raise InvalidBenchmark(
+                f"The supertask '{dataset_config.task.supertask}' is not supported."
             )
+
+    if "messages" in batch:
+        input_texts = [messages[-1]["content"] for messages in batch["messages"]]
+    else:
+        input_texts = batch["text"]
+
+    for input_text, prediction, label in zip(input_texts, extracted_labels, labels):
+        logger.info(
+            f"Input: {input_text!r}\nPrediction: {prediction!r}\nLabel: {label!r}"
+        )
diff --git a/src/scandeval/human_evaluation.py b/src/scandeval/human_evaluation.py
index e2737457f..664ac69b5 100644
--- a/src/scandeval/human_evaluation.py
+++ b/src/scandeval/human_evaluation.py
@@ -3,32 +3,37 @@
 import importlib.util
 import json
 import logging
+from collections import defaultdict
 from functools import partial
 from pathlib import Path
 
 import click
 from datasets import Dataset
-from transformers import AutoConfig, AutoTokenizer
-from transformers.utils import ModelOutput
+from transformers import AutoTokenizer
 
 from .benchmark_config_factory import build_benchmark_config
-from .benchmark_dataset import BenchmarkDataset
 from .benchmarker import BenchmarkResult
-from .config import ModelConfig
+from .data_loading import load_data
+from .data_models import GenerativeModelOutput
 from .dataset_configs import SPEED_CONFIG, get_all_dataset_configs
-from .dataset_factory import DatasetFactory
-from .enums import Framework, ModelType
+from .enums import Framework
 from .exceptions import NeedsExtraInstalled
 from .scores import aggregate_scores
+from .task_utils import (
+    question_answering,
+    sequence_classification,
+    text_to_text,
+    token_classification,
+)
 from .tasks import NER
-from .types import ScoreDict
-from .utils import create_model_cache_dir, enforce_reproducibility
+from .types import ComputeMetricsFunction, ExtractLabelsFunction, ScoreDict
+from .utils import enforce_reproducibility
 
 if importlib.util.find_spec("gradio") is not None:
     import gradio as gr
 
 
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 class HumanEvaluator:
@@ -59,7 +64,6 @@ def __init__(
         self.dummy_model_id = dummy_model_id
 
         self.sample_idx: int
-        self.benchmark_dataset: BenchmarkDataset
         self.active_dataset: Dataset
 
         self.dataset_configs = {
@@ -84,6 +88,9 @@ def __init__(
             }
         )
 
+        self.extract_labels_from_generation: ExtractLabelsFunction
+        self.compute_metrics: ComputeMetricsFunction
+
     def create_app(self) -> "gr.Blocks":
         """Create the Gradio app for human evaluation.
 
@@ -170,7 +177,9 @@ def create_app(self) -> "gr.Blocks":
             )
         return app
 
-    def update_dataset_choices(self, language: str, task: str) -> "gr.Dropdown":
+    def update_dataset_choices(
+        self, language: str | None, task: str | None
+    ) -> "gr.Dropdown":
         """Update the dataset choices based on the selected language and task.
 
         Args:
@@ -251,45 +260,38 @@ def update_dataset(
             framework=None,
             device=None,
             batch_size=1,
-            evaluate_train=False,
             raise_errors=False,
             cache_dir=".scandeval_cache",
-            token=None,
-            openai_api_key=None,
-            prefer_azure=False,
-            azure_openai_api_key=None,
-            azure_openai_endpoint=None,
-            azure_openai_api_version=None,
+            api_key=None,
             force=False,
             verbose=False,
             trust_remote_code=False,
             load_in_4bit=None,
             use_flash_attention=None,
             clear_model_cache=False,
-            only_validation_split=True,
+            evaluate_test_split=False,
             few_shot=True,
             num_iterations=iteration + 1,
             debug=False,
             run_with_cli=True,
         )
-        dataset_factory = DatasetFactory(benchmark_config=benchmark_config)
-        dataset_config = get_all_dataset_configs()[dataset_name]
-
-        model_id = f"human-{iteration}"
-        model_config = ModelConfig(
-            model_id=model_id,
-            revision="main",
-            framework=Framework.HUMAN,
-            task="text-generation",
-            languages=dataset_config.languages,
-            model_type=ModelType.HUMAN,
-            model_cache_dir=create_model_cache_dir(
-                cache_dir=benchmark_config.cache_dir, model_id=model_id
-            ),
-            adapter_base_model_id=None,
-        )
+        self.dataset_config = get_all_dataset_configs()[dataset_name]
+
+        # TODO: Is this needed?
+        # model_id = f"human-{iteration}"
+        # model_config = ModelConfig(
+        #     model_id=model_id,
+        #     revision="main",
+        #     framework=Framework.HUMAN,
+        #     task="text-generation",
+        #     languages=dataset_config.languages,
+        #     model_type=ModelType.HUMAN,
+        #     model_cache_dir=create_model_cache_dir(
+        #         cache_dir=benchmark_config.cache_dir, model_id=model_id
+        #     ),
+        #     adapter_base_model_id=None,
+        # )
 
-        self.benchmark_dataset = dataset_factory.build_dataset(dataset=dataset_config)
         self.sample_idx = 0
 
         dataset_path = (
@@ -299,7 +301,9 @@ def update_dataset(
             / f"human-{iteration}.csv"
         )
         if dataset_path.exists():
-            self.active_dataset = Dataset.from_csv(str(dataset_path))
+            active_dataset = Dataset.from_csv(str(dataset_path))
+            assert isinstance(active_dataset, Dataset)
+            self.active_dataset = active_dataset
             try:
                 while self.active_dataset["answer"][self.sample_idx] is not None:
                     self.sample_idx += 1
@@ -308,21 +312,79 @@ def update_dataset(
                 return blank_answer
         else:
             rng = enforce_reproducibility(framework=Framework.PYTORCH)
-            train, val, tests = self.benchmark_dataset._load_data(rng=rng)
-            _, _, tests = self.benchmark_dataset._load_prepared_data(
-                train=train,
-                val=val,
-                tests=tests,
-                model_config=model_config,
-                hf_model_config=AutoConfig.from_pretrained(self.dummy_model_id),
-                tokenizer=AutoTokenizer.from_pretrained(self.dummy_model_id),
-                benchmarking_generative_model=True,
+            datasets = load_data(
+                rng=rng,
+                dataset_config=self.dataset_config,
+                benchmark_config=benchmark_config,
             )
+            # TODO: Prepare data?
             self.active_dataset = (
-                tests[iteration]
-                .remove_columns(column_names=["input_ids", "attention_mask"])
-                .add_column(name="answer", column=[None] * len(tests[iteration]))
+                datasets[iteration]["test"]
+                .remove_columns(
+                    column_names=["input_ids", "attention_mask"],
+                    new_fingerprint=datasets[iteration]["test"]._fingerprint,
+                )
+                .add_column(
+                    name="answer",
+                    column=[None] * len(datasets[iteration]["test"]),
+                    new_fingerprint=datasets[iteration]["test"]._fingerprint,
+                )
             )
+            if self.dataset_config.task == NER:
+                labels_in_train: set[str] = {
+                    tag
+                    for tag_list in self.active_dataset["labels"]
+                    for tag in tag_list
+                }
+                self.has_misc_tags = (
+                    "B-MISC" in labels_in_train or "I-MISC" in labels_in_train
+                )
+
+        match self.dataset_config.task.supertask:
+            case "sequence-classification":
+                self.compute_metrics = partial(
+                    sequence_classification.compute_metrics,
+                    id2label=self.dataset_config.id2label,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=benchmark_config,
+                )
+                self.extract_labels_from_generation = partial(
+                    sequence_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "text-to-text":
+                self.compute_metrics = partial(
+                    text_to_text.compute_metrics,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=benchmark_config,
+                )
+                self.extract_labels_from_generation = (
+                    text_to_text.extract_labels_from_generation
+                )
+            case "token-classification":
+                self.compute_metrics = partial(
+                    token_classification.compute_metrics,
+                    has_misc_tags=self.has_misc_tags,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=benchmark_config,
+                )
+                self.extract_labels_from_generation = partial(
+                    token_classification.extract_labels_from_generation,
+                    dataset_config=self.dataset_config,
+                )
+            case "question-answering":
+                self.compute_metrics = partial(
+                    question_answering.compute_metrics,
+                    dataset_config=self.dataset_config,
+                    benchmark_config=benchmark_config,
+                )
+                self.extract_labels_from_generation = (
+                    question_answering.extract_labels_from_generation
+                )
+            case _:
+                raise NotImplementedError(
+                    f"Supertask {self.dataset_config.task.supertask} is not supported."
+                )
 
         task_examples, question = self.example_to_markdown(
             example=self.active_dataset[self.sample_idx]
@@ -333,9 +395,9 @@ def update_dataset(
             f"{task_examples}"
         )
 
-        if self.benchmark_dataset.dataset_config.task == NER:
+        if self.dataset_config.task == NER:
             ner_tags = list()
-            for ner_tag in dataset_config.prompt_label_mapping.values():
+            for ner_tag in self.dataset_config.prompt_label_mapping.values():
                 if ner_tag not in ner_tags:
                     ner_tags.append(ner_tag)
             return (
@@ -411,9 +473,7 @@ def reset_entities(self) -> "gr.Textbox":
             A blank answer.
         """
         ner_tags = list()
-        for (
-            ner_tag
-        ) in self.benchmark_dataset.dataset_config.prompt_label_mapping.values():
+        for ner_tag in self.dataset_config.prompt_label_mapping.values():
             if ner_tag not in ner_tags:
                 ner_tags.append(ner_tag)
         return gr.Textbox(json.dumps({ner_tag: [] for ner_tag in ner_tags}))
@@ -443,7 +503,7 @@ def submit_answer(
             return question, answer
 
         # Custom NER validation
-        if self.benchmark_dataset.dataset_config.task == NER:
+        if self.dataset_config.task == NER:
             try:
                 json.loads(answer)
             except json.JSONDecodeError:
@@ -460,9 +520,7 @@ def submit_answer(
                 )
                 return question, answer
 
-            ner_tags = list(
-                self.benchmark_dataset.dataset_config.prompt_label_mapping.values()
-            )
+            ner_tags = list(self.dataset_config.prompt_label_mapping.values())
             for ner_tag in ner_tags:
                 if ner_tag not in json.loads(answer).keys():
                     gr.Warning(
@@ -481,8 +539,12 @@ def submit_answer(
         # Store the user's answer
         answers = self.active_dataset["answer"]
         answers[self.sample_idx] = answer
-        self.active_dataset = self.active_dataset.remove_columns("answer").add_column(
-            name="answer", column=answers
+        self.active_dataset = self.active_dataset.remove_columns(
+            column_names=["answer"], new_fingerprint=self.active_dataset._fingerprint
+        ).add_column(
+            name="answer",
+            column=answers,
+            new_fingerprint=self.active_dataset._fingerprint,
         )
         logger.info(
             f"User submitted the answer {answer!r} to the question {question!r}, with "
@@ -505,11 +567,9 @@ def submit_answer(
                 example=self.active_dataset[self.sample_idx]
             )
 
-            if self.benchmark_dataset.dataset_config.task == NER:
+            if self.dataset_config.task == NER:
                 ner_tags = list()
-                for ner_tag in (
-                    self.benchmark_dataset.dataset_config.prompt_label_mapping.values()
-                ):
+                for ner_tag in self.dataset_config.prompt_label_mapping.values():
                     if ner_tag not in ner_tags:
                         ner_tags.append(ner_tag)
                 answer = json.dumps({ner_tag: [] for ner_tag in ner_tags})
@@ -557,16 +617,18 @@ def compute_and_log_scores(self) -> None:
             padding=True,
             return_tensors="pt",
         ).input_ids
-        model_output = ModelOutput(sequences=sequences)
-        all_preds = self.benchmark_dataset._extract_labels_from_generation(
-            input_batch=self.active_dataset.to_dict(),
-            model_output=model_output,
-            tokenizer=tokenizer,
+        model_output = GenerativeModelOutput(sequences=sequences)
+
+        active_dataset_dict = self.active_dataset.to_dict()
+        assert isinstance(active_dataset_dict, dict)
+
+        all_preds = self.extract_labels_from_generation(
+            input_batch=active_dataset_dict, model_output=model_output
         )
         ground_truth = self.active_dataset["label"]
-        itr_scores: dict[str, float] = self.benchmark_dataset._compute_metrics(
+        itr_scores: dict[str, float] = self.compute_metrics(
             model_outputs_and_labels=(all_preds, ground_truth),
-            id2label=self.benchmark_dataset.dataset_config.id2label,
+            id2label=self.dataset_config.id2label,
         )
 
         # We reverse the order, as the Info messages are printed in reverse order
@@ -580,7 +642,7 @@ def compute_and_log_scores(self) -> None:
             gr.Info(f"\n\n{metric_name}: {score:.2%}")
         gr.Info("You have completed this dataset! Here are your scores:")
         logger.info(
-            f"User completed the dataset {self.benchmark_dataset.dataset_config.name!r}"
+            f"User completed the dataset {self.dataset_config.name!r}"
             f", with the following scores: {itr_scores}"
         )
 
@@ -588,7 +650,7 @@ def compute_and_log_scores(self) -> None:
         # only a single iteration, so the results from the current annotation should be
         # added to the previous results.
         results_path = Path.cwd() / "scandeval_benchmark_results.jsonl"
-        results: ScoreDict = dict(raw=dict(test=list()))  # type: ignore[dict-item]
+        results: ScoreDict = defaultdict(list)
         if results_path.exists():
             all_results = [
                 json.loads(line.strip())
@@ -599,24 +661,23 @@ def compute_and_log_scores(self) -> None:
                 result
                 for result in all_results
                 if result["model"] == "human"
-                and result["dataset"] == self.benchmark_dataset.dataset_config.name
+                and result["dataset"] == self.dataset_config.name
             ]
             if human_result_candidates:
                 results = human_result_candidates[0]["results"]
 
         # Append to results
-        results["raw"]["test"].append(  # type: ignore[union-attr]
+        results["raw"].append(  # type: ignore[union-attr]
             {f"test_{metric_name}": score for metric_name, score in itr_scores.items()}
         )
 
         # Aggregate scores
         total_dict: dict[str, float] = dict()
-        for metric_cfg in self.benchmark_dataset.dataset_config.task.metrics:
-            agg_scores = aggregate_scores(
+        for metric_cfg in self.dataset_config.task.metrics:
+            test_score, test_se = aggregate_scores(
                 scores=results["raw"],  # type: ignore[arg-type]
                 metric_config=metric_cfg,
             )
-            test_score, test_se = agg_scores["test"]
             test_score, _ = metric_cfg.postprocessing_fn(test_score)
             test_se, _ = metric_cfg.postprocessing_fn(test_se)
             total_dict[f"test_{metric_cfg.name}"] = test_score
@@ -624,11 +685,10 @@ def compute_and_log_scores(self) -> None:
         results["total"] = total_dict
 
         benchmark_result = BenchmarkResult(
-            dataset=self.benchmark_dataset.dataset_config.name,
-            task=self.benchmark_dataset.dataset_config.task.name,
+            dataset=self.dataset_config.name,
+            task=self.dataset_config.task.name,
             dataset_languages=[
-                language.code
-                for language in self.benchmark_dataset.dataset_config.languages
+                language.code for language in self.dataset_config.languages
             ],
             model="human",
             results=results,
diff --git a/src/scandeval/languages.py b/src/scandeval/languages.py
index 3ef50cb03..4f680121c 100644
--- a/src/scandeval/languages.py
+++ b/src/scandeval/languages.py
@@ -5,7 +5,7 @@
 Last updated 19 June 2022.
 """
 
-from .config import Language
+from .data_models import Language
 
 
 def get_all_languages() -> dict[str, Language]:
diff --git a/src/scandeval/model_cache.py b/src/scandeval/model_cache.py
index 4471c7cba..bc855f2c8 100644
--- a/src/scandeval/model_cache.py
+++ b/src/scandeval/model_cache.py
@@ -1,37 +1,25 @@
 """ModelCache class for caching model outputs."""
 
+import hashlib
 import json
 import logging
-import math
 import sys
+import typing as t
 from collections import defaultdict
-from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING
+from dataclasses import asdict
 
 import pandas as pd
-import torch
 from tqdm.auto import tqdm
-from transformers.modeling_utils import ModelOutput
 
-if TYPE_CHECKING:
+from .data_models import GenerativeModelOutput, SingleGenerativeModelOutput
+
+if t.TYPE_CHECKING:
     from pathlib import Path
 
     from datasets import Dataset
 
-    from .protocols import Tokenizer
-
-
-logger = logging.getLogger(__package__)
 
-
-@dataclass
-class GenerativeModelOutput:
-    """The output of a generative model."""
-
-    completion: str
-    top_score_indices: list[list[int]] | None = None
-    top_score_values: list[list[float]] | None = None
-    vocab_size: int | None = None
+logger = logging.getLogger("scandeval")
 
 
 class ModelCache:
@@ -75,9 +63,9 @@ def load(self) -> None:
         with self.cache_path.open() as f:
             json_cache = json.load(f)
 
-        cache: dict[str, GenerativeModelOutput] = dict()
+        cache: dict[str, SingleGenerativeModelOutput] = dict()
         for key in json_cache:
-            cache[key] = GenerativeModelOutput(**json_cache[key])
+            cache[key] = SingleGenerativeModelOutput(**json_cache[key])
 
         self.cache = cache
 
@@ -90,7 +78,21 @@ def save(self) -> None:
         with self.cache_path.open("w") as f:
             json.dump(dumpable_cache, f)
 
-    def __getitem__(self, key: str) -> GenerativeModelOutput:
+    def _hash_key(self, key: str | list[dict[str, str]]) -> str:
+        """Hash the key to use as an index in the cache.
+
+        Args:
+            key:
+                The key to hash.
+
+        Returns:
+            The hashed key.
+        """
+        return hashlib.md5(string=str(key).encode()).hexdigest()
+
+    def __getitem__(
+        self, key: str | list[dict[str, str]]
+    ) -> SingleGenerativeModelOutput:
         """Get an item from the cache.
 
         Args:
@@ -100,9 +102,10 @@ def __getitem__(self, key: str) -> GenerativeModelOutput:
         Returns:
             The model output.
         """
-        return self.cache[key]
+        hashed_key = self._hash_key(key=key)
+        return self.cache[hashed_key]
 
-    def __setitem__(self, key: str, value: GenerativeModelOutput) -> None:
+    def __setitem__(self, key: t.Any, value: SingleGenerativeModelOutput) -> None:
         """Set an item in the cache.
 
         Args:
@@ -111,7 +114,8 @@ def __setitem__(self, key: str, value: GenerativeModelOutput) -> None:
             value:
                 The value to set in the cache.
         """
-        self.cache[key] = value
+        hashed_key = self._hash_key(key=key)
+        self.cache[hashed_key] = value
 
     def remove(self) -> None:
         """Remove the cache from memory and delete it from disk."""
@@ -123,68 +127,39 @@ def cached_texts(self) -> list[str]:
         return [key for key in self.cache.keys()]
 
     def add_to_cache(
-        self,
-        model_input: torch.Tensor,
-        model_output: ModelOutput,
-        tokenizer: "Tokenizer",
+        self, model_inputs: dict, model_output: GenerativeModelOutput
     ) -> None:
         """Add the model input/output to the cache.
 
         Args:
-            model_input:
-                The model input.
+            model_inputs:
+                The model inputs.
             model_output:
                 The model output.
-            tokenizer:
-                The tokenizer used to generate the tokens.
         """
-        model_input = model_input.detach().cpu()
-
-        # Extract the scores from the model output, to be cached. We only store the
-        # indices of the top scores, to save space. Further, we only store the scores
-        # if the generated sequence is shorter than the maximum length
-        store_scores = "scores" in model_output and self.max_generated_tokens < 8
-        if store_scores:
-            scores = torch.stack(
-                tensors=[
-                    score_tensor.detach().cpu().float()
-                    for score_tensor in model_output.scores
-                ],
-                dim=1,
-            )
-            top_scores = torch.topk(scores, k=10)
+        input_column = "messages" if "messages" in model_inputs else "text"
+        model_inputs = model_inputs[input_column]
 
         # Store the generated sequences in the cache, one by one
-        # TODO: This is a bit slow, should be optimized
         with tqdm(
-            iterable=model_input,
+            iterable=model_inputs,
             desc="Caching model outputs",
             leave=False,
             disable=hasattr(sys, "_called_from_test"),
         ) as pbar:
-            for sample_idx, sample in enumerate(pbar):
-                decoded_inputs = tokenizer.decode(
-                    token_ids=sample, skip_special_tokens=True
-                )
-                generated_ids = model_output.sequences[sample_idx].tolist()
-
-                # Set up the model output in a GenerativeModelOutput object
-                cached_model_output = GenerativeModelOutput(
-                    completion=tokenizer.decode(
-                        token_ids=generated_ids, skip_special_tokens=True
-                    )
+            for sample_idx, model_input in enumerate(pbar):
+                # Extract the scores from the model output, to be cached. We only store
+                # the indices of the top scores, to save space. Further, we only store
+                # the scores if the generated sequence is shorter than the maximum
+                # length
+                if model_output.scores is not None and self.max_generated_tokens < 8:
+                    assert model_output.scores is not None
+                    scores = model_output.scores[sample_idx]
+                else:
+                    scores = None
+                self[model_input] = SingleGenerativeModelOutput(
+                    sequence=model_output.sequences[sample_idx], scores=scores
                 )
-                if store_scores:
-                    cached_model_output.top_score_indices = top_scores.indices[
-                        sample_idx
-                    ].tolist()
-                    cached_model_output.top_score_values = top_scores.values[
-                        sample_idx
-                    ].tolist()
-                    cached_model_output.vocab_size = int(scores.shape[-1])
-
-                # Store the generated sequence in the cache
-                self[decoded_inputs] = cached_model_output
 
 
 def split_dataset_into_cached_and_non_cached(
@@ -203,7 +178,8 @@ def split_dataset_into_cached_and_non_cached(
     """
     # Get the sample indices of the non-cached examples, which are unique with respect
     # to the "text" column.
-    dataset_texts = pd.Series(dataset["text"])
+    input_column = "messages" if "messages" in dataset.column_names else "text"
+    dataset_texts = pd.Series(dataset[input_column])
     dataset_texts.drop_duplicates(inplace=True)
     unique_non_cached_ids = set(
         dataset_texts[~dataset_texts.isin(cache.cached_texts())].index.tolist()
@@ -220,8 +196,8 @@ def split_dataset_into_cached_and_non_cached(
 
 
 def load_cached_model_outputs(
-    cached_dataset: "Dataset", cache: ModelCache, tokenizer: "Tokenizer"
-) -> ModelOutput:
+    cached_dataset: "Dataset", cache: ModelCache
+) -> GenerativeModelOutput:
     """Load the cached model outputs.
 
     Args:
@@ -229,77 +205,19 @@ def load_cached_model_outputs(
             The dataset containing the cached examples.
         cache:
             The model output cache.
-        tokenizer:
-            The tokenizer used to generate the tokens.
 
     Returns:
         The model output containing the cached sequences.
     """
-    # Load the raw model outputs from the cache
-    cached_model_outputs: list[GenerativeModelOutput] = [
-        cache[prompt] for prompt in cached_dataset["text"]
+    input_column = "messages" if "messages" in cached_dataset.column_names else "text"
+    cached_model_outputs: list[SingleGenerativeModelOutput] = [
+        cache[prompt] for prompt in cached_dataset[input_column]
     ]
 
-    # Tokenize the cached sequences
-    tokenized_cached_sequences: list[torch.Tensor] = [
-        tokenizer(
-            text=cached_model_output.completion,
-            add_special_tokens=False,
-            return_tensors="pt",
-        ).input_ids.squeeze(dim=0)
-        for cached_model_output in cached_model_outputs
-    ]
+    cached_sequences = [model_output.sequence for model_output in cached_model_outputs]
 
-    # Pad the cached completions to the same length
-    cached_sequences = torch.nn.utils.rnn.pad_sequence(
-        sequences=tokenized_cached_sequences,
-        batch_first=True,
-        padding_value=tokenizer.pad_token_id,
-    )
+    if cached_model_outputs[0].scores is None:
+        return GenerativeModelOutput(sequences=cached_sequences)
 
-    # If we do not have any cached scores, then wrap the padded cached sequences in a
-    # ModelOutput and return it
-    if (
-        cached_model_outputs[0].top_score_indices is None
-        or cached_model_outputs[0].top_score_values is None
-        or cached_model_outputs[0].vocab_size is None
-    ):
-        return ModelOutput(sequences=cached_sequences)
-
-    # Otherwise, we format the cached scores into a tensor of shape [batch_size,
-    # num_sequences, vocab_size], wrap it in a ModelOutput with the padded cached
-    # sequences, and return it
-    cached_scores = torch.full(
-        size=(
-            len(cached_model_outputs),
-            max(
-                len(cached_model_output.top_score_indices)
-                for cached_model_output in cached_model_outputs
-                if cached_model_output.top_score_indices is not None
-            ),
-            cached_model_outputs[0].vocab_size,
-        ),
-        fill_value=-math.inf,
-    )
-    top_score_indices = torch.nn.utils.rnn.pad_sequence(
-        sequences=[
-            torch.tensor(cached_model_output.top_score_indices)
-            for cached_model_output in cached_model_outputs
-        ],
-        batch_first=True,
-        padding_value=tokenizer.pad_token_id,
-    )
-    top_score_values = torch.nn.utils.rnn.pad_sequence(
-        sequences=[
-            torch.tensor(cached_model_output.top_score_values)
-            for cached_model_output in cached_model_outputs
-        ],
-        batch_first=True,
-        padding_value=-math.inf,
-    )
-    for batch_idx in range(cached_scores.shape[0]):
-        for sequence_idx in range(cached_scores.shape[1]):
-            top_indices = top_score_indices[batch_idx, sequence_idx]
-            top_values = top_score_values[batch_idx, sequence_idx]
-            cached_scores[batch_idx, sequence_idx, top_indices] = top_values
-    return ModelOutput(sequences=cached_sequences, scores=cached_scores)
+    cached_scores = [model_output.scores or [] for model_output in cached_model_outputs]
+    return GenerativeModelOutput(sequences=cached_sequences, scores=cached_scores)
diff --git a/src/scandeval/model_config.py b/src/scandeval/model_config.py
index 16fe4ac88..01970c5ce 100644
--- a/src/scandeval/model_config.py
+++ b/src/scandeval/model_config.py
@@ -1,14 +1,14 @@
 """Functions related to getting the model configuration."""
 
 import importlib.util
-from typing import TYPE_CHECKING
+import typing as t
 
+from . import benchmark_modules
 from .enums import Framework
-from .exceptions import InvalidModel, NeedsExtraInstalled
-from .model_setups import MODEL_SETUP_CLASSES
+from .exceptions import InvalidModel, NeedsEnvironmentVariable, NeedsExtraInstalled
 
-if TYPE_CHECKING:
-    from .config import BenchmarkConfig, ModelConfig
+if t.TYPE_CHECKING:
+    from .data_models import BenchmarkConfig, ModelConfig
 
 
 def get_model_config(
@@ -29,24 +29,28 @@ def get_model_config(
         InvalidModel:
             If all model setups can handle the model, but the model does not exist.
     """
+    all_benchmark_modules = [
+        cls
+        for cls in benchmark_modules.__dict__.values()
+        if isinstance(cls, type)
+        and issubclass(cls, benchmark_modules.BenchmarkModule)
+        and cls is not benchmark_modules.BenchmarkModule
+    ]
+
     needs_extras: list[str] = list()
     needs_env_vars: list[str] = list()
-    for setup_class in MODEL_SETUP_CLASSES:
-        setup = setup_class(benchmark_config=benchmark_config)
-
-        exists_or_dict = setup.model_exists(model_id=model_id)
-        if isinstance(exists_or_dict, dict):
-            if "missing_extra" in exists_or_dict:
-                needs_extras.append(exists_or_dict["missing_extra"])
-            elif "missing_env_var" in exists_or_dict:
-                needs_env_vars.append(exists_or_dict["missing_env_var"])
-            else:
-                raise ValueError(
-                    "The dictionary returned by `model_exists` must contain either "
-                    "the key `missing_extra` or `missing_env_var`."
-                )
-        elif exists_or_dict:
-            model_config = setup.get_model_config(model_id=model_id)
+    for benchmark_module in all_benchmark_modules:
+        exists_or_err = benchmark_module.model_exists(
+            model_id=model_id, benchmark_config=benchmark_config
+        )
+        if isinstance(exists_or_err, NeedsExtraInstalled):
+            needs_extras.append(exists_or_err.extra)
+        elif isinstance(exists_or_err, NeedsEnvironmentVariable):
+            needs_env_vars.append(exists_or_err.env_var)
+        elif exists_or_err is True:
+            model_config = benchmark_module.get_model_config(
+                model_id=model_id, benchmark_config=benchmark_config
+            )
             if (
                 model_config.framework == Framework.JAX
                 and importlib.util.find_spec("jax") is None
diff --git a/src/scandeval/model_loading.py b/src/scandeval/model_loading.py
index 9242ccd19..d60f3c5fc 100644
--- a/src/scandeval/model_loading.py
+++ b/src/scandeval/model_loading.py
@@ -1,32 +1,27 @@
 """Functions related to the loading of models."""
 
-from typing import TYPE_CHECKING, Type
+import typing as t
 
-from .exceptions import InvalidBenchmark
-from .model_setups import (
-    FreshModelSetup,
-    HFModelSetup,
-    LocalModelSetup,
-    OpenAIModelSetup,
-)
-from .utils import (
-    GENERATIVE_DATASET_SUPERTASKS,
-    GENERATIVE_DATASET_TASKS,
-    model_is_generative,
+from .benchmark_modules import (
+    FreshEncoderModel,
+    HuggingFaceEncoderModel,
+    LiteLLMModel,
+    VLLMModel,
 )
+from .constants import GENERATIVE_DATASET_SUPERTASKS, GENERATIVE_DATASET_TASKS
+from .enums import ModelType
+from .exceptions import InvalidBenchmark
 
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-    from .protocols import GenerativeModel, ModelSetup, Tokenizer
+if t.TYPE_CHECKING:
+    from .benchmark_modules import BenchmarkModule
+    from .data_models import BenchmarkConfig, DatasetConfig, ModelConfig
 
 
 def load_model(
     model_config: "ModelConfig",
     dataset_config: "DatasetConfig",
     benchmark_config: "BenchmarkConfig",
-) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
+) -> "BenchmarkModule":
     """Load a model.
 
     Args:
@@ -38,19 +33,23 @@ def load_model(
             The benchmark configuration.
 
     Returns:
-        The tokenizer and model.
+        The model.
     """
-    model_type_to_model_setup_mapping: dict[str, Type["ModelSetup"]] = dict(
-        fresh=FreshModelSetup,
-        hf=HFModelSetup,
-        local=LocalModelSetup,
-        openai=OpenAIModelSetup,
-    )
-    setup_class = model_type_to_model_setup_mapping[model_config.model_type]
-    setup = setup_class(benchmark_config=benchmark_config)
+    # The order matters; the first model type that matches will be used. For this
+    # reason, they have been ordered in terms of the most common model types.
+    model_type_to_module_mapping: dict[ModelType, t.Type[BenchmarkModule]] = {
+        ModelType.HF_HUB_GENERATIVE: VLLMModel,
+        ModelType.HF_HUB_ENCODER: HuggingFaceEncoderModel,
+        ModelType.API: LiteLLMModel,
+        ModelType.LOCAL: VLLMModel,  # TODO: Add LocalModel
+        ModelType.FRESH: FreshEncoderModel,
+    }
+    model_class = model_type_to_module_mapping[model_config.model_type]
 
-    model, tokenizer = setup.load_model(
-        model_config=model_config, dataset_config=dataset_config
+    model = model_class(
+        model_config=model_config,
+        dataset_config=dataset_config,
+        benchmark_config=benchmark_config,
     )
 
     # Refuse to benchmark non-generative models on generative tasks
@@ -60,13 +59,11 @@ def load_model(
             or dataset_config.task.name in GENERATIVE_DATASET_TASKS
         )
         and model is not None
-        and not model_is_generative(model=model)
+        and not model.is_generative
     ):
         raise InvalidBenchmark(
             f"Cannot benchmark non-generative model {model_config.model_id!r} on "
             f"generative task {dataset_config.task.name!r}."
         )
 
-    # TODO: XMOD model setup: https://huggingface.co/facebook/xmod-base#input-language
-
-    return model, tokenizer
+    return model
diff --git a/src/scandeval/model_setups/__init__.py b/src/scandeval/model_setups/__init__.py
deleted file mode 100644
index a525fa85b..000000000
--- a/src/scandeval/model_setups/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-"""The different types of models that can be benchmarked."""
-
-from typing import TYPE_CHECKING, Type
-
-from .fresh import FreshModelSetup
-from .hf import HFModelSetup
-from .local import LocalModelSetup
-from .openai import OpenAIModelSetup
-
-if TYPE_CHECKING:
-    from ..protocols import ModelSetup
-
-
-# Note that the order of the model setup classes is important, as the first model setup
-# that can load a model will be used
-MODEL_SETUP_CLASSES: list[Type["ModelSetup"]] = [
-    FreshModelSetup,
-    LocalModelSetup,
-    HFModelSetup,
-    OpenAIModelSetup,
-]
diff --git a/src/scandeval/model_setups/fresh.py b/src/scandeval/model_setups/fresh.py
deleted file mode 100644
index c49efb17a..000000000
--- a/src/scandeval/model_setups/fresh.py
+++ /dev/null
@@ -1,177 +0,0 @@
-"""Model setup for fresh models."""
-
-import re
-from json import JSONDecodeError
-from typing import TYPE_CHECKING
-
-from transformers import (
-    AutoConfig,
-    AutoTokenizer,
-    ElectraForQuestionAnswering,
-    ElectraForSequenceClassification,
-    ElectraForTokenClassification,
-    XLMRobertaForQuestionAnswering,
-    XLMRobertaForSequenceClassification,
-    XLMRobertaForTokenClassification,
-)
-
-from ..config import ModelConfig
-from ..enums import Framework, ModelType
-from ..exceptions import InvalidBenchmark, InvalidModel
-from ..utils import block_terminal_output, create_model_cache_dir, model_is_generative
-from .utils import align_model_and_tokenizer, setup_model_for_question_answering
-
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerBase
-
-    from ..config import BenchmarkConfig, DatasetConfig
-    from ..protocols import GenerativeModel, Tokenizer
-
-
-FRESH_MODELS: list[str] = ["electra-small", "xlm-roberta-base"]
-
-
-class FreshModelSetup:
-    """Model setup for fresh models.
-
-    Attributes:
-        benchmark_config:
-            The benchmark configuration.
-    """
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the FreshModelSetup class.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration.
-        """
-        self.benchmark_config = benchmark_config
-
-    @staticmethod
-    def _strip_model_id(model_id: str) -> str:
-        return re.sub("(@.*$|^fresh-)", "", model_id)
-
-    def model_exists(self, model_id: str) -> bool | dict[str, str]:
-        """Check if a model ID denotes a fresh model.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            Whether the model exist, or a dictionary explaining why we cannot check
-            whether the model exists.
-        """
-        return self._strip_model_id(model_id=model_id) in FRESH_MODELS
-
-    def get_model_config(self, model_id: str) -> ModelConfig:
-        """Fetches configuration for a fresh model.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            The model configuration.
-        """
-        return ModelConfig(
-            model_id=self._strip_model_id(model_id=model_id),
-            framework=Framework.PYTORCH,
-            task="fill-mask",
-            languages=list(),
-            revision="main",
-            model_type=ModelType.FRESH,
-            model_cache_dir=create_model_cache_dir(
-                cache_dir=self.benchmark_config.cache_dir, model_id=model_id
-            ),
-            adapter_base_model_id=None,
-        )
-
-    def load_model(
-        self, model_config: ModelConfig, dataset_config: "DatasetConfig"
-    ) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-        """Load a fresh model.
-
-        Args:
-            model_config:
-                The model configuration.
-            dataset_config:
-                The dataset configuration.
-
-        Returns:
-            The tokenizer and model.
-        """
-        config: "PretrainedConfig"
-        block_terminal_output()
-
-        model_id = model_config.model_id
-        supertask = dataset_config.task.supertask
-
-        if model_config.model_id == "xlm-roberta-base":
-            if supertask == "sequence-classification":
-                model_cls = XLMRobertaForSequenceClassification
-            elif supertask == "token-classification":
-                model_cls = XLMRobertaForTokenClassification
-            elif supertask == "question-answering":
-                model_cls = XLMRobertaForQuestionAnswering
-            else:
-                raise InvalidBenchmark(
-                    f"Supertask {supertask} is not supported for model {model_id}"
-                )
-
-        elif model_id == "electra-small":
-            model_id = "google/electra-small-discriminator"
-            if supertask == "sequence-classification":
-                model_cls = ElectraForSequenceClassification
-            elif supertask == "token-classification":
-                model_cls = ElectraForTokenClassification
-            elif supertask == "question-answering":
-                model_cls = ElectraForQuestionAnswering
-            else:
-                raise InvalidBenchmark(
-                    f"Supertask {supertask} is not supported for model {model_id}"
-                )
-
-        else:
-            raise InvalidModel(f"Model {model_id} is not supported as a fresh class.")
-
-        config = AutoConfig.from_pretrained(
-            model_id,
-            token=self.benchmark_config.token,
-            num_labels=dataset_config.num_labels,
-            id2label=dataset_config.id2label,
-            label2id=dataset_config.label2id,
-            cache_dir=model_config.model_cache_dir,
-        )
-        model = model_cls(config)
-
-        if supertask == "question-answering":
-            model = setup_model_for_question_answering(model=model)
-
-        # Load the tokenizer. If the model is a subclass of a RoBERTa model then we
-        # have to add a prefix space to the tokens, by the way the model is constructed
-        prefix_models = ["Roberta", "GPT", "Deberta"]
-        prefix = any(model_type in type(model).__name__ for model_type in prefix_models)
-        try:
-            tokenizer: "PreTrainedTokenizerBase" = AutoTokenizer.from_pretrained(
-                model_id,
-                revision=model_config.revision,
-                token=self.benchmark_config.token,
-                add_prefix_space=prefix,
-                cache_dir=model_config.model_cache_dir,
-                use_fast=True,
-                verbose=False,
-            )
-        except (JSONDecodeError, OSError):
-            raise InvalidModel(f"Could not load tokenizer for model {model_id!r}.")
-
-        model, tokenizer = align_model_and_tokenizer(
-            model=model,
-            tokenizer=tokenizer,
-            generation_length=dataset_config.max_generated_tokens,
-            raise_errors=self.benchmark_config.raise_errors,
-            generative_model=model_is_generative(model=model),
-        )
-
-        return model, tokenizer
diff --git a/src/scandeval/model_setups/hf.py b/src/scandeval/model_setups/hf.py
deleted file mode 100644
index f5b2b6b9a..000000000
--- a/src/scandeval/model_setups/hf.py
+++ /dev/null
@@ -1,711 +0,0 @@
-"""Model setup for Hugging Face Hub models."""
-
-import importlib.util
-import logging
-import os
-from json import JSONDecodeError
-from time import sleep
-from typing import TYPE_CHECKING, Type
-
-import torch
-from huggingface_hub import HfApi, HfFileSystem
-from huggingface_hub import whoami as hf_whoami
-from huggingface_hub.hf_api import ModelInfo, RepositoryNotFoundError
-from huggingface_hub.utils import (
-    GatedRepoError,
-    HFValidationError,
-    LocalTokenNotFoundError,
-)
-from requests.exceptions import RequestException
-from transformers import AutoConfig, AutoTokenizer, BitsAndBytesConfig, PreTrainedModel
-from urllib3.exceptions import RequestError
-
-from ..config import ModelConfig
-from ..enums import Framework, ModelType
-from ..exceptions import (
-    FlashAttentionNotInstalled,
-    HuggingFaceHubDown,
-    InvalidBenchmark,
-    InvalidModel,
-    MissingHuggingFaceToken,
-    NeedsAdditionalArgument,
-    NeedsExtraInstalled,
-    NeedsManualDependency,
-    NoInternetConnection,
-)
-from ..languages import get_all_languages
-from ..protocols import GenerativeModel
-from ..utils import (
-    GENERATIVE_DATASET_SUPERTASKS,
-    GENERATIVE_DATASET_TASKS,
-    GENERATIVE_MODEL_TASKS,
-    block_terminal_output,
-    create_model_cache_dir,
-    get_class_by_name,
-    internet_connection_available,
-    model_is_generative,
-)
-from ..vllm_models import VLLMModel
-from .utils import align_model_and_tokenizer, setup_model_for_question_answering
-
-if TYPE_CHECKING:
-    from peft.config import PeftConfig
-    from transformers import PretrainedConfig, PreTrainedModel
-
-    from ..config import BenchmarkConfig, DatasetConfig
-    from ..protocols import Tokenizer
-
-if importlib.util.find_spec("peft") is not None:
-    from peft.config import PeftConfig
-
-
-logger = logging.getLogger(__package__)
-
-
-class HFModelSetup:
-    """Model setup for Hugging Face Hub models.
-
-    Args:
-        benchmark_config:
-            The benchmark configuration.
-
-    Attributes:
-        benchmark_config:
-            The benchmark configuration.
-    """
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the model setup.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration.
-        """
-        self.benchmark_config = benchmark_config
-
-    def model_exists(self, model_id: str) -> bool | dict[str, str]:
-        """Check if a model ID denotes a model on the Hugging Face Hub.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            Whether the model exist, or a dictionary explaining why we cannot check
-            whether the model exists.
-        """
-        # Extract the revision from the model_id, if present
-        model_id, revision = (
-            model_id.split("@") if "@" in model_id else (model_id, "main")
-        )
-
-        # Connect to the Hugging Face Hub API
-        hf_api = HfApi()
-
-        # Get the model info, and return it
-        try:
-            hf_api.model_info(
-                repo_id=model_id, revision=revision, token=self.benchmark_config.token
-            )
-            return True
-
-        except (GatedRepoError, LocalTokenNotFoundError):
-            try:
-                hf_whoami()
-                raise NeedsAdditionalArgument(
-                    cli_argument="--use-token",
-                    script_argument="token=True",
-                    run_with_cli=self.benchmark_config.run_with_cli,
-                )
-            except LocalTokenNotFoundError:
-                raise MissingHuggingFaceToken(
-                    run_with_cli=self.benchmark_config.run_with_cli
-                )
-
-        except (RepositoryNotFoundError, HFValidationError):
-            return False
-
-        # If fetching from the Hugging Face Hub failed in a different way then throw a
-        # reasonable exception
-        except OSError:
-            if internet_connection_available():
-                raise HuggingFaceHubDown()
-            else:
-                raise NoInternetConnection()
-
-    def get_model_config(self, model_id: str) -> ModelConfig:
-        """Fetches configuration for an OpenAI model.
-
-        Args:
-            model_id:
-                The model ID of the model.
-
-        Returns:
-            The model configuration.
-        """
-        # Extract the revision from the model IDs, if they are specified
-        if "@" in model_id:
-            model_id_without_revision, revision = model_id.split("@", 1)
-        else:
-            model_id_without_revision = model_id
-            revision = "main"
-
-        # Extract the author and model name from the model ID
-        author: str | None
-        if "/" in model_id_without_revision:
-            author, model_name = model_id_without_revision.split("/")
-        else:
-            author = None
-            model_name = model_id_without_revision
-
-        # Attempt to fetch model data from the Hugging Face Hub
-        try:
-            api: HfApi = HfApi()
-            fs = HfFileSystem()
-
-            # Fetch the model metadata
-            models = api.list_models(
-                author=author, model_name=model_name, token=self.benchmark_config.token
-            )
-
-            # Filter the models to only keep the one with the specified model ID
-            models = [
-                model for model in models if model.id == model_id_without_revision
-            ]
-
-            # Check that the model exists. If it does not then raise an error
-            if len(models) == 0:
-                raise InvalidModel(
-                    f"The model {model_id} does not exist on the Hugging Face Hub."
-                )
-
-            model: ModelInfo = models[0]
-            tags: list[str] = model.tags or list()
-
-            framework = Framework.PYTORCH
-            if "pytorch" in tags:
-                pass
-            elif "jax" in tags:
-                framework = Framework.JAX
-            elif "spacy" in tags:
-                raise InvalidModel("SpaCy models are not supported.")
-            elif "tf" in tags or "tensorflow" in tags or "keras" in tags:
-                raise InvalidModel("TensorFlow/Keras models are not supported.")
-
-            base_model_ids: list[str] = list()
-
-            # If the model is a finetuned model then we fetch the base model ID
-            has_base_model_tag = any(
-                tag.startswith("base_model:") and tag.count(":") == 1 for tag in tags
-            )
-            if has_base_model_tag:
-                base_model_id = [
-                    tag.split(":")[1]
-                    for tag in tags
-                    if tag.startswith("base_model:") and tag.count(":") == 1
-                ][0]
-                base_model_ids.append(base_model_id)
-
-            # If the model is an adapter model then we fetch the base model ID
-            is_adapter = "adapter_config.json" in [
-                path.split("/")[-1]
-                for path in fs.ls(path=model_id, detail=False)
-                if isinstance(path, str)
-            ]
-            adapter_base_model_id: str | None = None
-            if is_adapter:
-                if importlib.util.find_spec("peft") is None:
-                    raise NeedsManualDependency(package="peft")
-                peft_config = PeftConfig.from_pretrained(model_id)
-                adapter_base_model_id = peft_config.base_model_name_or_path
-                if adapter_base_model_id is not None:
-                    base_model_ids.append(adapter_base_model_id)
-
-            # Add tags from the base models
-            for base_model_id in base_model_ids:
-                base_model_author, base_model_name = base_model_id.split("/")
-                base_models = [
-                    model
-                    for model in api.list_models(
-                        author=base_model_author,
-                        model_name=base_model_name,
-                        token=self.benchmark_config.token,
-                    )
-                    if model.id == base_model_id
-                ]
-                if len(base_models) > 0:
-                    tags += base_models[0].tags or list()
-                    tags = list(set(tags))
-
-            model_task: str | None = model.pipeline_tag
-            if model_task is None:
-                generative_tags = [
-                    "trl",
-                    "mistral",
-                    "text-generation-inference",
-                    "unsloth",
-                    "text-generation",
-                    "llama",
-                ]
-                if any(tag in tags for tag in generative_tags):
-                    model_task = "text-generation"
-                else:
-                    model_task = "fill-mask"
-
-            language_mapping = get_all_languages()
-            language_codes = list(language_mapping.keys())
-
-            model_config = ModelConfig(
-                model_id=model_id,
-                framework=framework,
-                task=model_task,
-                languages=[
-                    language_mapping[tag] for tag in tags if tag in language_codes
-                ],
-                revision=revision,
-                model_type=ModelType.HF,
-                model_cache_dir=create_model_cache_dir(
-                    cache_dir=self.benchmark_config.cache_dir, model_id=model_id
-                ),
-                adapter_base_model_id=adapter_base_model_id,
-            )
-
-        # If fetching from the Hugging Face Hub failed then throw a reasonable
-        # exception
-        except RequestException:
-            if internet_connection_available():
-                raise HuggingFaceHubDown()
-            else:
-                raise NoInternetConnection()
-
-        return model_config
-
-    def load_model(
-        self, model_config: ModelConfig, dataset_config: "DatasetConfig"
-    ) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-        """Load an OpenAI model.
-
-        Args:
-            model_config:
-                The model configuration.
-            dataset_config:
-                The dataset configuration.
-
-        Returns:
-            The tokenizer and model.
-        """
-        config: "PretrainedConfig"
-        block_terminal_output()
-
-        model_id = model_config.model_id
-        supertask = dataset_config.task.supertask
-        from_flax = model_config.framework == Framework.JAX
-        ignore_mismatched_sizes = False
-
-        config = self._load_hf_model_config(
-            model_id=model_config.adapter_base_model_id or model_id,
-            num_labels=dataset_config.num_labels,
-            id2label=dataset_config.id2label,
-            label2id=dataset_config.label2id,
-            revision=model_config.revision,
-            model_cache_dir=model_config.model_cache_dir,
-        )
-
-        quantization = None
-        if hasattr(config, "quantization_config"):
-            quantization = config.quantization_config.get("quant_method", None)
-        if quantization == "gptq" and importlib.util.find_spec("auto_gptq") is None:
-            raise NeedsManualDependency(package="auto-gptq")
-        if quantization == "gptq" and importlib.util.find_spec("optimum") is None:
-            raise NeedsManualDependency(package="optimum")
-        if quantization == "awq" and importlib.util.find_spec("awq") is None:
-            raise NeedsManualDependency(package="autoawq")
-
-        if self.benchmark_config.load_in_4bit is not None:
-            load_in_4bit = self.benchmark_config.load_in_4bit
-        else:
-            load_in_4bit = (
-                model_config.task in GENERATIVE_MODEL_TASKS
-                and self.benchmark_config.device == torch.device("cuda")
-                and (
-                    not hasattr(config, "quantization_config")
-                    or config.quantization_config is None
-                )
-            )
-
-        if load_in_4bit and importlib.util.find_spec("bitsandbytes") is None:
-            raise NeedsExtraInstalled(extra="generative")
-
-        use_bf16 = (
-            self.benchmark_config.device == torch.device("cuda")
-            and torch.cuda.is_bf16_supported()
-            and config.to_dict().get("torch_dtype") == "bfloat16"
-        )
-        bnb_config = (
-            BitsAndBytesConfig(
-                load_in_4bit=load_in_4bit,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_compute_dtype=torch.bfloat16 if use_bf16 else torch.float16,
-            )
-            if load_in_4bit
-            else None
-        )
-
-        use_vllm = (
-            model_config.task in GENERATIVE_MODEL_TASKS
-            and self.benchmark_config.device == torch.device("cuda")
-            and os.getenv("USE_VLLM", "1") == "1"
-        )
-
-        if use_vllm and importlib.util.find_spec("vllm") is None:
-            raise NeedsExtraInstalled(extra="generative")
-
-        model: "PreTrainedModel | GenerativeModel | None" = None
-        if use_vllm:
-            try:
-                model = VLLMModel(
-                    model_config=model_config,
-                    hf_model_config=config,
-                    model_cache_dir=model_config.model_cache_dir,
-                    trust_remote_code=self.benchmark_config.trust_remote_code,
-                    adapter_base_model_id=model_config.adapter_base_model_id,
-                )
-            except ValueError as e:
-                # If the model is too large to fit on the GPU then we simply throw an
-                # informative error message
-                oom_error_message = "No available memory for the cache blocks"
-                if oom_error_message in str(e):
-                    raise InvalidModel("The model is too large to load on the GPU.")
-
-                if self.benchmark_config.raise_errors:
-                    raise e
-
-                # Otherwise some other error occurred, and we log it and try to load
-                # the model with Hugging Face instead
-                use_vllm = False
-                logger.info(
-                    "Failed to benchmark with vLLM - trying with the Hugging Face "
-                    f"implementation instead. The error raised was {e!r}"
-                )
-
-        if not use_vllm:
-            if self.benchmark_config.use_flash_attention is None:
-                flash_attention = model_config.task in GENERATIVE_MODEL_TASKS
-            else:
-                flash_attention = self.benchmark_config.use_flash_attention
-
-            model_kwargs = dict(
-                config=config,
-                from_flax=from_flax,
-                ignore_mismatched_sizes=ignore_mismatched_sizes,
-                revision=model_config.revision,
-                token=self.benchmark_config.token,
-                cache_dir=model_config.model_cache_dir,
-                trust_remote_code=self.benchmark_config.trust_remote_code,
-                quantization_config=bnb_config,
-                torch_dtype=self._get_torch_dtype(config=config),
-                attn_implementation="flash_attention_2" if flash_attention else None,
-                device_map=(
-                    "cuda:0"
-                    if (
-                        hasattr(config, "quantization_config")
-                        and config.quantization_config.get("quant_method") == "gptq"
-                    )
-                    else None
-                ),
-            )
-
-            # These are used when a timeout occurs
-            attempts_left = 5
-
-            while True:
-                try:
-                    # Get the model class associated with the supertask
-                    if model_config.task in ["text-generation", "conversational"]:
-                        model_cls_supertask = "causal-l-m"
-                    elif model_config.task == "text2text-generation":
-                        model_cls_supertask = "seq-2-seq-l-m"
-                    elif (
-                        dataset_config.task.name in GENERATIVE_DATASET_TASKS
-                        or supertask in GENERATIVE_DATASET_SUPERTASKS
-                    ):
-                        raise InvalidBenchmark(
-                            f"The {dataset_config.task.name!r} task is not supported "
-                            f"for the model {model_id!r}."
-                        )
-                    else:
-                        model_cls_supertask = supertask
-                    model_cls_or_none: Type["PreTrainedModel"] | None = (
-                        get_class_by_name(
-                            class_name=f"auto-model-for-{model_cls_supertask}",
-                            module_name="transformers",
-                        )
-                    )
-
-                    # If the model class could not be found then raise an error
-                    if not model_cls_or_none:
-                        raise InvalidBenchmark(
-                            f"The supertask {supertask!r} does not correspond to a "
-                            "Hugging Face AutoModel type (such as "
-                            "`AutoModelForSequenceClassification`)."
-                        )
-
-                    # If the model is a DeBERTaV2 model then we ensure that
-                    # `pooler_hidden_size` is the same size as `hidden_size`
-                    if config.model_type == "deberta-v2":
-                        config.pooler_hidden_size = config.hidden_size
-
-                    try:
-                        model_or_tuple = model_cls_or_none.from_pretrained(
-                            model_config.model_id, **model_kwargs
-                        )
-                    except ImportError as e:
-                        if "flash attention" in str(e).lower():
-                            raise FlashAttentionNotInstalled()
-                        else:
-                            raise e
-                    except (KeyError, RuntimeError) as e:
-                        if not model_kwargs["ignore_mismatched_sizes"]:
-                            logger.debug(
-                                f"{type(e).__name__} occurred during the loading "
-                                f"of the {model_id!r} model. Retrying with "
-                                "`ignore_mismatched_sizes` set to True."
-                            )
-                            model_kwargs["ignore_mismatched_sizes"] = True
-                            continue
-                        else:
-                            raise InvalidModel(str(e))
-                    except (TimeoutError, RequestError):
-                        attempts_left -= 1
-                        if attempts_left == 0:
-                            raise InvalidModel(
-                                "The model could not be loaded after 5 attempts."
-                            )
-                        logger.info(f"Couldn't load the model {model_id!r}. Retrying.")
-                        sleep(5)
-                        continue
-                    except ValueError as e:
-                        if "already quantized" in str(e):
-                            model_kwargs["quantization_config"] = None
-                            model_or_tuple = model_cls_or_none.from_pretrained(
-                                model_config.model_id, **model_kwargs
-                            )
-                        elif "does not support Flash Attention" in str(e):
-                            model_kwargs["attn_implementation"] = None
-                            continue
-                        else:
-                            raise e
-
-                    if isinstance(model_or_tuple, tuple):
-                        model = model_or_tuple[0]
-                    else:
-                        model = model_or_tuple
-                    break
-
-                except (OSError, ValueError) as e:
-                    # If `from_flax` is False but only Flax models are available then
-                    # try again with `from_flax` set to True
-                    if (
-                        not from_flax
-                        and "Use `from_flax=True` to load this model" in str(e)
-                    ):
-                        from_flax = True
-                        continue
-
-                    self._handle_loading_exception(exception=e, model_id=model_id)
-
-        assert isinstance(model, (PreTrainedModel, GenerativeModel))
-        model.eval()
-        if not load_in_4bit:
-            model.to(self.benchmark_config.device)
-
-        generative_model = model_is_generative(model=model)
-
-        if isinstance(model, PreTrainedModel) and supertask == "question-answering":
-            model = setup_model_for_question_answering(model=model)
-
-        tokenizer = self._load_tokenizer(
-            model=model,
-            model_id=model_config.adapter_base_model_id or model_id,
-            generative_model=generative_model,
-        )
-
-        if use_vllm and isinstance(model, VLLMModel):
-            model.set_tokenizer(tokenizer=tokenizer)
-
-        model, tokenizer = align_model_and_tokenizer(
-            model=model,
-            tokenizer=tokenizer,
-            generative_model=generative_model,
-            generation_length=dataset_config.max_generated_tokens,
-            raise_errors=self.benchmark_config.raise_errors,
-        )
-
-        return model, tokenizer
-
-    def _get_torch_dtype(self, config: "PretrainedConfig") -> str | torch.dtype:
-        """Get the torch dtype, used for loading the model.
-
-        Args:
-            config:
-                The Hugging Face model configuration.
-
-        Returns:
-            The torch dtype.
-        """
-        using_cuda = self.benchmark_config.device == torch.device("cuda")
-        torch_dtype_is_set = config.to_dict().get("torch_dtype") is not None
-        bf16_available = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
-        if using_cuda and torch_dtype_is_set:
-            return "auto"
-        elif using_cuda and bf16_available:
-            return torch.bfloat16
-        elif using_cuda:
-            return torch.float16
-        return torch.float32
-
-    def _load_hf_model_config(
-        self,
-        model_id: str,
-        num_labels: int,
-        id2label: dict[int, str],
-        label2id: dict[str, int],
-        revision: str,
-        model_cache_dir: str | None,
-    ) -> "PretrainedConfig":
-        """Load the Hugging Face model configuration.
-
-        Args:
-            model_id:
-                The Hugging Face model ID.
-            num_labels:
-                The number of labels in the dataset.
-            id2label:
-                The mapping from label IDs to labels.
-            label2id:
-                The mapping from labels to label IDs.
-            revision:
-                The revision of the model.
-            model_cache_dir:
-                The directory to cache the model in.
-
-        Returns:
-            The Hugging Face model configuration.
-        """
-        while True:
-            try:
-                config = AutoConfig.from_pretrained(
-                    model_id,
-                    num_labels=num_labels,
-                    id2label=id2label,
-                    label2id=label2id,
-                    revision=revision,
-                    token=self.benchmark_config.token,
-                    trust_remote_code=self.benchmark_config.trust_remote_code,
-                    cache_dir=model_cache_dir,
-                )
-                if config.eos_token_id is not None and config.pad_token_id is None:
-                    if isinstance(config.eos_token_id, list):
-                        config.pad_token_id = config.eos_token_id[0]
-                    else:
-                        config.pad_token_id = config.eos_token_id
-                return config
-            except KeyError as e:
-                key = e.args[0]
-                raise InvalidModel(
-                    f"The model config for the model {model_id!r} could not be "
-                    f"loaded, as the key {key!r} was not found in the config."
-                )
-            except OSError as e:
-                # TEMP: When the model is gated then we cannot set cache dir, for some
-                # reason (transformers==4.38.2). This should be included back in when
-                # this is fixed.
-                if "gated repo" in str(e):
-                    model_cache_dir = None
-                    continue
-                raise InvalidModel(
-                    f"Couldn't load model config for {model_id!r}. The error was "
-                    f"{e!r}. Skipping"
-                )
-            except (TimeoutError, RequestError):
-                logger.info(f"Couldn't load model config for {model_id!r}. Retrying.")
-                sleep(5)
-                continue
-            except ValueError as e:
-                requires_trust_remote_code = "trust_remote_code" in str(e)
-                if requires_trust_remote_code:
-                    raise NeedsAdditionalArgument(
-                        cli_argument="--trust-remote-code",
-                        script_argument="trust_remote_code=True",
-                        run_with_cli=self.benchmark_config.run_with_cli,
-                    )
-                raise e
-
-    def _load_tokenizer(
-        self,
-        model: "PreTrainedModel | GenerativeModel",
-        model_id: str,
-        generative_model: bool,
-    ) -> "Tokenizer":
-        """Load the tokenizer.
-
-        Args:
-            model:
-                The model, which is used to determine whether to add a prefix space to
-                the tokens.
-            model_id:
-                The model identifier. Used for logging.
-            generative_model:
-                Whether the model is a generative model.
-
-        Returns:
-            The loaded tokenizer.
-        """
-        loading_kwargs: dict[str, bool | str] = dict(
-            use_fast=True,
-            verbose=False,
-            trust_remote_code=self.benchmark_config.trust_remote_code,
-        )
-
-        # If the model is a subclass of a certain model types then we have to add a
-        # prefix space to the tokens, by the way the model is constructed.
-        prefix_models = ["Roberta", "GPT", "Deberta"]
-        add_prefix = any(
-            model_type in type(model).__name__ for model_type in prefix_models
-        )
-        if add_prefix:
-            loading_kwargs["add_prefix_space"] = True
-
-        padding_side = "left" if generative_model else "right"
-        loading_kwargs["padding_side"] = padding_side
-        loading_kwargs["truncation_side"] = padding_side
-
-        while True:
-            try:
-                return AutoTokenizer.from_pretrained(model_id, **loading_kwargs)
-            except (JSONDecodeError, OSError, TypeError):
-                raise InvalidModel(f"Could not load tokenizer for model {model_id!r}.")
-            except (TimeoutError, RequestError):
-                logger.info(f"Couldn't load tokenizer for {model_id!r}. Retrying.")
-                sleep(5)
-                continue
-
-    @staticmethod
-    def _handle_loading_exception(exception: Exception, model_id: str) -> None:
-        if "checkpoint seems to be incorrect" in str(exception):
-            raise InvalidModel(f"The model {model_id!r} has an incorrect checkpoint.")
-        if "trust_remote_code" in str(exception):
-            raise InvalidModel(
-                f"Loading the model {model_id!r} needs to trust remote code. "
-                "If you trust the suppliers of this model, then you can enable "
-                "this by setting the `--trust-remote-code` flag."
-            )
-        raise InvalidModel(
-            f"The model {model_id} either does not exist on the Hugging Face "
-            "Hub, or it has no frameworks registered, or it is a private "
-            "model. If it *does* exist on the Hub and is a public model then "
-            "please ensure that it has a framework registered. If it is a "
-            "private model then enable the `--use-token` flag and make "
-            "sure that you are logged in to the Hub via the "
-            f"`huggingface-cli login` command. The error raised was {exception!r}"
-        )
diff --git a/src/scandeval/model_setups/local.py b/src/scandeval/model_setups/local.py
deleted file mode 100644
index 4142cfddf..000000000
--- a/src/scandeval/model_setups/local.py
+++ /dev/null
@@ -1,174 +0,0 @@
-"""Model setup for local Hugging Face Hub models."""
-
-import importlib.util
-import logging
-from pathlib import Path
-from typing import TYPE_CHECKING
-
-from transformers import AutoConfig
-from transformers.models.auto.modeling_auto import (
-    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
-    MODEL_FOR_MASKED_LM_MAPPING_NAMES,
-    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
-)
-
-from ..config import ModelConfig
-from ..enums import Framework, ModelType
-from ..exceptions import InvalidModel, NeedsManualDependency
-from ..utils import create_model_cache_dir
-from .hf import HFModelSetup
-
-if TYPE_CHECKING:
-    from peft.config import PeftConfig
-    from transformers import PreTrainedModel
-
-    from ..config import BenchmarkConfig, DatasetConfig
-    from ..protocols import GenerativeModel, Tokenizer
-
-if importlib.util.find_spec("peft") is not None:
-    from peft.config import PeftConfig
-
-
-logger = logging.getLogger(__package__)
-
-
-class LocalModelSetup:
-    """Model setup for local Hugging Face Hub models.
-
-    Attributes:
-        benchmark_config:
-            The benchmark configuration.
-    """
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the LocalModelSetup class.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration.
-        """
-        self.benchmark_config = benchmark_config
-
-    def model_exists(self, model_id: str) -> bool | dict[str, str]:
-        """Check if a model exists locally.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            Whether the model exist, or a dictionary explaining why we cannot check
-            whether the model exists.
-        """
-        model_dir = Path(model_id)
-
-        if not model_dir.exists():
-            return False
-
-        try:
-            AutoConfig.from_pretrained(model_id)
-        except OSError:
-            return False
-
-        return (
-            model_dir.glob("*.bin") is not None
-            or model_dir.glob("*.pt") is not None
-            or model_dir.glob("*.msgpack") is not None
-            or model_dir.glob("*.safetensors") is not None
-        )
-
-    def get_model_config(self, model_id: str) -> ModelConfig:
-        """Fetches configuration for a local Hugging Face model.
-
-        Args:
-            model_id:
-                The model ID of the model.
-
-        Returns:
-            The model configuration.
-        """
-        framework = self.benchmark_config.framework
-        if framework is None:
-            try:
-                exts = {f.suffix for f in Path(model_id).iterdir()}
-                if ".bin" in exts:
-                    framework = Framework.PYTORCH
-                elif ".msgpack" in exts:
-                    framework = Framework.JAX
-                elif ".whl" in exts:
-                    raise InvalidModel("SpaCy models are not supported.")
-                elif ".h5" in exts:
-                    raise InvalidModel("TensorFlow/Keras models are not supported.")
-            except OSError as e:
-                logger.info(f"Cannot list files for local model `{model_id}`!")
-                if self.benchmark_config.raise_errors:
-                    raise e
-
-        if framework is None:
-            msg = f"Assuming 'pytorch' as the framework for local model `{model_id}`. "
-            if self.benchmark_config.run_with_cli:
-                msg += (
-                    "If this is not the case then please use the --framework argument "
-                    "to override."
-                )
-            else:
-                msg += (
-                    "If this is not the case then please use the `framework` argument "
-                    "in the `Benchmarker` class to override."
-                )
-            logger.info(msg)
-            framework = Framework.PYTORCH
-
-        hf_model_config = AutoConfig.from_pretrained(model_id)
-        model_type = hf_model_config.model_type.lower()
-        if model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
-            task = "text-generation"
-        elif model_type in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
-            task = "text2text-generation"
-        elif model_type in MODEL_FOR_MASKED_LM_MAPPING_NAMES:
-            task = "fill-mask"
-        else:
-            task = "unknown"
-
-        is_adapter = "adapter_config.json" in [
-            path.name for path in Path(model_id).glob("*.json")
-        ]
-        adapter_base_model_id: str | None = None
-        if is_adapter:
-            if importlib.util.find_spec("peft") is None:
-                raise NeedsManualDependency(package="peft")
-            peft_config = PeftConfig.from_pretrained(model_id)
-            adapter_base_model_id = peft_config.base_model_name_or_path
-
-        model_config = ModelConfig(
-            model_id=model_id,
-            revision="main",
-            framework=framework,
-            task=task,
-            languages=list(),
-            model_type=ModelType.LOCAL,
-            model_cache_dir=create_model_cache_dir(
-                cache_dir=self.benchmark_config.cache_dir, model_id=model_id
-            ),
-            adapter_base_model_id=adapter_base_model_id,
-        )
-        return model_config
-
-    def load_model(
-        self, model_config: ModelConfig, dataset_config: "DatasetConfig"
-    ) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-        """Load a local Hugging Face model.
-
-        Args:
-            model_config:
-                The model configuration.
-            dataset_config:
-                The dataset configuration.
-
-        Returns:
-            The tokenizer and model.
-        """
-        hf_model_setup = HFModelSetup(benchmark_config=self.benchmark_config)
-        return hf_model_setup.load_model(
-            model_config=model_config, dataset_config=dataset_config
-        )
diff --git a/src/scandeval/model_setups/openai.py b/src/scandeval/model_setups/openai.py
deleted file mode 100644
index c4244f85b..000000000
--- a/src/scandeval/model_setups/openai.py
+++ /dev/null
@@ -1,246 +0,0 @@
-"""Model setup for OpenAI models."""
-
-import importlib.util
-import logging
-import re
-from typing import TYPE_CHECKING
-
-from transformers import PretrainedConfig
-
-from ..config import ModelConfig
-from ..enums import Framework, ModelType
-from ..openai_models import OpenAIModel, OpenAITokenizer
-from ..utils import create_model_cache_dir
-
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-    from ..config import BenchmarkConfig, DatasetConfig
-    from ..protocols import GenerativeModel, Tokenizer
-
-if importlib.util.find_spec("openai") is not None:
-    import openai
-
-    # Older versions of `openai` doesn't have the `models` module, so we need to check
-    # that, as it will cause errors later otherwise
-    openai.models
-
-
-logger = logging.getLogger(__package__)
-
-
-# This is a list of the major models that OpenAI has released
-CACHED_OPENAI_MODEL_IDS: list[str] = [
-    "ada|babbage|curie|davinci",
-    "(code|text)-(ada|babbage|curie|davinci)-[0-9]{3}",
-    "gpt-3.5-turbo(-16k|-instruct)?(-[0-9]{4})?(-preview)?",
-    "gpt-4(-[0-9]{4})?(-preview)?",
-    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?",
-    "gpt-4-32k(-[0-9]{4})?(-preview)?",
-    "gpt-4o(-mini)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?",
-]
-
-
-VOCAB_SIZE_MAPPING = {
-    "(text-)?(ada|babbage|curie|davinci)(-001)?": 50_257,
-    "(code|text)-davinci-00[2-9]": 50_281,
-    "gpt-3.5-turbo(-16k)?(-[0-9]{4})?": 100_256,
-    "gpt-4-(32k)?(-[0-9]{4})?": 100_256,
-    "gpt-4-[0-9]{4}-preview": 100_256,
-    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 100_256,
-    "gpt-4-(vision|turbo)(-preview)?": 100_256,
-    "gpt-3.5-turbo-instruct(-[0-9]{4})?": 100_256,
-    "gpt-4o(-mini)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 200_019,
-}
-
-
-MODEL_MAX_LENGTH_MAPPING = {
-    "(text-)?(ada|babbage|curie|davinci)(-001)?": 2_050,
-    "text-davinci-00[2-9]": 4_098,
-    "code-davinci-00[1-9]": 8_002,
-    "gpt-3.5-turbo-0613": 4_096,
-    "gpt-3.5-turbo(-[0-9]{4})?": 16_385,
-    "gpt-3.5-turbo-16k(-[0-9]{4})?": 16_384,
-    "gpt-4(-[0-9]{4})?": 8_191,
-    "gpt-4-32k(-[0-9]{4})?": 32_767,
-    "gpt-4-[0-9]{4}-preview": 128_000,
-    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 128_000,
-    "gpt-4-(vision|turbo)(-preview)?": 128_000,
-    "gpt-3.5-turbo-instruct(-[0-9]{4})?": 4_095,
-    "gpt-4o(-mini)?(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": 128_000,
-}
-
-
-NUM_PARAMS_MAPPING = {
-    "(text-)?ada(-001)?": 350_000_000,
-    "(text-)?babbage(-001)?": 3_000_000_000,
-    "(text-)?curie(-001)?": 13_000_000_000,
-    "((text|code)-)?davinci(-00[1-9])?": 175_000_000_000,
-    "gpt-(3.5|4)-turbo-((16|32)k)?(-[0-9]{4})?": -1,
-    "gpt-4-[0-9]{4}-preview": -1,
-    "gpt-4-turbo(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
-    "gpt-4-(vision|turbo)(-preview)?": -1,
-    "gpt-3.5-turbo-instruct(-[0-9]{4})?": -1,
-    "gpt-4o(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
-    "gpt-4o-mini(-[0-9]{4}-[0-9]{2}-[0-9]{2})?": -1,
-}
-
-
-class OpenAIModelSetup:
-    """Model setup for OpenAI models.
-
-    Attributes:
-        benchmark_config:
-            The benchmark configuration.
-    """
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the model setup.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration.
-        """
-        self.benchmark_config = benchmark_config
-
-    def model_exists(self, model_id: str) -> bool | dict[str, str]:
-        """Check if a model ID denotes an OpenAI model.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            Whether the model exists, or a dictionary explaining why we cannot check
-            whether the model exists.
-        """
-        if importlib.util.find_spec("openai") is None:
-            return dict(missing_extra="openai")
-
-        # The model ID for the Azure OpenAI API is the deployment name and therefore
-        # different from the model ID used in the OpenAI API. We'll just assume that
-        # the model exists in this case.
-        if self.benchmark_config.azure_openai_api_key is not None:
-            return True
-
-        all_models: list[openai.types.model.Model] = list()
-        try:
-            all_models = list(openai.models.list())
-            return model_id in [model.id for model in all_models]
-        except openai.OpenAIError as e:
-            model_exists = any(
-                [
-                    re.match(pattern=model_pattern, string=model_id) is not None
-                    for model_pattern in CACHED_OPENAI_MODEL_IDS
-                ]
-            )
-            if not model_exists:
-                if "OPENAI_API_KEY" in str(e):
-                    return dict(missing_env_var="OPENAI_API_KEY")
-                elif "AZURE_OPENAI_API_KEY" in str(e):
-                    return dict(missing_env_var="AZURE_OPENAI_API_KEY")
-                elif "AZURE_OPENAI_ENDPOINT" in str(e):
-                    return dict(missing_env_var="AZURE_OPENAI_ENDPOINT")
-            return model_exists
-
-    def get_model_config(self, model_id: str) -> ModelConfig:
-        """Fetches configuration for an OpenAI model.
-
-        Args:
-            model_id:
-                The model ID of the model.
-
-        Returns:
-            The model configuration.
-        """
-        return ModelConfig(
-            model_id=model_id,
-            revision="main",
-            framework=Framework.API,
-            task="text-generation",
-            languages=list(),
-            model_type=ModelType.OPENAI,
-            model_cache_dir=create_model_cache_dir(
-                cache_dir=self.benchmark_config.cache_dir, model_id=model_id
-            ),
-            adapter_base_model_id=None,
-        )
-
-    def load_model(
-        self, model_config: ModelConfig, dataset_config: "DatasetConfig"
-    ) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-        """Load an OpenAI model.
-
-        Args:
-            model_config:
-                The model configuration.
-            dataset_config:
-                The dataset configuration.
-
-        Returns:
-            The tokenizer and model.
-        """
-        hf_model_config = PretrainedConfig.from_pretrained("gpt2")
-
-        vocab_sizes = [
-            vocab_size
-            for pattern, vocab_size in VOCAB_SIZE_MAPPING.items()
-            if re.match(pattern=pattern, string=model_config.model_id)
-        ]
-        hf_model_config.vocab_size = vocab_sizes[0] if vocab_sizes else 100_256
-
-        model_lengths = [
-            model_length
-            for pattern, model_length in MODEL_MAX_LENGTH_MAPPING.items()
-            if re.match(pattern=f"^{pattern}$", string=model_config.model_id)
-        ]
-        hf_model_config.model_max_length = model_lengths[0] if model_lengths else -1
-
-        num_params = [
-            num_param
-            for pattern, num_param in NUM_PARAMS_MAPPING.items()
-            if re.match(pattern=pattern, string=model_config.model_id)
-        ]
-        hf_model_config.num_params = num_params[0] if num_params else -1
-
-        hf_model_config.id2label = dataset_config.id2label
-        hf_model_config.label2id = dataset_config.label2id
-        hf_model_config.eos_token_id = hf_model_config.vocab_size - 1
-        hf_model_config.bos_token_id = hf_model_config.vocab_size - 1
-        hf_model_config.pad_token_id = hf_model_config.vocab_size - 1
-
-        # Check if the vocab size is correct, and if not then correct it
-        tok = OpenAITokenizer(
-            model_config=model_config, hf_model_config=hf_model_config
-        )
-        for idx in range(hf_model_config.vocab_size - 1, 0, -1):
-            try:
-                tok.decode([idx])
-                hf_model_config.vocab_size = idx + 1
-                break
-            except Exception:
-                pass
-        else:
-            raise ValueError(
-                f"Couldn't find vocab size for the model {model_config.model_id!r}"
-            )
-
-        tokenizer = OpenAITokenizer(
-            model_config=model_config, hf_model_config=hf_model_config
-        )
-        model = OpenAIModel(
-            model_config=model_config,
-            hf_model_config=hf_model_config,
-            dataset_config=dataset_config,
-            benchmark_config=self.benchmark_config,
-            tokenizer=tokenizer,
-        )
-
-        # If the model is a chat model then we need to reduce the maximum context
-        # length by 7 tokens, as these are used in the chat prompt
-        if model.is_chat_model:
-            hf_model_config.model_max_length -= 7
-            tokenizer.hf_model_config = hf_model_config
-            model.config = hf_model_config
-
-        return model, tokenizer
diff --git a/src/scandeval/model_setups/utils.py b/src/scandeval/model_setups/utils.py
deleted file mode 100644
index fc14db045..000000000
--- a/src/scandeval/model_setups/utils.py
+++ /dev/null
@@ -1,214 +0,0 @@
-"""Utility functions related to setting up models."""
-
-from typing import TYPE_CHECKING, Any
-
-import torch
-import torch.nn as nn
-
-from ..exceptions import InvalidModel
-from ..protocols import GenerativeModel, Tokenizer
-from ..utils import DUMMY_FILL_VALUE, get_model_max_length
-from ..vllm_models import VLLMModel
-
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-
-def get_children_of_module(
-    name: str, module: nn.Module
-) -> nn.Module | dict[str, Any] | None:
-    """Get the children of a module.
-
-    Args:
-        name:
-            The name of the module.
-        module:
-            The module to get the children of.
-
-    Returns:
-        The children of the module, or None if the module has no children.
-    """
-    if len(list(module.children())) == 0:
-        if name == "token_type_embeddings":
-            return module
-        else:
-            return None
-    else:
-        submodules = dict()
-        for subname, submodule in module.named_children():
-            children = get_children_of_module(name=subname, module=submodule)
-            if children:
-                submodules[subname] = children
-        return submodules
-
-
-def setup_model_for_question_answering(model: "PreTrainedModel") -> "PreTrainedModel":
-    """Setup a model for question answering.
-
-    Args:
-        model:
-            The model to setup.
-
-    Returns:
-        The setup model.
-    """
-    # Get the models' token type embedding children, if they exist
-    children = get_children_of_module(name="model", module=model)
-
-    # If the model has token type embeddings then get them
-    if children:
-        # Get the list of attributes that are token type embeddings
-        attribute_list = list()
-        done = False
-        while not done:
-            for key, value in children.items():
-                attribute_list.append(key)
-                if isinstance(value, dict):
-                    children = value
-                else:
-                    done = True
-                break
-
-        # Get the token type embeddings
-        token_type_embeddings = model
-        for attribute in attribute_list:
-            token_type_embeddings = getattr(token_type_embeddings, attribute)
-
-        # If the token type embeddings has shape (1, ...) then set the shape to
-        # (2, ...) by randomly initializing the second token type embedding
-        if token_type_embeddings.weight.data.shape[0] == 1:
-            token_type_embeddings.weight.data = torch.cat(
-                (
-                    token_type_embeddings.weight.data,
-                    torch.rand_like(token_type_embeddings.weight.data),
-                ),
-                dim=0,
-            )
-            token_type_embeddings.num_embeddings = 2
-
-        # Set the model config to use the new type vocab size
-        model.config.type_vocab_size = 2
-
-    return model
-
-
-def align_model_and_tokenizer(
-    model: "PreTrainedModel | GenerativeModel",
-    tokenizer: "Tokenizer",
-    generative_model: bool,
-    generation_length: int,
-    raise_errors: bool = False,
-) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-    """Aligns the model and the tokenizer.
-
-    Args:
-        model:
-            The model to fix.
-        tokenizer:
-            The tokenizer to fix.
-        generative_model:
-            Whether the model is a generative model.
-        generation_length:
-            The length of the generation, which depends on the benchmark dataset. Only
-            relevant if the model is a generative model.
-        raise_errors:
-            Whether to raise errors instead of trying to fix them silently.
-
-    Returns:
-        The fixed model and tokenizer.
-    """
-    model_max_length = get_model_max_length(model=model, tokenizer=tokenizer)
-
-    # If the model is a generative model then we need to subtract the generation length
-    # from the maximum length, to allow it to keep generating
-    if generative_model:
-        model_max_length -= generation_length
-
-    # Ensure that the model max length is at least 5,000, to avoid OOM errors
-    model_max_length = min(model_max_length, 5_000)
-
-    if model_max_length > 0:
-        tokenizer.model_max_length = model_max_length
-    elif generative_model:
-        tokenizer.model_max_length = 5_000
-    else:
-        tokenizer.model_max_length = 512
-
-    # If we're not dealing with a generative model then we move it to CPU to avoid OOM
-    # errors
-    device: torch.device = model.device if generative_model else torch.device("cpu")
-    model_device = model.device
-    model.to(device)
-
-    # Manually check that this model max length is valid for the model, and adjust
-    # otherwise
-    initial_max_length = tokenizer.model_max_length
-    for max_length in range(initial_max_length, 0, -1):
-        tokenizer.model_max_length = max_length
-        dummy_inputs = torch.full(
-            size=(1, max_length),
-            fill_value=DUMMY_FILL_VALUE,
-            dtype=torch.long,
-            device=device,
-        )
-
-        with torch.inference_mode():
-            try:
-                model(dummy_inputs)
-                break
-
-            # This handles the case where the model is a sequence-to-sequence model, as
-            # they require text labels to be passed in
-            except ValueError as e:
-                if "decoder_input_ids" not in str(e) or isinstance(model, VLLMModel):
-                    raise e
-                model(input_ids=dummy_inputs, labels=torch.zeros(1, 1).long())
-                break
-
-            # This happens if `max_length` is too large
-            except IndexError:
-                continue
-
-    # If there is a mismatch between the vocab size according to the tokenizer and
-    # the vocab size according to the model, we raise an error
-    if hasattr(model.config, "vocab_size") and hasattr(tokenizer, "vocab_size"):
-        if model.config.vocab_size < tokenizer.vocab_size:
-            if raise_errors:
-                raise InvalidModel(
-                    "The vocab size of the tokenizer is larger than the vocab size of "
-                    "the model. As the --raise-errors option was specified, the "
-                    "embeddings of the model will not be automatically adjusted."
-                )
-            if isinstance(model, VLLMModel):
-                raise InvalidModel(
-                    "The vocab size of the tokenizer is larger than the vocab size of "
-                    "the model."
-                )
-            if hasattr(model, "resize_token_embeddings"):
-                model.resize_token_embeddings(new_num_tokens=tokenizer.vocab_size + 1)
-
-    # For generative models, the `transformers` package requires the pad token to be
-    # identical to the eos token, if the latter exists. Otherwise, if both the pad and
-    # eos token are not defined, then we attempt to set the padding token to the sep
-    # token. If a sep token doesn't exist either, we raise an error.
-    if generative_model:
-        tokenizer.padding_side = "left"
-        if tokenizer.eos_token is not None:
-            tokenizer.pad_token = tokenizer.eos_token
-        elif tokenizer.pad_token is None:
-            if tokenizer.sep_token is not None:
-                tokenizer.pad_token = tokenizer.sep_token
-            else:
-                raise InvalidModel(
-                    "The tokenizer does not have a padding token and does not have a "
-                    "SEP token or EOS token to use as a padding token."
-                )
-        model.config.pad_token_id = tokenizer.pad_token_id
-
-    if tokenizer.bos_token is None and tokenizer.eos_token is not None:
-        tokenizer.bos_token = tokenizer.eos_token
-        tokenizer.bos_token_id = tokenizer.eos_token_id
-
-    model.to(model_device)
-
-    return model, tokenizer
diff --git a/src/scandeval/named_entity_recognition.py b/src/scandeval/named_entity_recognition.py
deleted file mode 100644
index 21d9664fa..000000000
--- a/src/scandeval/named_entity_recognition.py
+++ /dev/null
@@ -1,720 +0,0 @@
-"""Named entity recognition benchmark dataset."""
-
-import importlib.util
-import itertools as it
-import json
-import logging
-import random
-import re
-from copy import deepcopy
-from functools import partial
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-from evaluate import EvaluationModule
-from transformers import DataCollatorWithPadding
-from transformers.data.data_collator import DataCollatorForTokenClassification
-
-from .benchmark_dataset import BenchmarkDataset
-from .exceptions import InvalidBenchmark, NeedsExtraInstalled
-from .generation import extract_raw_predictions
-from .utils import (
-    GENERATIVE_MODEL_TASKS,
-    convert_prompt_to_instruction,
-    model_is_generative,
-    raise_if_model_output_contains_nan_values,
-)
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-    from datasets.dataset_dict import DatasetDict
-    from transformers import BatchEncoding, PreTrainedModel
-    from transformers.modeling_utils import ModelOutput
-
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Labels, Predictions
-
-if importlib.util.find_spec("demjson3") is not None:
-    import demjson3
-
-
-logger = logging.getLogger(__package__)
-
-
-class NamedEntityRecognition(BenchmarkDataset):
-    """Named entity recognition benchmark dataset.
-
-    Args:
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
-
-    Attributes:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-    """
-
-    def _process_data(self, dataset_dict: "DatasetDict") -> "DatasetDict":
-        """Process the data.
-
-        Args:
-            dataset_dict:
-                The dataset dictionary.
-
-        Returns:
-            The processed dataset dictionary.
-        """
-        # Check what labels are present in the dataset, and store if MISC tags are not
-        # present
-        labels_in_train: set[str] = {
-            tag for tag_list in dataset_dict["train"]["labels"] for tag in tag_list
-        }
-        self.has_misc_tags = "B-MISC" in labels_in_train or "I-MISC" in labels_in_train
-
-        return dataset_dict
-
-    def _compute_metrics(
-        self,
-        model_outputs_and_labels: tuple["Predictions", "Labels"],
-        id2label: dict[int, str],
-    ) -> dict[str, float]:
-        """Compute the metrics needed for evaluation.
-
-        Args:
-            model_outputs_and_labels:
-                The first array contains the probability predictions and the second
-                array contains the true labels.
-            id2label:
-                Conversion of indices to labels.
-
-        Returns:
-            A dictionary with the names of the metrics as keys and the metric values as
-            values.
-        """
-        model_outputs, labels = model_outputs_and_labels
-
-        raise_if_model_output_contains_nan_values(model_output=model_outputs)
-
-        predictions: list[list[str]]
-        if not isinstance(model_outputs[0][0], str):
-            raw_predictions: list[list[int]] = np.argmax(
-                model_outputs, axis=-1
-            ).tolist()
-
-            # Remove ignored index (special tokens)
-            predictions = [
-                [
-                    id2label[pred_id]
-                    for pred_id, lbl_id in zip(pred, label)
-                    if lbl_id != -100
-                ]
-                for pred, label in zip(raw_predictions, labels)
-            ]
-            labels = [
-                [
-                    (
-                        id2label[int(lbl_id)]
-                        if isinstance(lbl_id, int) or isinstance(lbl_id, np.int_)
-                        else lbl_id
-                    )
-                    for lbl_id in label
-                    if lbl_id != -100
-                ]
-                for label in labels
-            ]
-
-        else:
-            predictions = model_outputs  # type: ignore[assignment]
-
-        # Replace predicted tag with either MISC or O tags if they are not part of the
-        # dataset
-        labels_without_misc = {
-            label
-            for label in self.dataset_config.id2label.values()
-            if label not in {"b-misc", "i-misc"}
-        }
-        ner_tag: str
-        for i, prediction_list in enumerate(predictions):
-            for j, ner_tag in enumerate(prediction_list):
-                if ner_tag not in labels_without_misc:
-                    if self.has_misc_tags and ner_tag[:2] == "b-":
-                        predictions[i][j] = "b-misc"
-                    elif self.has_misc_tags and ner_tag[:2] == "i-":
-                        predictions[i][j] = "i-misc"
-                    else:
-                        predictions[i][j] = "o"
-
-        # Remove MISC labels from predictions
-        predictions_no_misc = deepcopy(predictions)
-        for i, prediction_list in enumerate(predictions_no_misc):
-            for j, ner_tag in enumerate(prediction_list):
-                if ner_tag[-4:] == "misc":
-                    predictions_no_misc[i][j] = "o"
-
-        # Remove MISC labels from labels
-        labels_no_misc: list[list[str]] = deepcopy(labels)  # type: ignore[arg-type]
-        for i, label_list in enumerate(labels_no_misc):
-            for j, ner_tag in enumerate(label_list):
-                if (
-                    isinstance(ner_tag, str)
-                    and len(ner_tag) >= 4
-                    and ner_tag[-4:] == "misc"
-                ):
-                    labels_no_misc[i][j] = "o"
-
-        # Compute the metrics
-        # We manually set the F1 metric to be 100% if both the labels and the models
-        # have no NER tags in them, since this causes an error with the `compute`
-        # method otherwise
-        predictions_all_zero = all(
-            all(ner_tag == "o" for ner_tag in prediction_list)
-            for prediction_list in predictions
-        )
-        labels_all_zero = all(
-            all(ner_tag == "o" for ner_tag in label_list) for label_list in labels
-        )
-        if predictions_all_zero and labels_all_zero:
-            results = dict(overall_f1=1.0)
-        else:
-            metric = self._metrics["micro_f1"]
-            assert isinstance(metric, EvaluationModule)
-            results = metric.compute(predictions=predictions, references=labels)
-
-        # Compute the metrics without MISC tags
-        # We manually set the F1 metric to be 100% if both the labels and the models
-        # have no NER tags in them, since this causes an error with the `compute`
-        # method otherwise
-        predictions_no_misc_all_zero = all(
-            all(ner_tag == "o" for ner_tag in prediction_list)
-            for prediction_list in predictions_no_misc
-        )
-        labels_no_misc_all_zero = all(
-            all(ner_tag == "o" for ner_tag in label_list)
-            for label_list in labels_no_misc
-        )
-        if predictions_no_misc_all_zero and labels_no_misc_all_zero:
-            results_no_misc = dict(overall_f1=1.0)
-        else:
-            metric = self._metrics["micro_f1_no_misc"]
-            assert isinstance(metric, EvaluationModule)
-            results_no_misc = metric.compute(
-                predictions=predictions_no_misc, references=labels_no_misc
-            )
-
-        # Raise error if the metrics are invalid
-        if results is None or results_no_misc is None:
-            raise InvalidBenchmark(
-                "The predictions and labels are not of the same length."
-            )
-
-        return dict(
-            micro_f1_no_misc=results_no_misc["overall_f1"],
-            micro_f1=results["overall_f1"],
-        )
-
-    def _tokenize_and_align_labels(
-        self, examples: dict, tokenizer: "Tokenizer", label2id: dict[str, int]
-    ) -> "BatchEncoding":
-        """Tokenise all texts and align the labels with them.
-
-        Args:
-            examples:
-                The examples to be tokenised.
-            tokenizer:
-                A pretrained tokenizer.
-            label2id (dict):
-                A dictionary that converts NER tags to IDs.
-
-        Returns:
-            A dictionary containing the tokenized data as well as labels.
-        """
-        # Tokenize the texts. We use the `is_split_into_words` argument here because
-        # the texts in our dataset are lists of words (with a label for each word)
-        tokenized_inputs = tokenizer(
-            examples["tokens"], is_split_into_words=True, truncation=True, padding=True
-        )
-
-        # Extract a mapping between all the tokens and their corresponding word. If the
-        # tokenizer is of a "fast" variant then this can be accessed through the
-        # `word_ids` method. Otherwise, we have to extract it manually.
-        all_labels: list[list[int]] = list()
-        labels: list[str]
-        word_ids: list[int | None]
-        for i, labels in enumerate(examples["labels"]):
-            # Try to get the word IDs from the tokenizer
-            try:
-                word_ids = tokenized_inputs.word_ids(batch_index=i)
-
-            # If the tokenizer is not of a "fast" variant, we have to extract the word
-            # IDs manually
-            except ValueError:
-                # Get the list of words in the document
-                words: list[str] = examples["tokens"][i]
-
-                # Get the list of token IDs in the document
-                tok_ids: list[int] = tokenized_inputs.input_ids[i]
-
-                # Decode the token IDs
-                tokens = tokenizer.convert_ids_to_tokens(tok_ids)
-                assert isinstance(tokens, list)
-
-                # Remove prefixes from the tokens
-                prefixes_to_remove = ["▁", "##"]
-                for tok_idx, tok in enumerate(tokens):
-                    if tok:
-                        for prefix in prefixes_to_remove:
-                            if tok.startswith(prefix):
-                                tokens[tok_idx] = tok[len(prefix) :]
-
-                # Replace UNK tokens with the correct word
-                tokens = self._handle_unk_tokens(
-                    tokenizer=tokenizer, tokens=tokens, words=words
-                )
-
-                # Get list of special tokens. Some tokenizers do not record these
-                # properly, which is why we convert the values to their indices and
-                # then back to strings
-                sp_toks = [
-                    tokenizer.convert_ids_to_tokens(
-                        tokenizer.convert_tokens_to_ids(sp_tok)
-                    )
-                    for sp_tok in tokenizer.special_tokens_map.values()
-                ]
-
-                # Replace special tokens with `None`
-                tokens_with_none = [None if tok in sp_toks else tok for tok in tokens]
-
-                # Get the alignment between the words and the tokens, on a character
-                # level
-                word_idxs = [
-                    word_idx for word_idx, word in enumerate(words) for _ in str(word)
-                ]
-                token_idxs = [
-                    tok_idx
-                    for tok_idx, tok_or_none in enumerate(tokens_with_none)
-                    for _ in str(tok_or_none)
-                    if tok_or_none is not None
-                ]
-                alignment = list(zip(word_idxs, token_idxs))
-
-                # Raise error if there are not as many characters in the words as in
-                # the tokens. This can be due to the use of a different prefix.
-                if len(word_idxs) != len(token_idxs):
-                    raise InvalidBenchmark(
-                        "The tokens could not be aligned with the words during manual "
-                        "word-token alignment. It seems that the tokenizer is neither "
-                        "of the fast variant nor of a SentencePiece/WordPiece variant."
-                    )
-
-                # Get the aligned word IDs
-                word_ids = list()
-                for tok_idx, tok_or_none in enumerate(tokens_with_none):
-                    if tok_or_none is None or tok_or_none == "":
-                        word_ids.append(None)
-                    else:
-                        word_idx = [
-                            word_idx
-                            for word_idx, token_idx in alignment
-                            if token_idx == tok_idx
-                        ][0]
-                        word_ids.append(word_idx)
-
-            previous_word_idx: int | None = None
-            label_ids: list[int] = list()
-            for word_id in word_ids:
-                # Special tokens have a word id that is None. We set the label to -100
-                # so they are automatically ignored in the loss function
-                if word_id is None:
-                    label_ids.append(-100)
-
-                # We set the label for the first token of each word
-                elif word_id != previous_word_idx:
-                    label = labels[word_id]
-                    try:
-                        label_id = label2id[label.lower()]
-                    except KeyError:
-                        msg = f"The label {label} was not found in the model's config."
-                        raise InvalidBenchmark(msg)
-                    label_ids.append(label_id)
-
-                # For the other tokens in a word, we set the label to -100
-                else:
-                    label_ids.append(-100)
-
-                previous_word_idx = word_id
-
-            all_labels.append(label_ids)
-        tokenized_inputs["labels"] = all_labels
-        return tokenized_inputs
-
-    def _handle_unk_tokens(
-        self, tokenizer: "Tokenizer", tokens: list[str], words: list[str]
-    ) -> list[str]:
-        """Replace unknown tokens in the tokens with the corresponding word.
-
-        Args:
-            tokenizer:
-                The tokenizer used to tokenize the words.
-            tokens:
-                The list of tokens.
-            words:
-                The list of words.
-
-        Returns:
-            The list of tokens with unknown tokens replaced by the corresponding word.
-        """
-        # Locate the token indices of the unknown tokens
-        token_unk_idxs = [
-            i for i, tok in enumerate(tokens) if tok == tokenizer.unk_token
-        ]
-
-        # Locate the word indices of the words which contain an unknown token
-        word_unk_idxs = [
-            i
-            for i, word in enumerate(words)
-            if tokenizer.unk_token
-            in tokenizer.convert_ids_to_tokens(
-                tokenizer.encode(word, add_special_tokens=False)
-            )
-        ]
-
-        # Iterate over the token index and word index pairs
-        for tok_idx, word_idx in zip(token_unk_idxs, word_unk_idxs):
-            # Fetch the word
-            word = words[word_idx]
-
-            # Tokenize the word, which is now a list containing at least one UNK token
-            tokens_with_unk = tokenizer.convert_ids_to_tokens(
-                tokenizer.encode(word, add_special_tokens=False)
-            )
-
-            # Iterate over the tokens in the word
-            for possible_unk_token in tokens_with_unk:
-                # If the token is not an UNK token then we remove the first occurence
-                # of the content of this token from the word. The result of the `word`
-                # variable will be the content of the UNK token.
-                # NOTE: This is a bit hacky and not bulletproof. For instance, if the
-                # word is "1925-1950" and the tokenizer splits it into ["[UNK]", "-",
-                # "19", "50"], then the result will be 2519 instead of 1925. This
-                # happens almost never, however, so we can live with it.
-                if possible_unk_token != tokenizer.unk_token:
-                    word = word.replace(possible_unk_token, "", 1)
-
-            # Replace the token with the word
-            tokens[tok_idx] = word
-
-        # Return the tokens
-        return tokens
-
-    def _preprocess_data(self, dataset: "Dataset", **kwargs) -> "Dataset":
-        """Preprocess a dataset by tokenizing and aligning the labels.
-
-        Args:
-            dataset:
-                The dataset to preprocess.
-            kwargs:
-                Extra keyword arguments containing objects used in preprocessing the
-                dataset.
-
-        Returns:
-            Hugging Face dataset:
-                The preprocessed dataset.
-        """
-        if kwargs["model_config"].task in GENERATIVE_MODEL_TASKS:
-            if "few_shot_examples" in kwargs:
-                few_shot_examples = kwargs["few_shot_examples"]
-                few_shot_fn = partial(
-                    self._apply_few_shot_prompt, few_shot_examples=few_shot_examples
-                )
-                dataset = dataset.map(
-                    few_shot_fn,
-                    batched=True,
-                    load_from_cache_file=False,
-                    keep_in_memory=True,
-                )
-
-            def tokenise(examples: dict) -> "BatchEncoding":
-                return kwargs["tokenizer"](
-                    text=examples["text"], truncation=True, padding=False
-                )
-
-            tokenised_dataset = dataset.map(
-                tokenise, batched=True, load_from_cache_file=False, keep_in_memory=True
-            )
-
-        else:
-            map_fn = partial(
-                self._tokenize_and_align_labels,
-                tokenizer=kwargs["tokenizer"],
-                label2id=kwargs["hf_model_config"].label2id,
-            )
-            tokenised_dataset = dataset.map(
-                map_fn, batched=True, load_from_cache_file=False, keep_in_memory=True
-            )
-
-        return tokenised_dataset
-
-    def _load_data_collator(
-        self,
-        tokenizer: "Tokenizer | None" = None,
-        model: "PreTrainedModel | GenerativeModel | None" = None,
-    ):
-        """Load the data collator used to prepare samples during finetuning.
-
-        Args:
-            tokenizer:
-                A pretrained tokenizer. Can be None if the tokenizer is not used in the
-                initialisation of the data collator. Defaults to None.
-            model:
-                A pretrained model. Can be None if the model is not used in the
-                initialisation of the data collator. Defaults to None.
-
-        Returns:
-            Hugging Face data collator:
-                The data collator.
-        """
-        if model_is_generative(model=model):
-            return DataCollatorWithPadding(tokenizer=tokenizer)
-        else:
-            return DataCollatorForTokenClassification(
-                tokenizer=tokenizer, label_pad_token_id=-100
-            )
-
-    def _extract_few_shot_examples(
-        self, train_dataset: "Dataset", random_seed: int
-    ) -> list[dict[str, Any]]:
-        """Extract few-shot examples from the training dataset.
-
-        Args:
-            train_dataset:
-                The training dataset.
-            random_seed:
-                The random seed to use when extracting the few-shot examples.
-
-        Returns:
-            list[dict[str, Any]]:
-                The few-shot examples.
-        """
-        shuffled_train = train_dataset.shuffle(seed=random_seed)
-        num_few_shots = self.dataset_config.num_few_shot_examples
-        labels = it.cycle(
-            [
-                label.lower()
-                for label in self.dataset_config.task.labels
-                if label.lower().startswith("b-")
-            ]
-        )
-        few_shot_examples: list[dict[str, Any]] = list()
-
-        # We pick the few-shot examples one at a time rather than all at once since
-        # we're working with a bootstrapped training dataset, meaning that it will have
-        # duplicates. This ensures that we don't have any duplicates in the few-shot
-        # examples
-        while len(few_shot_examples) < num_few_shots:
-            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["text"] != example["text"]
-            )
-
-        random.seed(random_seed)
-        random.shuffle(few_shot_examples)
-        return few_shot_examples
-
-    def _apply_few_shot_prompt(
-        self, examples: dict, few_shot_examples: list[dict], tokenizer: "Tokenizer"
-    ) -> dict:
-        """Apply a few-shot prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            few_shot_examples:
-                The examples to be included in the few-shot prompt.
-            tokenizer:
-                The tokenizer to use to encode the few-shot prompt.
-
-        Returns:
-            The examples with the few-shot prompt applied.
-        """
-        prompt_labels = self.dataset_config.prompt_label_mapping.values()
-
-        def create_label(example: dict) -> str:
-            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)
-
-        # Build the few-shot part of the prompt
-        few_shot_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=" ".join(example["tokens"]).replace("\n", " ").strip(),
-                label=create_label(example),
-            )
-            for example in few_shot_examples
-        ]
-        prompt_prefix = ""
-        if self.dataset_config.prompt_prefix:
-            prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"
-        few_shot_prompt = prompt_prefix + "\n\n".join(few_shot_prompts)
-
-        # Add the texts from the examples to the prompts
-        new_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=" ".join(tokens).replace("\n", " ").strip(), label=""
-            )
-            for tokens in examples["tokens"]
-        ]
-        examples["text"] = [
-            few_shot_prompt + "\n\n" + new_prompt for new_prompt in new_prompts
-        ]
-
-        return examples
-
-    def _apply_instruction_prompt(self, examples: dict, tokenizer: "Tokenizer") -> dict:
-        """Apply an instruction prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            tokenizer:
-                The tokenizer to use to encode the instruction prompt.
-
-        Returns:
-            The examples with the instruction prompt applied.
-        """
-        prompts = [
-            self.dataset_config.instruction_prompt.format(
-                text=re.sub(r"\n+", "\n", text).strip()
-            )
-            for text in examples["text"]
-        ]
-        prompts = [
-            convert_prompt_to_instruction(prompt=prompt, tokenizer=tokenizer)
-            for prompt in prompts
-        ]
-        examples["text"] = prompts
-        return examples
-
-    def _extract_labels_from_generation(
-        self,
-        input_batch: dict[str, list],
-        model_output: "ModelOutput",
-        tokenizer: "Tokenizer",
-    ) -> list[Any]:
-        """Extract the predicted labels from the generated output.
-
-        Args:
-            input_batch:
-                The input batch, where the keys are the feature names and the values
-                are lists with the feature values.
-            model_output:
-                The raw generated output of the model.
-            tokenizer:
-                The tokenizer used together with the model.
-
-        Returns:
-            list:
-                The predicted labels.
-        """
-        if importlib.util.find_spec("demjson3") is None:
-            raise NeedsExtraInstalled(extra="generative")
-
-        raw_predictions = extract_raw_predictions(
-            generated_sequences=model_output["sequences"], tokenizer=tokenizer
-        )
-
-        # Attempt to extract the JSON dictionary from the predictions
-        json_regex = r"\{.+?\}"
-        json_matches = [
-            re.search(pattern=json_regex, string=raw_prediction, flags=re.DOTALL)
-            or raw_prediction
-            for raw_prediction in raw_predictions
-        ]
-        raw_predictions = [
-            json_match.group() if isinstance(json_match, re.Match) else json_match
-            for json_match in json_matches
-        ]
-
-        tokens = input_batch["tokens"]
-        predicted_labels: list[list[str]] = [
-            ["o"] * len(token_ids) for token_ids in tokens
-        ]
-        for idx, raw_prediction in enumerate(raw_predictions):
-            try:
-                json_output = demjson3.decode(txt=raw_prediction)
-                if not isinstance(json_output, dict):
-                    logger.debug(
-                        "The model output is not a JSON dictionary, so cannot parse "
-                        f"it. Skipping. Here is the output: {raw_prediction}"
-                    )
-                    continue
-                elif not all(isinstance(key, str) for key in json_output.keys()):
-                    logger.debug(
-                        "The model output is not a JSON dictionary with string keys, "
-                        "so cannot parse it. Skipping. Here is the output: "
-                        f"{raw_prediction}"
-                    )
-                    continue
-                elif not all(isinstance(value, list) for value in json_output.values()):
-                    logger.debug(
-                        "The model output is not a JSON dictionary with list values, "
-                        "so cannot parse it. Skipping. Here is the output: "
-                        f"{raw_prediction}"
-                    )
-                    continue
-                prediction_dict: dict[str, list[str]] = json_output
-            except demjson3.JSONDecodeError:
-                logger.debug(
-                    "The model output is not valid JSON, so cannot parse it. Skipping. "
-                    f"Here is the output: {raw_prediction!r}"
-                )
-                continue
-
-            prompt_label_mapping = self.dataset_config.prompt_label_mapping
-            for prompt_tag_name, named_entities in prediction_dict.items():
-                try:
-                    tag_name = [
-                        tag[2:]
-                        for tag, prompt_tag in prompt_label_mapping.items()
-                        if prompt_tag == prompt_tag_name
-                    ][0]
-                except IndexError:
-                    logger.debug(
-                        "The model produced an invalid prompt tag name, "
-                        f"{prompt_tag_name}. Skipping."
-                    )
-                    continue
-
-                named_entities = [str(named_entity) for named_entity in named_entities]
-                for named_entity in named_entities:
-                    for ne_idx, named_entity_word in enumerate(named_entity.split()):
-                        for token_idx, token in enumerate(tokens[idx]):
-                            if named_entity_word in token:
-                                if ne_idx == 0:
-                                    predicted_labels[idx][token_idx] = f"b-{tag_name}"
-                                elif (
-                                    predicted_labels[idx][token_idx] == "o"
-                                    and predicted_labels[idx][token_idx - 1][2:]
-                                    == tag_name
-                                ):
-                                    predicted_labels[idx][token_idx] = f"i-{tag_name}"
-        return predicted_labels
diff --git a/src/scandeval/openai_models.py b/src/scandeval/openai_models.py
deleted file mode 100644
index 3202b6937..000000000
--- a/src/scandeval/openai_models.py
+++ /dev/null
@@ -1,613 +0,0 @@
-"""Model and tokenizer wrapper for OpenAI models."""
-
-import importlib.util
-import logging
-import math
-from typing import TYPE_CHECKING, Literal
-
-import torch
-from torch import LongTensor
-from torch.nn.utils.rnn import pad_sequence
-from transformers import BatchEncoding, GenerationConfig
-from transformers.modeling_utils import ModelOutput
-
-from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-from .exceptions import InvalidBenchmark, InvalidModel, NeedsExtraInstalled
-from .tasks import NER
-from .types import is_list_of_int, is_list_of_list_of_int
-
-if TYPE_CHECKING:
-    import tiktoken
-    from torch import Tensor
-    from transformers import PretrainedConfig
-
-if importlib.util.find_spec("tiktoken") is not None:
-    import tiktoken
-
-if importlib.util.find_spec("openai") is not None:
-    from openai import AzureOpenAI, BadRequestError, NotFoundError, OpenAI
-    from openai.types.chat import ChatCompletionUserMessageParam
-    from openai.types.chat.chat_completion_token_logprob import TopLogprob
-    from openai.types.shared.response_format_json_object import ResponseFormatJSONObject
-
-logger = logging.getLogger(__package__)
-
-
-class OpenAITokenizer:
-    """An OpenAI tokenizer.
-
-    Attributes:
-        model_config:
-            The model configuration.
-        hf_model_config:
-            The Hugging Face model configuration.
-        encoding:
-            The encoding.
-    """
-
-    unk_token: str | None = "<unk>"
-    unk_token_id = -1
-    pad_token: str | None = "<pad>"
-    padding_side: Literal["right", "left"] = "left"
-    is_fast = False
-    chat_template: str | None = None
-
-    def __init__(
-        self, model_config: ModelConfig, hf_model_config: "PretrainedConfig"
-    ) -> None:
-        """Initialize the tokenizer.
-
-        Args:
-            model_config:
-                The model configuration.
-            hf_model_config:
-                The Hugging Face model configuration.
-        """
-        self.model_config = model_config
-        self.hf_model_config = hf_model_config
-        self.model_max_length = self.hf_model_config.model_max_length
-
-        self.bos_token_id: int = self.hf_model_config.bos_token_id or -1
-        self.cls_token_id: int = self.bos_token_id
-        self.eos_token_id: int = self.hf_model_config.eos_token_id or -1
-        self.sep_token_id: int = self.eos_token_id
-        self.pad_token_id: int = self.hf_model_config.pad_token_id or -1
-
-        self.bos_token: str | None = self.encoding.decode([self.bos_token_id])
-        self.cls_token: str | None = self.bos_token
-        self.eos_token: str | None = self.encoding.decode([self.eos_token_id])
-        self.sep_token: str | None = self.eos_token
-
-        assert self.bos_token is not None
-        assert self.cls_token is not None
-        assert self.eos_token is not None
-        assert self.sep_token is not None
-        assert self.pad_token is not None
-
-        self.vocab_size = self.encoding.max_token_value + 1
-        self.special_tokens_map: dict[str, str] = dict(
-            bos_token=self.bos_token,
-            cls_token=self.cls_token,
-            eos_token=self.eos_token,
-            sep_token=self.sep_token,
-            pad_token=self.pad_token,
-        )
-
-    @property
-    def encoding(self) -> "tiktoken.Encoding":
-        """Return the underlying tiktoken encoding."""
-        try:
-            return tiktoken.encoding_for_model(model_name=self.model_config.model_id)
-        except KeyError:
-            # For Azure, the model_id is the deployment name. I do not know how to
-            # dynamically get the currently deployed model so assuming Azure only
-            # supports the latest models.
-            # TODO: This needs to be fixed, as gpt-4o uses a different encoding
-            return tiktoken.get_encoding("cl100k_base")
-
-    def __call__(self, text: str | list[str], **kwargs) -> BatchEncoding:
-        """Tokenize text.
-
-        Args:
-            text:
-                The text to tokenize.
-            **kwargs:
-                Additional keyword arguments.
-
-        Returns:
-            The tokenized text.
-        """
-        truncation = kwargs.get("truncation", False)
-        start_idx = -self.model_max_length if truncation else 0
-
-        text_list = [text] if isinstance(text, str) else text
-        encoded_inputs = [
-            BatchEncoding(
-                dict(
-                    input_ids=self.encoding.encode(
-                        text,
-                        allowed_special={
-                            self.bos_token,
-                            self.eos_token,
-                            self.cls_token,
-                            self.sep_token,
-                            self.pad_token,
-                        },
-                    )[start_idx:]
-                )
-            )
-            for text in text_list
-        ]
-        return self.pad(encoded_inputs=encoded_inputs)
-
-    def decode(self, token_ids: list[int], **kwargs) -> str:
-        """Decode token IDs.
-
-        Args:
-            token_ids:
-                The token IDs to decode.
-            **kwargs:
-                Additional keyword arguments.
-
-        Returns:
-            The decoded text.
-        """
-        token_ids = [
-            token_id for token_id in token_ids if token_id != self.pad_token_id
-        ]
-        return self.encoding.decode(tokens=token_ids)
-
-    def batch_decode(self, sequences: list[list[int]], **kwargs) -> list[str]:
-        """Decode batched token IDs.
-
-        Args:
-            sequences:
-                The token IDs to decode.
-            **kwargs:
-                Additional keyword arguments.
-
-        Returns:
-            The decoded text.
-        """
-        return [self.decode(token_ids=sequence) for sequence in sequences]
-
-    def encode(self, text: str, **kwargs) -> list[int]:
-        """Encode text.
-
-        Args:
-            text:
-                The text to encode.
-            **kwargs:
-                Additional keyword arguments.
-
-        Returns:
-            The encoded text.
-        """
-        return self(text, **kwargs).input_ids.tolist()[0]
-
-    def convert_ids_to_tokens(
-        self, ids: int | list[int], skip_special_tokens: bool = False
-    ) -> str | list[str]:
-        """Convert token IDs to tokens.
-
-        Args:
-            ids:
-                The token IDs to convert.
-            skip_special_tokens:
-                Whether to skip special tokens. Defaults to False.
-
-        Returns:
-            The tokens.
-        """
-        ids_list = [ids] if isinstance(ids, int) else ids
-        tokens = [self.decode(token_ids=[i]) for i in ids_list]
-        if skip_special_tokens:
-            tokens = [
-                token
-                for token in tokens
-                if token not in self.special_tokens_map.values()
-            ]
-        if len(tokens) == 1:
-            return tokens[0]
-        return tokens
-
-    def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
-        """Convert tokens to token IDs.
-
-        Args:
-            tokens:
-                The tokens to convert.
-
-        Returns:
-            The token IDs.
-        """
-        if isinstance(tokens, str):
-            tokens = [tokens]
-        ids = [self.encode(text=token)[0] for token in tokens]
-        if len(ids) == 1:
-            return ids[0]
-        return ids
-
-    def pad(
-        self,
-        encoded_inputs: (
-            BatchEncoding
-            | list[BatchEncoding]
-            | dict[str, list[int]]
-            | dict[str, list[list[int]]]
-            | list[dict[str, list[int]]]
-        ),
-        **kwargs,
-    ) -> BatchEncoding:
-        """Pad encoded inputs.
-
-        Args:
-            encoded_inputs:
-                Tokenized inputs. Can represent one input (BatchEncoding or Dict[str,
-                List[int]]) or a batch of tokenized inputs (list of BatchEncoding,
-                Dict[str, List[List[int]]] or List[Dict[str, List[int]]]) so you can
-                use this method during preprocessing as well as in a PyTorch Dataloader
-                collate function.
-            **kwargs:
-                Additional keyword arguments.
-        """
-        # Single example
-        if isinstance(encoded_inputs, BatchEncoding):
-            return encoded_inputs
-        elif isinstance(encoded_inputs, dict) and is_list_of_int(
-            encoded_inputs["input_ids"]
-        ):
-            return BatchEncoding(data=encoded_inputs)
-
-        # Batch of examples
-        if isinstance(encoded_inputs, dict) and is_list_of_list_of_int(
-            encoded_inputs["input_ids"]
-        ):
-            input_ids = encoded_inputs["input_ids"]
-        else:
-            assert isinstance(encoded_inputs, list)
-            input_ids = [list(example["input_ids"]) for example in encoded_inputs]
-
-        # Flip the token IDs in the lists, since `pad_sequence` pads to the right by
-        # default, and we want padding to the left
-        flipped_input_ids: list["Tensor"] = []
-        for input_id_list in input_ids:
-            input_id_list.reverse()
-            flipped_input_ids.append(LongTensor(input_id_list))
-
-        padded_input_ids = (
-            pad_sequence(
-                sequences=flipped_input_ids,
-                batch_first=True,
-                padding_value=self.pad_token_id,
-            )
-            .flip(dims=[1])
-            .long()
-        )
-
-        return BatchEncoding(dict(input_ids=padded_input_ids))
-
-    def apply_chat_template(
-        self, conversation: list[dict[Literal["role", "content"], str]], **kwargs
-    ) -> str | list[int]:
-        """Apply a chat template to a conversation.
-
-        Args:
-            conversation:
-                The conversation to apply the chat template to.
-            **kwargs:
-                Additional keyword arguments.
-
-        Returns:
-            The chat template.
-        """
-        raise NotImplementedError("Chat templates are not supported for OpenAI models.")
-
-
-class OpenAIModel:
-    """An OpenAI model.
-
-    Attributes:
-        model_config:
-            The model configuration.
-        config:
-            The Hugging Face model configuration.
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
-        tokenizer:
-            The tokenizer.
-        device:
-            The device to use, is always CPU.
-        is_chat_model:
-            Whether the model is a chat model.
-    """
-
-    def __init__(
-        self,
-        model_config: ModelConfig,
-        hf_model_config: "PretrainedConfig",
-        dataset_config: DatasetConfig,
-        benchmark_config: BenchmarkConfig,
-        tokenizer: OpenAITokenizer,
-    ) -> None:
-        """Initialize the model.
-
-        Args:
-            model_config:
-                The model configuration.
-            hf_model_config:
-                The Hugging Face model configuration.
-            dataset_config:
-                The dataset configuration.
-            benchmark_config:
-                The benchmark configuration.
-            tokenizer:
-                The tokenizer.
-        """
-        if importlib.util.find_spec("openai") is None:
-            raise NeedsExtraInstalled(extra="openai")
-
-        self.model_config = model_config
-        self.config = hf_model_config
-        self.dataset_config = dataset_config
-        self.benchmark_config = benchmark_config
-        self.tokenizer = tokenizer
-        self.device = torch.device("cpu")
-        self.client = self._initialize_openai_client()
-        self.is_chat_model = self._is_chat_model()
-        self.supports_json_mode = self._supports_json_mode()
-
-    def _initialize_openai_client(self) -> "OpenAI | AzureOpenAI":
-        """Initialize and return the OpenAI client.
-
-        Returns:
-            The OpenAI client.
-
-        Raises:
-            InvalidBenchmark:
-                If the OpenAI API key is not specified.
-        """
-        if self.benchmark_config.openai_api_key is not None:
-            return OpenAI(api_key=self.benchmark_config.openai_api_key, max_retries=60)
-        elif self.benchmark_config.azure_openai_api_key is not None:
-            if (
-                self.benchmark_config.azure_openai_endpoint is None
-                or self.benchmark_config.azure_openai_api_version is None
-            ):
-                if self.benchmark_config.run_with_cli:
-                    argument_message = (
-                        "`--azure-openai-endpoint` and `--azure_openai_api_version` "
-                        "flags"
-                    )
-                else:
-                    argument_message = (
-                        "`azure_openai_endpoint` and `azure_openai_api_version` "
-                        "arguments in the `Benchmarker` class"
-                    )
-                raise InvalidModel(
-                    "Azure OpenAI models require an endpoint and API version to be "
-                    "specified. Please specify the endpoint with the "
-                    f"{argument_message}, or specify the environment variables "
-                    "`AZURE_OPENAI_ENDPOINT` and `AZURE_OPENAI_API_VERSION`."
-                )
-            return AzureOpenAI(
-                api_key=self.benchmark_config.azure_openai_api_key,
-                azure_endpoint=self.benchmark_config.azure_openai_endpoint,
-                api_version=self.benchmark_config.azure_openai_api_version,
-                max_retries=60,
-            )
-        elif self.benchmark_config.run_with_cli:
-            raise InvalidModel(
-                "OpenAI models require an API key to be specified. Please specify the "
-                "`--openai-api-key` argument (or `--azure-openai-api-key` and "
-                "`--azure-openai-endpoint` arguments) or specify the environment "
-                "variables `OPENAI_API_KEY` (or `AZURE_OPENAI_API_KEY` and "
-                "`AZURE_OPENAI_ENDPOINT`)."
-            )
-        else:
-            raise InvalidModel(
-                "OpenAI models require an API key to be specified. Please specify the "
-                "`openai_api_key` argument to the `Benchmarker` class (or "
-                "`azure_openai_api_key` and `azure_openai_endpoint` arguments) or "
-                "specify the environment variables `OPENAI_API_KEY` (or "
-                "`AZURE_OPENAI_API_KEY` and `AZURE_OPENAI_ENDPOINT`)."
-            )
-
-    def _is_chat_model(self) -> bool:
-        """Returns whether the model is a chat model."""
-        try:
-            self.client.completions.create(
-                model=self.model_config.model_id, prompt="Test", max_tokens=1
-            )
-            return False
-        except (NotFoundError, BadRequestError) as e:
-            chat_model_strings = [
-                "This is a chat model",
-                "The completion operation does not work with the specified model",
-            ]
-            if any(string in str(e) for string in chat_model_strings):
-                return True
-            raise e
-
-    def _supports_json_mode(self) -> bool:
-        """Returns whether the model supports JSON mode."""
-        if not self.is_chat_model:
-            return False
-
-        try:
-            self.client.chat.completions.create(  # type: ignore[call-overload]
-                model=self.model_config.model_id,
-                messages=[
-                    ChatCompletionUserMessageParam(role="user", content="Test json")
-                ],
-                max_tokens=1,
-                response_format=ResponseFormatJSONObject(type="json_object"),
-            )
-            return True
-        except BadRequestError as e:
-            no_json_mode_strings = ["not supported with this model"]
-            if any(string in str(e) for string in no_json_mode_strings):
-                return False
-            raise e
-
-    def generate(
-        self,
-        inputs: "Tensor",
-        generation_config: GenerationConfig | None = None,
-        **generation_kwargs,
-    ) -> "Tensor | LongTensor | ModelOutput":
-        """Generate text using the model.
-
-        Args:
-            inputs:
-                The input IDs, of shape (batch_size, sequence_length).
-            generation_config:
-                The generation configuration. If None then a default GenerationConfig
-                will be used. Defaults to None.
-            **generation_kwargs:
-                Additional keyword arguments. Can also be used to override
-                generation configuration.
-
-        Returns:
-            The model output.
-        """
-        if generation_config is None:
-            generation_config = GenerationConfig(**generation_kwargs)
-        else:
-            for key, value in generation_kwargs.items():
-                setattr(generation_config, key, value)
-
-        multiple_inputs = inputs.size(dim=0) > 1
-        if multiple_inputs:
-            raise ValueError(
-                "OpenAI models do not support multiple inputs. Please use a batch "
-                "size of 1."
-            )
-
-        two_dimensional_input = len(inputs.size()) == 2
-        if two_dimensional_input:
-            inputs = inputs[0]
-
-        prompt = self.tokenizer.decode(
-            [
-                token_id
-                for token_id in inputs.tolist()
-                if token_id != self.config.pad_token_id
-            ]
-        )
-
-        model_id = self.model_config.model_id
-        max_tokens: int = generation_config.max_new_tokens or 1
-        temperature = (
-            0.0 if not generation_config.do_sample else generation_config.temperature
-        )
-        generation_kwargs = dict(
-            model=model_id,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=generation_config.top_p,
-            n=generation_config.num_return_sequences,
-            frequency_penalty=generation_config.repetition_penalty - 1.0,
-            stop=["\n\n", self.tokenizer.eos_token, self.tokenizer.pad_token],
-            seed=4242,
-        )
-
-        if generation_config.output_scores:
-            generation_kwargs["logprobs"] = True
-            generation_kwargs["top_logprobs"] = 10
-
-        if (
-            self.dataset_config.task == NER
-            and self.supports_json_mode
-            and "json" in prompt.lower()
-        ):
-            generation_kwargs["response_format"] = dict(type="json_object")
-
-        try:
-            if not self.is_chat_model:
-                model_output = self.client.completions.create(
-                    prompt=prompt, **generation_kwargs
-                )
-                generation_output = model_output.choices[0].text.strip()
-            else:
-                model_output = self.client.chat.completions.create(
-                    messages=[dict(role="system", content=prompt)], **generation_kwargs
-                )
-                generation_output = model_output.choices[0].message.content.strip()
-        except BadRequestError as e:
-            raise InvalidBenchmark(
-                f"Encountered error during OpenAI generation: {str(e)!r}\n"
-                f"The prompt causing it was {prompt!r}."
-            )
-
-        completion_ids = self.tokenizer([generation_output]).input_ids.tolist()
-        output = LongTensor(completion_ids)
-
-        if generation_config.return_dict_in_generate:
-            output_dict: dict[str, Tensor | tuple] = dict(sequences=output)
-
-            # Add logprobs scores to the output
-            if generation_config.output_scores:
-                # Create a list with placeholder logprobs for every token generated.
-                # Each tensor in the list will be of shape (batch_size, vocab_size)
-                batch_size = 1
-                vocab_size = self.config.vocab_size
-                logprobs_list = model_output.choices[0].logprobs.content
-                seq_len = len(logprobs_list)
-                scores = [
-                    torch.full(size=(batch_size, vocab_size), fill_value=-math.inf)
-                    for _ in range(seq_len)
-                ]
-
-                # Fill in the logprobs for each generated token. The logprobs from the
-                # OpenAI output only contain the logprobs for the top-k tokens, so we
-                # only fill in these and leave the rest at ~0% probability
-                for gen_token_idx, logprobs in enumerate(logprobs_list):
-                    # Remove duplicate logprobs, since, e.g., "negative" and " negative"
-                    # can both be included in the list, and we only want the one with
-                    # the highest logprob
-                    top_logprobs: list[TopLogprob] = list()
-                    for logprob_obj in logprobs.top_logprobs:
-                        token = logprob_obj.token.strip()
-                        if token in {lg.token.strip() for lg in top_logprobs}:
-                            continue
-                        top_logprobs.append(logprob_obj)
-
-                    for logprob_obj in top_logprobs:
-                        logprob = logprob_obj.logprob
-                        token = logprob_obj.token.strip()
-                        if token:
-                            token_idx = self.tokenizer.encode(text=token)[0]
-                            scores[gen_token_idx][0, token_idx] = logprob
-
-                output_dict["scores"] = tuple(scores)
-
-            return ModelOutput(output_dict)
-
-        return output
-
-    def to(self, device: torch.device) -> "OpenAIModel":
-        """Move the model to a device.
-
-        Args:
-            device:
-                The device to move the model to.
-
-        Returns:
-            The model.
-        """
-        return self
-
-    def eval(self) -> "OpenAIModel":
-        """Put the model in evaluation mode.
-
-        Returns:
-            The model.
-        """
-        return self
-
-    def __call__(self, *args, **kwargs):
-        """Call the model."""
-        raise NotImplementedError("OpenAI models do not support the `__call__` method.")
diff --git a/src/scandeval/protocols.py b/src/scandeval/protocols.py
deleted file mode 100644
index 15a0e48cf..000000000
--- a/src/scandeval/protocols.py
+++ /dev/null
@@ -1,282 +0,0 @@
-"""Protocols used throughout the project."""
-
-from typing import TYPE_CHECKING, Literal, Protocol, Union, runtime_checkable
-
-if TYPE_CHECKING:
-    import torch
-    from transformers import (
-        BatchEncoding,
-        GenerationConfig,
-        PretrainedConfig,
-        PreTrainedModel,
-    )
-    from transformers.utils import ModelOutput
-
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-
-
-@runtime_checkable
-class Tokenizer(Protocol):
-    """A protocol for a tokenizer."""
-
-    cls_token: str | None
-    sep_token: str | None
-    bos_token: str | None
-    eos_token: str | None
-    pad_token: str | None
-    unk_token: str | None
-    cls_token_id: int
-    sep_token_id: int
-    bos_token_id: int
-    eos_token_id: int
-    pad_token_id: int
-    unk_token_id: int
-    is_fast: bool
-    chat_template: str | None
-    special_tokens_map: dict[str, str]
-    model_max_length: int
-    vocab_size: int
-    padding_side: Literal["right", "left"]
-
-    def __call__(self, text: str | list[str], **kwargs) -> "BatchEncoding":
-        """Call the tokenizer.
-
-        Args:
-            text:
-                The text to tokenize.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The encoded inputs.
-        """
-        ...
-
-    def decode(self, token_ids: list[int], **kwargs) -> str:
-        """Decode a list of token IDs.
-
-        Args:
-            token_ids:
-                The token IDs to decode.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The decoded string.
-        """
-        ...
-
-    def batch_decode(self, sequences: list[list[int]], **kwargs) -> list[str]:
-        """Decode a batch of token IDs.
-
-        Args:
-            sequences:
-                The token IDs to decode.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The decoded strings.
-        """
-        ...
-
-    def encode(self, text: str, **kwargs) -> list[int]:
-        """Encode one or more texts.
-
-        Args:
-            text:
-                The text to encode.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The encoded token IDs.
-        """
-        ...
-
-    def convert_ids_to_tokens(
-        self, ids: int | list[int], skip_special_tokens: bool = False
-    ) -> str | list[str]:
-        """Convert a list of token IDs to tokens.
-
-        Args:
-            ids:
-                The token IDs to convert.
-            skip_special_tokens:
-                Whether to skip special tokens.
-
-        Returns:
-            The tokens.
-        """
-        ...
-
-    def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
-        """Convert a list of tokens to token IDs.
-
-        Args:
-            tokens:
-                The tokens to convert.
-
-        Returns:
-            The token IDs.
-        """
-        ...
-
-    def pad(
-        self,
-        encoded_inputs: Union[
-            "BatchEncoding",
-            list["BatchEncoding"],
-            dict[str, list[int]],
-            dict[str, list[list[int]]],
-            list[dict[str, list[int]]],
-        ],
-        **kwargs,
-    ) -> "BatchEncoding":
-        """Pad a batch of encoded inputs.
-
-        Args:
-            encoded_inputs:
-                The encoded inputs to pad.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The padded encoded inputs.
-        """
-        ...
-
-    def apply_chat_template(
-        self, conversation: list[dict[Literal["role", "content"], str]], **kwargs
-    ) -> str | list[int]:
-        """Apply a chat template to a conversation.
-
-        Args:
-            conversation:
-                The conversation.
-            **kwargs:
-                Keyword arguments to pass to the tokenizer.
-
-        Returns:
-            The conversation as a string.
-        """
-        ...
-
-
-@runtime_checkable
-class GenerativeModel(Protocol):
-    """A protocol for a generative model."""
-
-    config: "PretrainedConfig"
-    device: "torch.device"
-
-    def to(self, device: "torch.device") -> "GenerativeModel":
-        """Move the model to a device.
-
-        Args:
-            device:
-                The device to move the model to.
-
-        Returns:
-            The model.
-        """
-        ...
-
-    def eval(self) -> "GenerativeModel":
-        """Put the model in evaluation mode.
-
-        Returns:
-            The model.
-        """
-        ...
-
-    def generate(
-        self,
-        inputs: "torch.Tensor",
-        generation_config: "GenerationConfig | None" = None,
-        **generation_kwargs,
-    ) -> "ModelOutput | torch.Tensor":
-        """Generate text.
-
-        Args:
-            inputs:
-                The input IDs.
-            generation_config:
-                The generation configuration.
-            **generation_kwargs:
-                Keyword arguments to pass to the generation method.
-
-        Returns:
-            The generated text.
-        """
-        ...
-
-    def __call__(
-        self, input_ids: "torch.Tensor", labels: "torch.Tensor | None" = None
-    ) -> "ModelOutput":
-        """Call the model.
-
-        Args:
-            input_ids:
-                The input IDs.
-            labels:
-                The labels.
-
-        Returns:
-            The model output.
-        """
-        ...
-
-
-class ModelSetup(Protocol):
-    """A protocol for a general model setup."""
-
-    def __init__(self, benchmark_config: "BenchmarkConfig") -> None:
-        """Initialize the model setup.
-
-        Args:
-            benchmark_config:
-                The benchmark configuration.
-        """
-        ...
-
-    def model_exists(self, model_id: str) -> bool | dict[str, str]:
-        """Check whether a model exists.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            Whether the model exist, or a dictionary explaining why we cannot check
-            whether the model exists.
-        """
-        ...
-
-    def get_model_config(self, model_id: str) -> "ModelConfig":
-        """Get the model configuration.
-
-        Args:
-            model_id:
-                The model ID.
-
-        Returns:
-            The model configuration.
-        """
-        ...
-
-    def load_model(
-        self, model_config: "ModelConfig", dataset_config: "DatasetConfig"
-    ) -> tuple["PreTrainedModel | GenerativeModel", "Tokenizer"]:
-        """Load a model.
-
-        Args:
-            model_config:
-                The model configuration.
-            dataset_config:
-                The dataset configuration.
-
-        Returns:
-            The tokenizer and model.
-        """
-        ...
diff --git a/src/scandeval/question_answering.py b/src/scandeval/question_answering.py
deleted file mode 100644
index d5905a523..000000000
--- a/src/scandeval/question_answering.py
+++ /dev/null
@@ -1,606 +0,0 @@
-"""Question-answering benchmark dataset."""
-
-import logging
-import re
-from functools import partial
-from typing import TYPE_CHECKING, Any, Type
-
-import numpy as np
-from evaluate import EvaluationModule
-from transformers.data.data_collator import DataCollatorWithPadding
-
-from .benchmark_dataset import BenchmarkDataset
-from .exceptions import InvalidBenchmark
-from .generation import extract_raw_predictions
-from .question_answering_trainer import QuestionAnsweringTrainer
-from .utils import (
-    convert_prompt_to_instruction,
-    get_special_token_metadata,
-    raise_if_model_output_contains_nan_values,
-)
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-    from transformers.modeling_utils import ModelOutput, PreTrainedModel
-    from transformers.tokenization_utils_base import BatchEncoding
-    from transformers.trainer import Trainer
-
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Labels, Predictions
-
-logger = logging.getLogger(__package__)
-
-
-class QuestionAnswering(BenchmarkDataset):
-    """Question-answering benchmark dataset.
-
-    Args:
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
-
-    Attributes:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-    """
-
-    def _preprocess_data(self, dataset: "Dataset", **kwargs) -> "Dataset":
-        """Preprocess a dataset by tokenizing and aligning the labels.
-
-        Args:
-            dataset:
-                The dataset to preprocess.
-            kwargs:
-                Extra keyword arguments containing objects used in preprocessing the
-                dataset.
-
-        Returns:
-            The preprocessed dataset.
-        """
-        split: str = kwargs.pop("split")
-        tokenizer: "Tokenizer" = kwargs.pop("tokenizer")
-        generative_model: bool = kwargs.pop("generative_model")
-
-        # If the tokenizer is not a fast variant then raise an error
-        if not tokenizer.is_fast and not generative_model:
-            raise InvalidBenchmark(
-                "Question-answering benchmarks require a fast tokenizer."
-            )
-
-        if generative_model:
-            preprocess_fn = partial(
-                prepare_examples_for_generation, tokenizer=tokenizer
-            )
-        elif split == "test":
-            preprocess_fn = partial(prepare_test_examples, tokenizer=tokenizer)
-        else:
-            preprocess_fn = partial(prepare_train_examples, tokenizer=tokenizer)
-
-        # Preprocess the data and return it
-        try:
-            if generative_model:
-                cols_to_remove = [
-                    col for col in dataset.column_names if col not in ["id", "text"]
-                ]
-            else:
-                cols_to_remove = dataset.column_names
-            preprocessed = dataset.map(
-                preprocess_fn,
-                batched=True,
-                batch_size=10,
-                remove_columns=cols_to_remove,
-                load_from_cache_file=False,
-                keep_in_memory=True,
-            )
-
-        except NotImplementedError as e:
-            raise InvalidBenchmark(str(e))
-
-        # The Trainer hides the columns that are not used by the model (here `id` and
-        # `offset_mapping` which we will need for our post-processing), so we put them
-        # back
-        preprocessed.set_format(
-            type=preprocessed.format["type"], columns=list(preprocessed.features.keys())
-        )
-
-        # Return the preprocessed dataset
-        return preprocessed
-
-    def _get_trainer_class(self) -> Type["Trainer"]:
-        return QuestionAnsweringTrainer
-
-    def _get_evaluate_inputs(
-        self, dataset: "Dataset", prepared_dataset: "Dataset", metric_key_prefix: str
-    ):
-        return dict(
-            orig_eval_dataset=dataset,
-            eval_dataset=prepared_dataset,
-            metric_key_prefix=metric_key_prefix,
-        )
-
-    def _load_data_collator(
-        self,
-        tokenizer: "Tokenizer | None" = None,
-        model: "PreTrainedModel | GenerativeModel | None" = None,
-    ):
-        """Load the data collator used to prepare samples during finetuning.
-
-        Args:
-            tokenizer:
-                A pretrained tokenizer. Can be None if the tokenizer is not used in the
-                initialisation of the data collator. Defaults to None.
-            model:
-                A pretrained model. Can be None if the model is not used in the
-                initialisation of the data collator. Defaults to None.
-
-        Returns:
-            The data collator.
-        """
-        return DataCollatorWithPadding(tokenizer=tokenizer)
-
-    def _compute_metrics(
-        self,
-        model_outputs_and_labels: tuple["Predictions", "Labels"],
-        id2label: dict[int, str],
-    ) -> dict[str, float]:
-        """Compute the metrics needed for evaluation.
-
-        Args:
-            model_outputs_and_labels:
-                The first sequence contains the model outputs and the second sequence
-                contains the true labels.
-            id2label:
-                Conversion of indices to labels.
-
-        Returns:
-            A dictionary with the names of the metrics as keys and the metric values as
-            values.
-        """
-        model_outputs, labels = model_outputs_and_labels
-
-        raise_if_model_output_contains_nan_values(model_output=model_outputs)
-
-        model_output_dtype = np.asarray(model_outputs).dtype
-        if model_output_dtype in [np.float16, np.float32, np.float64]:
-            predictions = np.asarray(model_outputs).argmax(axis=-1)
-        else:
-            predictions = model_outputs
-
-        results: dict[str, float] = dict()
-        for cfg in self.dataset_config.task.metrics:
-            metric = self._metrics[cfg.name]
-            assert isinstance(metric, EvaluationModule)
-            score_dict: dict[str, float] | None = metric.compute(
-                predictions=predictions, references=labels, **cfg.compute_kwargs
-            )
-
-            # The metric returns None if we are running on multi-GPU and the current
-            # process is not the main process
-            if score_dict is not None:
-                scores = score_dict[cfg.results_key]
-                if isinstance(scores, list):
-                    scores = sum(scores) / len(scores)
-                results[cfg.name] = scores
-
-        return results
-
-    def _extract_few_shot_examples(
-        self, train_dataset: "Dataset", random_seed: int
-    ) -> list[dict[str, Any]]:
-        """Extract few-shot examples from the training dataset.
-
-        Args:
-            train_dataset:
-                The training dataset.
-            random_seed:
-                The random seed to use when extracting the few-shot examples.
-
-        Returns:
-            The few-shot examples.
-        """
-        # Locate the maximum number of tokens that constitutes a short example. We
-        # start with 512 tokens and double it until there is at least `num_few_shots`
-        # many examples that are shorter than the maximum number of tokens.
-        max_num_tokens = 512
-        while True:
-            train_with_short_examples = train_dataset.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
-            max_num_tokens *= 2
-
-        train_with_short_examples = train_dataset.filter(
-            lambda example: len(example["context"]) < max_num_tokens
-        )
-        shuffled_train = train_with_short_examples.shuffle(seed=random_seed)
-        num_few_shots = self.dataset_config.num_few_shot_examples
-        few_shot_examples: list[dict[str, Any]] = list()
-
-        # We pick the few-shot examples one at a time rather than all at once since
-        # we're working with a bootstrapped training dataset, meaning that it will have
-        # duplicates. This ensures that we don't have any duplicates in the few-shot
-        # examples
-        while len(few_shot_examples) < num_few_shots:
-            example = shuffled_train.select(range(1))[0]
-            few_shot_examples.append(example)
-            shuffled_train = shuffled_train.filter(
-                lambda x: x["context"] != example["context"]
-            )
-
-        return few_shot_examples
-
-    def _apply_few_shot_prompt(
-        self, examples: dict, few_shot_examples: list[dict], tokenizer: "Tokenizer"
-    ) -> dict:
-        """Apply a few-shot prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            few_shot_examples:
-                The examples to be included in the few-shot prompt.
-            tokenizer:
-                The tokenizer to use to encode the few-shot prompt.
-
-        Returns:
-            The examples with the few-shot prompt applied.
-        """
-        # Build the few-shot part of the prompt
-        few_shot_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=example["context"].replace("\n", " ").strip(),
-                question=example["question"].strip(),
-                label=example["answers"]["text"][0],
-            )
-            for example in few_shot_examples
-        ]
-        prompt_prefix = ""
-        if self.dataset_config.prompt_prefix:
-            prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"
-        few_shot_prompt = prompt_prefix + "\n\n".join(few_shot_prompts)
-
-        # Add the texts from the examples to the prompts
-        new_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=context.replace("\n", " ").strip(), question=question, label=""
-            )
-            for context, question in zip(examples["context"], examples["question"])
-        ]
-        examples["text"] = [
-            few_shot_prompt + "\n\n" + new_prompt for new_prompt in new_prompts
-        ]
-
-        return examples
-
-    def _apply_instruction_prompt(self, examples: dict, tokenizer: "Tokenizer") -> dict:
-        """Apply an instruction prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            tokenizer:
-                The tokenizer to use to encode the instruction prompt.
-
-        Returns:
-            The examples with the instruction prompt applied.
-        """
-        prompts = [
-            self.dataset_config.instruction_prompt.format(
-                text=re.sub(r"\n+", "\n", text).strip(), question=question.strip()
-            )
-            for (text, question) in zip(examples["context"], examples["question"])
-        ]
-        prompts = [
-            convert_prompt_to_instruction(prompt=prompt, tokenizer=tokenizer)
-            for prompt in prompts
-        ]
-        examples["text"] = prompts
-        return examples
-
-    def _extract_labels_from_generation(
-        self,
-        input_batch: dict[str, list],
-        model_output: "ModelOutput",
-        tokenizer: "Tokenizer",
-    ) -> list[Any]:
-        """Extract the predicted labels from the generated output.
-
-        Args:
-            input_batch:
-                The input batch, where the keys are the feature names and the values
-                are lists with the feature values.
-            model_output:
-                The raw generated output of the model.
-            tokenizer:
-                The tokenizer used together with the model.
-
-        Returns:
-            The predicted labels.
-        """
-        raw_predictions = extract_raw_predictions(
-            generated_sequences=model_output["sequences"], tokenizer=tokenizer
-        )
-
-        predictions = [
-            dict(
-                id=id,
-                prediction_text=predicted_answer.lower(),
-                no_answer_probability=0.0,
-            )
-            for id, predicted_answer in zip(input_batch["id"], raw_predictions)
-        ]
-
-        return predictions
-
-
-def prepare_train_examples(
-    examples: "BatchEncoding", tokenizer: "Tokenizer"
-) -> "BatchEncoding":
-    """Prepare the features for training.
-
-    Args:
-        examples:
-            The examples to prepare.
-        tokenizer:
-            The tokenizer to use to prepare the examples.
-
-    Returns:
-        The prepared examples.
-    """
-    # Some of the questions have lots of whitespace on the left, which is not useful
-    # and will make the truncation of the context fail (the tokenized question will
-    # take a lots of space). So we remove that left whitespace
-    examples["question"] = [q.lstrip() for q in examples["question"]]
-
-    # Extract special token metadata from the tokenizer
-    special_token_metadata = get_special_token_metadata(tokenizer=tokenizer)
-    has_cls_token = special_token_metadata["has_cls_token"]
-    has_sep_token = special_token_metadata["has_sep_token"]
-    cls_token_id = special_token_metadata["cls_token_id"]
-    cls_token = special_token_metadata["cls_token"]
-    sep_token = special_token_metadata["sep_token"]
-
-    # If the tokenizer is not adding special tokens, then we add them manually
-    if not has_cls_token and not has_sep_token:
-        examples["question"] = [
-            f"{cls_token}{q}{sep_token}" for q in examples["question"]
-        ]
-        examples["context"] = [f"{c}{sep_token}" for c in examples["context"]]
-
-    # Set the stride used during tokenization, when the context is long enough to be
-    # split into several features. Since we are always keeping the question tokens, we
-    # need to make sure that the stride does not exceed the resulting maximum context
-    # length.
-    max_question_tokens = max(len(tokenizer(q).input_ids) for q in examples["question"])
-    num_special_tokens = int(has_cls_token) + int(has_sep_token)
-    stride = tokenizer.model_max_length // 4
-    max_length = tokenizer.model_max_length - stride
-    stride = min(stride, max_length - max_question_tokens - num_special_tokens)
-    max_length = tokenizer.model_max_length - stride
-
-    # Tokenize our examples with truncation and padding, but keep the overflows using a
-    # stride. This results in one example possible giving several features when a
-    # context is long, each of those features having a context that overlaps a bit the
-    # context of the previous feature.
-    tokenized_examples = tokenizer(
-        text=examples["question"],
-        text_pair=examples["context"],
-        truncation="only_second",
-        max_length=max_length,
-        stride=stride,
-        return_overflowing_tokens=True,
-        return_offsets_mapping=True,
-        padding="max_length",
-    )
-
-    # Since one example might give us several features if it has a long context, we
-    # need a map from a feature to its corresponding example. This key gives us just
-    # that
-    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
-
-    # The offset mappings will give us a map from token to character position in the
-    # original context. This will help us compute the start_positions and
-    # end_positions.
-    offset_mapping = tokenized_examples.pop("offset_mapping")
-
-    # Initialise the start- and end positions of the answers
-    tokenized_examples["start_positions"] = list()
-    tokenized_examples["end_positions"] = list()
-
-    for i, offsets in enumerate(offset_mapping):
-        # Get the input IDs for the current example
-        input_ids = tokenized_examples.input_ids[i]
-
-        # We will label impossible answers with the index of the CLS token
-        cls_index = input_ids.index(cls_token_id)
-
-        # Grab the sequence corresponding to that example (to know what is the context
-        # and what is the question).
-        sequence_ids = tokenized_examples.sequence_ids(i)
-
-        # Manually ensure that the special tokens are set to None in `sequence_ids`
-        for special_token in tokenizer.special_tokens_map.keys():
-            if hasattr(tokenizer, f"{special_token}_id"):
-                special_token_id = getattr(tokenizer, f"{special_token}_id")
-                if special_token_id is not None:
-                    sequence_ids = [
-                        None if token_id == special_token_id else seq_id
-                        for token_id, seq_id in zip(input_ids, sequence_ids)
-                    ]
-
-        # One example can give several spans, this is the index of the example
-        # containing this span of text.
-        sample_index = sample_mapping[i]
-        answers = examples["answers"][sample_index]
-
-        # If no answers are given, set the cls_index as answer.
-        if len(answers["answer_start"]) == 0:
-            tokenized_examples.start_positions.append(cls_index)
-            tokenized_examples.end_positions.append(cls_index)
-
-        else:
-            # Start/end character index of the answer in the text.
-            start_char = answers["answer_start"][0]
-            end_char = start_char + len(answers["text"][0])
-
-            # Start token index of the current span in the text.
-            token_start_index = 0
-            while sequence_ids[token_start_index] != 1:
-                token_start_index += 1
-
-            # End token index of the current span in the text.
-            token_end_index = len(input_ids) - 1
-            while sequence_ids[token_end_index] != 1:
-                token_end_index -= 1
-
-            # Detect if the answer is out of the span (in which case this feature is
-            # labeled with the CLS index).
-            if not (
-                offsets[token_start_index][0] <= start_char
-                and offsets[token_end_index][1] >= end_char
-            ):
-                tokenized_examples.start_positions.append(cls_index)
-                tokenized_examples.end_positions.append(cls_index)
-
-            # Otherwise move the token_start_index and token_end_index to the two ends
-            # of the answer. Note: we could go after the last offset if the answer is
-            # the last word (edge case).
-            else:
-                while (
-                    token_start_index <= token_end_index
-                    and offsets[token_start_index][0] <= start_char
-                ):
-                    token_start_index += 1
-                token_start_index -= 1
-                tokenized_examples.start_positions.append(token_start_index)
-                while (
-                    token_start_index <= token_end_index
-                    and offsets[token_end_index][1] >= end_char
-                ):
-                    token_end_index -= 1
-                token_end_index += 1
-                tokenized_examples.end_positions.append(token_end_index)
-                assert token_end_index >= token_start_index
-
-    return tokenized_examples
-
-
-def prepare_test_examples(
-    examples: "BatchEncoding", tokenizer: "Tokenizer"
-) -> "BatchEncoding":
-    """Prepare test examples.
-
-    Args:
-        examples:
-            Dictionary of test examples.
-        tokenizer:
-            The tokenizer used to preprocess the examples.
-
-    Returns:
-        The prepared test examples.
-    """
-    # Some of the questions have lots of whitespace on the left, which is not useful
-    # and will make the truncation of the context fail (the tokenized question will
-    # take a lots of space). So we remove that left whitespace
-    examples["question"] = [q.lstrip() for q in examples["question"]]
-
-    # Extract special token metadata from the tokenizer
-    special_token_metadata = get_special_token_metadata(tokenizer=tokenizer)
-    has_cls_token = special_token_metadata["has_cls_token"]
-    has_sep_token = special_token_metadata["has_sep_token"]
-    cls_token = special_token_metadata["cls_token"]
-    sep_token = special_token_metadata["sep_token"]
-
-    # If the tokenizer is not adding special tokens, then we add them manually
-    if not has_cls_token and not has_sep_token:
-        examples["question"] = [
-            f"{cls_token}{q}{sep_token}" for q in examples["question"]
-        ]
-        examples["context"] = [f"{c}{sep_token}" for c in examples["context"]]
-
-    # Set the stride used during tokenization, when the context is long enough to be
-    # split into several features. Since we are always keeping the question tokens, we
-    # need to make sure that the stride does not exceed the resulting maximum context
-    # length.
-    max_question_tokens = max(len(tokenizer(q).input_ids) for q in examples["question"])
-    num_special_tokens = int(has_cls_token) + int(has_sep_token)
-    stride = tokenizer.model_max_length // 4
-    max_length = tokenizer.model_max_length - stride
-    stride = min(stride, max_length - max_question_tokens - num_special_tokens)
-    max_length = tokenizer.model_max_length - stride
-
-    # Tokenize our examples with truncation and maybe padding, but keep the overflows
-    # using a stride. This results in one example possible giving several features when
-    # a context is long, each of those features having a context that overlaps a bit
-    # the context of the previous feature.
-    tokenized_examples = tokenizer(
-        text=examples["question"],
-        text_pair=examples["context"],
-        truncation="only_second",
-        max_length=max_length,
-        stride=stride,
-        return_overflowing_tokens=True,
-        return_offsets_mapping=True,
-        padding="max_length",
-    )
-
-    # Since one example might give us several features if it has a long context, we
-    # need a map from a feature to its corresponding example. This key gives us just
-    # that.
-    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
-
-    # We keep the id that gave us this feature and we will store the offset mappings.
-    tokenized_examples["id"] = list()
-
-    for i in range(len(tokenized_examples.input_ids)):
-        # Grab the sequence corresponding to that example (to know what is the context
-        # and what is the question).
-        sequence_ids = tokenized_examples.sequence_ids(i)
-        context_index = 1
-
-        # One example can give several spans, this is the index of the example
-        # containing this span of text.
-        sample_index = sample_mapping[i]
-        tokenized_examples.id.append(examples["id"][sample_index])
-
-        # Set to (-1, -1) the offset_mapping that are not part of the context so it's
-        # easy to determine if a token position is part of the context or not.
-        tokenized_examples.offset_mapping[i] = [
-            (o if sequence_ids[k] == context_index else (-1, -1))
-            for k, o in enumerate(tokenized_examples.offset_mapping[i])
-        ]
-
-    return tokenized_examples
-
-
-def prepare_examples_for_generation(
-    examples: "BatchEncoding", tokenizer: "Tokenizer"
-) -> "BatchEncoding":
-    """Prepare test examples.
-
-    Args:
-        examples:
-            Dictionary of test examples.
-        tokenizer:
-            The tokenizer used to preprocess the examples.
-
-    Returns:
-        The prepared test examples.
-    """
-    tokenized_examples = tokenizer(text=examples["text"], truncation=True)
-    tokenized_examples["label"] = [
-        dict(
-            id=id,
-            answers=dict(
-                answer_start=answer_dct["answer_start"],
-                text=[answer_text.lower() for answer_text in answer_dct["text"]],
-            ),
-        )
-        for id, answer_dct in zip(examples["id"], examples["answers"])
-    ]
-    return tokenized_examples
diff --git a/src/scandeval/question_answering_trainer.py b/src/scandeval/question_answering_trainer.py
deleted file mode 100644
index eb6026eaf..000000000
--- a/src/scandeval/question_answering_trainer.py
+++ /dev/null
@@ -1,330 +0,0 @@
-"""Question answering Trainer subclass."""
-
-from collections import defaultdict
-from typing import TYPE_CHECKING
-
-import numpy as np
-from transformers.trainer import Trainer
-
-from .utils import get_special_token_metadata
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-
-
-class QuestionAnsweringTrainer(Trainer):
-    """Trainer subclass for question answering tasks."""
-
-    def __init__(self, *args, **kwargs) -> None:
-        """Initialize the trainer."""
-        super().__init__(*args, **kwargs)
-
-        # Get the CLS token id for the tokenizer
-        special_token_metadata = get_special_token_metadata(self.tokenizer)
-        self.cls_token_id = special_token_metadata["cls_token_id"]
-
-        # Set the label names
-        self.label_names = ["start_positions", "end_positions"]
-
-    def evaluate(
-        self,
-        eval_dataset: "Dataset | None" = None,
-        orig_eval_dataset: "Dataset | None" = None,
-        ignore_keys: list[str] | None = None,
-        metric_key_prefix: str = "eval",
-    ) -> dict[str, float] | None:
-        """Evaluate the model on the given dataset.
-
-        Args:
-            eval_dataset:
-                The dataset to evaluate on. If None, then use the stored evaluation
-                dataset.
-            orig_eval_dataset:
-                The original evaluation dataset, before any postprocessing. If None,
-                then use the stored original evaluation dataset.
-            ignore_keys:
-                The keys to ignore when computing the metrics.
-            metric_key_prefix:
-                The prefix to use for the metric keys.
-
-        Returns:
-            The metrics computed on the evaluation dataset.
-        """
-        eval_dataloader = self.get_eval_dataloader(eval_dataset)
-
-        # Temporarily disable metric computation, we will do it in the loop here.
-        compute_metrics = self.compute_metrics  # type: ignore[has-type]
-        self.compute_metrics = None
-        eval_loop = (
-            self.prediction_loop
-            if self.args.use_legacy_prediction_loop
-            else self.evaluation_loop
-        )
-        try:
-            output = eval_loop(
-                eval_dataloader,
-                description="Evaluation",
-                prediction_loss_only=True if compute_metrics is None else None,
-                ignore_keys=ignore_keys,
-                metric_key_prefix=metric_key_prefix,
-            )
-        finally:
-            self.compute_metrics = compute_metrics
-
-        if orig_eval_dataset is not None:
-            preds_and_labels = postprocess_predictions_and_labels(
-                predictions=output.predictions,
-                dataset=orig_eval_dataset,
-                prepared_dataset=eval_dataset,
-                cls_token_index=self.cls_token_id,
-            )
-            output.metrics.update(self.compute_metrics(preds_and_labels))
-
-            # Prefix all keys with metric_key_prefix + '_'
-            for key in list(output.metrics.keys()):
-                if not key.startswith(f"{metric_key_prefix}_"):
-                    output.metrics[f"{metric_key_prefix}_{key}"] = output.metrics.pop(
-                        key
-                    )
-
-        # Only the main node log the results by default
-        if self.args.should_log:
-            self.log(output.metrics)
-
-        self.control = self.callback_handler.on_evaluate(
-            self.args,
-            self.state,
-            self.control,  # type: ignore[has-type]
-            output.metrics,
-        )
-        return output.metrics
-
-
-def postprocess_predictions_and_labels(
-    predictions: list,
-    dataset: "Dataset",
-    prepared_dataset: "Dataset",
-    cls_token_index: int,
-) -> tuple[list[dict], list[dict]]:
-    """Postprocess the predictions and labels, to allow easier metric computation.
-
-    Args:
-        predictions:
-            A pair of (start_logits, end_logits) predictions.
-        dataset:
-            The dataset containing the examples.
-        prepared_dataset:
-            The dataset containing the prepared examples.
-        cls_token_index:
-            The index of the CLS token.
-
-    Returns:
-        The postprocessed predictions and labels.
-    """
-    # Extract the logits from the predictions
-    all_start_logits = predictions[0]
-    all_end_logits = predictions[1]
-
-    # Build a map from an example to its corresponding features, being the blocks of
-    # text from the context that we're feeding into the model. An example can have
-    # multiple features/blocks if it has a long context.
-    id_to_index = {k: i for i, k in enumerate(dataset["id"])}
-    features_per_example = defaultdict(list)
-    for i, feature in enumerate(prepared_dataset):
-        id = feature["id"]
-        example_index = id_to_index[id]
-        features_per_example[example_index].append(i)
-
-    # Loop over all the examples
-    predictions = list()
-    labels = list()
-    for example_index, example in enumerate(dataset):
-        # Extract the best valid answer associated with the current example
-        best_answer = find_best_answer(
-            all_start_logits=all_start_logits,
-            all_end_logits=all_end_logits,
-            prepared_dataset=prepared_dataset,
-            feature_indices=features_per_example[example_index],
-            context=example["context"],
-            max_answer_length=30,
-            num_best_logits=20,
-            min_null_score=0.0,
-            cls_token_index=cls_token_index,
-        )
-
-        # Create the final prediction dictionary, to be added to the list of
-        # predictions
-        prediction = dict(
-            id=example["id"], prediction_text=best_answer, no_answer_probability=0.0
-        )
-
-        # Add the answer to the list of predictions
-        predictions.append(prediction)
-
-        # Create the associated reference dictionary, to be added to the list of
-        # references
-        label = dict(
-            id=example["id"],
-            answers=dict(
-                text=example["answers"]["text"],
-                answer_start=example["answers"]["answer_start"],
-            ),
-        )
-
-        # Add the answer and label to the list of predictions and labels, respectively
-        labels.append(label)
-
-    return predictions, labels
-
-
-def find_best_answer(
-    all_start_logits: np.ndarray,
-    all_end_logits: np.ndarray,
-    prepared_dataset: "Dataset",
-    feature_indices: list[int],
-    context: str,
-    max_answer_length: int,
-    num_best_logits: int,
-    min_null_score: float,
-    cls_token_index: int,
-) -> str:
-    """Find the best answer for a given example.
-
-    Args:
-        all_start_logits:
-            The start logits for all the features.
-        all_end_logits:
-            The end logits for all the features.
-        prepared_dataset:
-            The dataset containing the prepared examples.
-        feature_indices:
-            The indices of the features associated with the current example.
-        context:
-            The context of the example.
-        max_answer_length:
-            The maximum length of the answer.
-        num_best_logits:
-            The number of best logits to consider.
-        min_null_score:
-            The minimum score an answer can have.
-        cls_token_index:
-            The index of the CLS token.
-
-    Returns:
-        The best answer for the example.
-    """
-    # Loop through all the features associated to the current example
-    valid_answers = list()
-    for feature_index in feature_indices:
-        # Get the features associated with the current example
-        features = prepared_dataset[feature_index]
-
-        # Get the predictions of the model for this feature
-        start_logits = all_start_logits[feature_index]
-        end_logits = all_end_logits[feature_index]
-
-        # Update minimum null prediction
-        cls_index = features["input_ids"].index(cls_token_index)
-        feature_null_score = (start_logits[cls_index] + end_logits[cls_index]).item()
-        if min_null_score < feature_null_score:
-            min_null_score = feature_null_score
-
-        # Find the valid answers for the feature
-        valid_answers_for_feature = find_valid_answers(
-            start_logits=start_logits,
-            end_logits=end_logits,
-            offset_mapping=features["offset_mapping"],
-            context=context,
-            max_answer_length=max_answer_length,
-            num_best_logits=num_best_logits,
-            min_null_score=min_null_score,
-        )
-        valid_answers.extend(valid_answers_for_feature)
-
-    # In the very rare edge case we have not a single non-null prediction, we create a
-    # fake prediction to avoid failure
-    if not valid_answers:
-        return ""
-
-    # Otherwise, we select the answer with the largest score as the best answer, and
-    # return it
-    best_answer_dict = sorted(valid_answers, key=lambda x: x["score"], reverse=True)[0]
-    return best_answer_dict["text"]
-
-
-def find_valid_answers(
-    start_logits: np.ndarray,
-    end_logits: np.ndarray,
-    offset_mapping: list[tuple[int, int]],
-    context: str,
-    max_answer_length: int,
-    num_best_logits: int,
-    min_null_score: float,
-) -> list[dict]:
-    """Find the valid answers from the start and end indexes.
-
-    Args:
-        start_logits:
-            The logits for the start of the answer.
-        end_logits:
-            The logits for the end of the answer.
-        offset_mapping:
-            The offset mapping, being a list of pairs of integers for each token index,
-            containing the start and end character index in the original context.
-        context:
-            The context of the example.
-        max_answer_length:
-            The maximum length of the answer.
-        num_best_logits:
-            The number of best logits to consider. Note that this function will run in
-            O(`num_best_logits` ^ 2) time.
-        min_null_score:
-            The minimum score an answer can have.
-
-    Returns:
-        A list of the valid answers, each being a dictionary with keys "text" and
-        "score", the score being the sum of the start and end logits.
-    """
-    # Fetch the top-k predictions for the start- and end token indices
-    start_indexes = np.argsort(start_logits)[-1 : -num_best_logits - 1 : -1].tolist()
-    end_indexes = np.argsort(end_logits)[-1 : -num_best_logits - 1 : -1].tolist()
-
-    # We loop over all combinations of starting and ending indexes for valid answers
-    valid_answers = list()
-    for start_index in start_indexes:
-        for end_index in end_indexes:
-            # If the starting or ending index is out-of-scope, meaning that they are
-            # either out of bounds or correspond to part of the input_ids that are not
-            # in the context, then we skip this index
-            if (
-                start_index >= len(offset_mapping)
-                or end_index >= len(offset_mapping)
-                or tuple(offset_mapping[start_index]) == (-1, -1)
-                or tuple(offset_mapping[end_index]) == (-1, -1)
-            ):
-                continue
-
-            # Do not consider answers with a length that is either negative or greater
-            # than the context length
-            max_val = max_answer_length + start_index - 1
-            if end_index < start_index or end_index > max_val:
-                continue
-
-            # If we got to this point then the answer is valid, so we store the
-            # corresponding start- and end character indices in the original context,
-            # and from these extract the answer
-            start_char = offset_mapping[start_index][0]
-            end_char = offset_mapping[end_index][1]
-            text = context[start_char:end_char]
-
-            # Compute the score of the answer, being the sum of the start and end
-            # logits. Intuitively, this indicates how likely the answer is to be
-            # correct, and allows us to pick the best valid answer.
-            score = start_logits[start_index] + end_logits[end_index]
-
-            # Add the answer to the list of valid answers, if the score is greater
-            # than the minimum null score
-            if score > min_null_score:
-                valid_answers.append(dict(score=score, text=text))
-
-    return valid_answers
diff --git a/src/scandeval/scores.py b/src/scandeval/scores.py
index db1d33373..d1d97ebd2 100644
--- a/src/scandeval/scores.py
+++ b/src/scandeval/scores.py
@@ -1,22 +1,22 @@
 """Aggregation of raw scores into the mean and a confidence interval."""
 
 import logging
+import typing as t
 import warnings
-from typing import TYPE_CHECKING
 
 import numpy as np
 
-if TYPE_CHECKING:
-    from .config import MetricConfig
+if t.TYPE_CHECKING:
+    from .data_models import MetricConfig
     from .types import ScoreDict
 
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 def log_scores(
     dataset_name: str,
     metric_configs: list["MetricConfig"],
-    scores: dict[str, list[dict[str, float]]],
+    scores: list[dict[str, float]],
     model_id: str,
 ) -> "ScoreDict":
     """Log the scores.
@@ -27,8 +27,8 @@ def log_scores(
         metric_configs:
             List of metrics to log.
         scores:
-            The scores that are to be logged. This is a dict with keys 'train' and
-            'test', with values being lists of dictionaries full of scores.
+            The scores that are to be logged. This is a list of dictionaries full of
+            scores.
         model_id:
             The full Hugging Face Hub path to the pretrained transformer model.
 
@@ -40,47 +40,20 @@ def log_scores(
     logger.info(f"Finished evaluation of {model_id} on {dataset_name}.")
 
     total_dict: dict[str, float] = dict()
-
-    # Logging of the aggregated scores
     for metric_cfg in metric_configs:
-        agg_scores = aggregate_scores(scores=scores, metric_config=metric_cfg)
-        test_score, test_se = agg_scores["test"]
+        test_score, test_se = aggregate_scores(scores=scores, metric_config=metric_cfg)
         test_score, test_score_str = metric_cfg.postprocessing_fn(test_score)
         test_se, test_se_str = metric_cfg.postprocessing_fn(test_se)
-        msg = f"{metric_cfg.pretty_name}:"
-
-        if "train" in agg_scores.keys():
-            train_score, train_se = agg_scores["train"]
-            train_score, train_score_str = metric_cfg.postprocessing_fn(train_score)
-            train_se, train_se_str = metric_cfg.postprocessing_fn(train_se)
-            msg += f"\n  - Test: {test_score_str} ± {test_se_str}"
-            msg += f"\n  - Train: {train_score_str} ± {train_se_str}"
-
-            # Store the aggregated train scores
-            total_dict[f"train_{metric_cfg.name}"] = train_score
-            total_dict[f"train_{metric_cfg.name}_se"] = train_se
-
-        else:
-            msg += f" {test_score_str} ± {test_se_str}"
-
-        # Store the aggregated test scores
         total_dict[f"test_{metric_cfg.name}"] = test_score
         total_dict[f"test_{metric_cfg.name}_se"] = test_se
+        logger.info(f"{metric_cfg.pretty_name}: {test_score_str} ± {test_se_str}")
 
-        # Log the scores
-        logger.info(msg)
-
-    # Define a dict with both the raw scores and the aggregated scores
-    all_scores: dict[str, dict[str, float] | dict[str, list[dict[str, float]]]]
-    all_scores = dict(raw=scores, total=total_dict)
-
-    # Return the extended scores
-    return all_scores
+    return dict(raw=scores, total=total_dict)
 
 
 def aggregate_scores(
-    scores: dict[str, list[dict[str, float]]], metric_config: "MetricConfig"
-) -> dict[str, tuple[float, float]]:
+    scores: list[dict[str, float]], metric_config: "MetricConfig"
+) -> tuple[float, float]:
     """Helper function to compute the mean with confidence intervals.
 
     Args:
@@ -92,51 +65,26 @@ def aggregate_scores(
             metric from `scores`.
 
     Returns:
-        Dictionary with keys among 'train' and 'test', with corresponding values being
-        a pair of floats, containing the score and the radius of its 95% confidence
+        A pair of floats, containing the score and the radius of its 95% confidence
         interval.
     """
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
 
-        results: dict[str, tuple[float, float]] = dict()
-
-        if "train" in scores.keys():
-            train_scores = [
-                (
-                    dct[metric_config.name]
-                    if metric_config.name in dct
-                    else dct[f"train_{metric_config.name}"]
-                )
-                for dct in scores["train"]
-            ]
-            train_score = np.mean(train_scores).item()
-
-            if len(train_scores) > 1:
-                sample_std = np.std(train_scores, ddof=1)
-                train_se = sample_std / np.sqrt(len(train_scores))
-            else:
-                train_se = np.nan
-
-            results["train"] = (train_score, 1.96 * train_se)
-
-        if "test" in scores.keys():
-            test_scores = [
-                (
-                    dct[metric_config.name]
-                    if metric_config.name in dct
-                    else dct[f"test_{metric_config.name}"]
-                )
-                for dct in scores["test"]
-            ]
-            test_score = np.mean(test_scores).item()
-
-            if len(test_scores) > 1:
-                sample_std = np.std(test_scores, ddof=1)
-                test_se = sample_std / np.sqrt(len(test_scores))
-            else:
-                test_se = np.nan
-
-            results["test"] = (test_score, 1.96 * test_se)
-
-        return results
+        test_scores = [
+            (
+                dct[metric_config.name]
+                if metric_config.name in dct
+                else dct[f"test_{metric_config.name}"]
+            )
+            for dct in scores
+        ]
+        test_score = np.mean(test_scores).item()
+
+        if len(test_scores) > 1:
+            sample_std = np.std(test_scores, ddof=1)
+            test_se = sample_std / np.sqrt(len(test_scores))
+        else:
+            test_se = np.nan
+
+        return (test_score, 1.96 * test_se)
diff --git a/src/scandeval/sequence_classification.py b/src/scandeval/sequence_classification.py
deleted file mode 100644
index e6afb16a5..000000000
--- a/src/scandeval/sequence_classification.py
+++ /dev/null
@@ -1,460 +0,0 @@
-"""Sequence classification benchmark dataset."""
-
-import importlib.util
-import itertools as it
-import logging
-import random
-import re
-from functools import partial
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-import torch
-from evaluate import EvaluationModule
-from transformers.data.data_collator import DataCollatorWithPadding
-
-from .benchmark_dataset import BenchmarkDataset
-from .exceptions import InvalidBenchmark, NeedsExtraInstalled
-from .generation import extract_raw_predictions
-from .utils import (
-    GENERATIVE_MODEL_TASKS,
-    convert_prompt_to_instruction,
-    get_special_token_metadata,
-    raise_if_model_output_contains_nan_values,
-    should_prefix_space_be_added_to_labels,
-)
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-    from transformers import BatchEncoding, PreTrainedModel
-    from transformers.modeling_utils import ModelOutput
-
-    from .config import DatasetConfig
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Labels, Predictions
-
-if importlib.util.find_spec("Levenshtein") is not None:
-    import Levenshtein
-
-
-logger = logging.getLogger(__package__)
-
-
-class SequenceClassification(BenchmarkDataset):
-    """Sequence classification benchmark dataset.
-
-    Args:
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
-
-    Attributes:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-    """
-
-    def _preprocess_data(self, dataset: "Dataset", **kwargs) -> "Dataset":
-        """Preprocess a dataset by tokenizing and aligning the labels.
-
-        Args:
-            dataset:
-                The dataset to preprocess.
-            kwargs:
-                Extra keyword arguments containing objects used in preprocessing the
-                dataset.
-
-        Returns:
-            The preprocessed dataset.
-        """
-        tokenizer: "Tokenizer" = kwargs["tokenizer"]
-
-        # Extract special token metadata from the tokenizer
-        special_token_metadata = get_special_token_metadata(tokenizer=tokenizer)
-        has_cls_token = special_token_metadata["has_cls_token"]
-        has_sep_token = special_token_metadata["has_sep_token"]
-        cls_token = special_token_metadata["cls_token"]
-        sep_token = special_token_metadata["sep_token"]
-
-        def tokenise(examples: dict) -> "BatchEncoding":
-            # If the tokenizer is not adding special tokens, then we add them manually.
-            # We don't need this when performing few-shot evaluations, so in that case
-            # we don't add the special tokens.
-            if (
-                not has_cls_token
-                and not has_sep_token
-                and cls_token is not None
-                and sep_token is not None
-                and kwargs["model_config"].task not in GENERATIVE_MODEL_TASKS
-            ):
-                examples["text"] = [
-                    f"{cls_token}{doc}{sep_token}" for doc in examples["text"]
-                ]
-
-            return tokenizer(text=examples["text"], truncation=True, padding=False)
-
-        tokenised = dataset.map(tokenise, batched=True, load_from_cache_file=False)
-
-        if kwargs["model_config"].task not in GENERATIVE_MODEL_TASKS:
-            numericalise = partial(
-                self._create_numerical_labels,
-                label2id=kwargs["hf_model_config"].label2id,
-            )
-            return tokenised.map(
-                numericalise,
-                batched=True,
-                load_from_cache_file=False,
-                keep_in_memory=True,
-            ).remove_columns(["text"])
-        else:
-            return tokenised
-
-    def _create_numerical_labels(self, examples: dict, label2id: dict) -> dict:
-        try:
-            examples["label"] = [label2id[lbl.lower()] for lbl in examples["label"]]
-        except KeyError:
-            raise InvalidBenchmark(
-                f"One of the labels in the dataset, {examples['label'].lower()}, does "
-                f"not occur in the label2id dictionary {label2id}."
-            )
-        return examples
-
-    def _load_data_collator(
-        self,
-        tokenizer: "Tokenizer | None" = None,
-        model: "PreTrainedModel | GenerativeModel | None" = None,
-    ):
-        """Load the data collator used to prepare samples during finetuning.
-
-        Args:
-            tokenizer:
-                A pretrained tokenizer. Can be None if the tokenizer is not used in the
-                initialisation of the data collator. Defaults to None.
-            model:
-                A pretrained model. Can be None if the model is not used in the
-                initialisation of the data collator. Defaults to None.
-
-        Returns:
-            The data collator.
-        """
-        return DataCollatorWithPadding(tokenizer, padding="longest")
-
-    def _compute_metrics(
-        self,
-        model_outputs_and_labels: tuple["Predictions", "Labels"],
-        id2label: dict[int, str],
-    ) -> dict[str, float]:
-        """Compute the metrics needed for evaluation.
-
-        Args:
-            model_outputs_and_labels:
-                The first sequence contains the model outputs and the second sequence
-                contains the true labels.
-            id2label:
-                Conversion of indices to labels.
-
-        Returns:
-            A dictionary with the names of the metrics as keys and the metric values as
-            values.
-        """
-        model_outputs, labels = model_outputs_and_labels
-        label2id = {label: idx for idx, label in id2label.items()}
-
-        raise_if_model_output_contains_nan_values(model_output=model_outputs)
-
-        model_output_dtype = np.asarray(model_outputs).dtype
-        if model_output_dtype in [np.float16, np.float32, np.float64]:
-            predictions = np.asarray(model_outputs).argmax(axis=-1)
-        else:
-            predictions = model_outputs
-
-        prompt_label_to_label_mapping = {
-            prompt_label: label
-            for label, prompt_label in self.dataset_config.prompt_label_mapping.items()
-        }
-        predictions = [
-            (
-                label2id[prompt_label_to_label_mapping[pred.lower()]]
-                if isinstance(pred, str)
-                else pred
-            )
-            for pred in predictions
-        ]
-
-        label_ids = [
-            label2id[label.lower()] if isinstance(label, str) else label
-            for label in labels
-        ]
-
-        results: dict[str, float] = dict()
-        for cfg in self.dataset_config.task.metrics:
-            metric = self._metrics[cfg.name]
-            assert isinstance(metric, EvaluationModule)
-            score_dict: dict[str, float] | None = metric.compute(
-                predictions=predictions, references=label_ids, **cfg.compute_kwargs
-            )
-
-            # The metric returns None if we are running on multi-GPU and the current
-            # process is not the main process
-            if score_dict is not None:
-                scores = score_dict[cfg.results_key]
-                if isinstance(scores, list):
-                    scores = sum(scores) / len(scores)
-                results[cfg.name] = scores
-
-        return results
-
-    def _extract_few_shot_examples(
-        self, train_dataset: "Dataset", random_seed: int
-    ) -> list[dict[str, Any]]:
-        """Extract few-shot examples from the training dataset.
-
-        Args:
-            train_dataset:
-                The training dataset.
-            random_seed:
-                The random seed to use when extracting the few-shot examples.
-
-        Returns:
-            The few-shot examples.
-        """
-        shuffled_train = train_dataset.shuffle(seed=random_seed)
-        num_few_shots = self.dataset_config.num_few_shot_examples
-        labels = it.cycle(self.dataset_config.task.labels)
-        few_shot_examples: list[dict[str, Any]] = list()
-
-        # We pick the few-shot examples one at a time rather than all at once since
-        # we're working with a bootstrapped training dataset, meaning that it will have
-        # duplicates. This ensures that we don't have any duplicates in the few-shot
-        # examples
-        while len(few_shot_examples) < num_few_shots:
-            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"]
-            )
-
-        random.seed(random_seed)
-        random.shuffle(few_shot_examples)
-        return few_shot_examples
-
-    def _apply_few_shot_prompt(
-        self, examples: dict, few_shot_examples: list[dict], tokenizer: "Tokenizer"
-    ) -> dict:
-        """Apply a few-shot prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            few_shot_examples:
-                The examples to be included in the few-shot prompt.
-            tokenizer:
-                The tokenizer to use to encode the few-shot prompt.
-
-        Returns:
-            The examples with the few-shot prompt applied.
-        """
-        # Build the few-shot part of the prompt
-        label_mapping = self.dataset_config.prompt_label_mapping
-        few_shot_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=re.sub(r"\n+", "\n", example["text"]).strip(),
-                label=label_mapping[example["label"].lower()],
-            )
-            for example in few_shot_examples
-        ]
-        prompt_prefix = ""
-        if self.dataset_config.prompt_prefix:
-            prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"
-        few_shot_prompt = prompt_prefix + "\n\n".join(few_shot_prompts)
-
-        # Add the texts from the examples to the prompts. We remove newlines from the
-        # examples as they have the special function to separate the few-shot examples
-        # from one another
-        new_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=re.sub(r"\n+", "\n", text).strip(), label=""
-            )
-            for text in examples["text"]
-        ]
-
-        final_prompts = [
-            few_shot_prompt + "\n\n" + new_prompt for new_prompt in new_prompts
-        ]
-
-        examples["text"] = final_prompts
-
-        return examples
-
-    def _apply_instruction_prompt(self, examples: dict, tokenizer: "Tokenizer") -> dict:
-        """Apply an instruction prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            tokenizer:
-                The tokenizer to use to encode the instruction prompt.
-
-        Returns:
-            The examples with the instruction prompt applied.
-        """
-        prompts = [
-            self.dataset_config.instruction_prompt.format(
-                text=re.sub(r"\n+", "\n", text).strip()
-            )
-            for text in examples["text"]
-        ]
-        prompts = [
-            convert_prompt_to_instruction(prompt=prompt, tokenizer=tokenizer)
-            for prompt in prompts
-        ]
-        examples["text"] = prompts
-        return examples
-
-    def _extract_labels_from_generation(
-        self,
-        input_batch: dict[str, list],
-        model_output: "ModelOutput",
-        tokenizer: "Tokenizer",
-    ) -> list[Any]:
-        """Extract the predicted labels from the generated output.
-
-        Args:
-            input_batch:
-                The input batch, where the keys are the feature names and the values
-                are lists with the feature values.
-            model_output:
-                The raw generated output of the model.
-            tokenizer:
-                The tokenizer used together with the model.
-
-        Returns:
-            The predicted labels.
-        """
-        if "scores" in model_output:
-            if isinstance(model_output["scores"], tuple):
-                all_logprobs = torch.stack(model_output["scores"], dim=1)
-            else:
-                all_logprobs = model_output["scores"]
-            return get_closest_logprobs_labels(
-                generation_logprobs=all_logprobs,
-                tokenizer=tokenizer,
-                dataset_config=self.dataset_config,
-            )
-        else:
-            return get_closest_word_edit_labels(
-                generated_sequences=model_output["sequences"],
-                tokenizer=tokenizer,
-                dataset_config=self.dataset_config,
-            )
-
-
-def get_closest_logprobs_labels(
-    generation_logprobs: torch.Tensor,
-    tokenizer: "Tokenizer",
-    dataset_config: "DatasetConfig",
-) -> list[str]:
-    """Get the labels with the highest predicted logprob value.
-
-    In case a candidate label is split into multiple tokens, we only use the first
-    token to compute the logprob value. E.g., if the candidate label "positive" is
-    tokenised as ["pos", "itive"], we only use the logprob value of "pos" to
-    represent the logprob value of the entire label.
-
-    Args:
-        generation_logprobs:
-            The logprobs of the generated tokens.
-        tokenizer:
-            The tokenizer used to generate the tokens.
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-
-    Returns:
-        The predicted labels.
-    """
-    candidate_labels = [
-        dataset_config.prompt_label_mapping[lbl]
-        for lbl in dataset_config.id2label.values()
-    ]
-
-    add_prefix_space_to_labels = should_prefix_space_be_added_to_labels(
-        labels_to_be_generated=candidate_labels, tokenizer=tokenizer
-    )
-
-    # Shape: [batch_size, num_candidate_labels]
-    pred_logprobs = torch.empty(
-        generation_logprobs.shape[0],
-        len(candidate_labels),
-        device=generation_logprobs.device,
-    )
-
-    for idx, candidate_label in enumerate(candidate_labels):
-        # We only use the first token to represent the logprob value of the entire
-        # label.
-        label_ready_for_tokenization = candidate_label.lower()
-        if add_prefix_space_to_labels:
-            label_ready_for_tokenization = " " + label_ready_for_tokenization
-        candidate_label_ids: list[list[int]] = tokenizer(
-            [label_ready_for_tokenization], add_special_tokens=False
-        )["input_ids"]
-        candidate_label_id: int = candidate_label_ids[0][0]
-        pred_logprobs[:, idx] = generation_logprobs[:, 0, candidate_label_id]
-
-    # Shape: [batch_size,]
-    predicted_label_ids = pred_logprobs.argmax(dim=1)
-
-    predicted_labels = [candidate_labels[idx] for idx in predicted_label_ids]
-
-    return predicted_labels
-
-
-def get_closest_word_edit_labels(
-    generated_sequences: torch.Tensor,
-    tokenizer: "Tokenizer",
-    dataset_config: "DatasetConfig",
-) -> list[str]:
-    """Get the labels with the smallest edit distance to the predicted labels.
-
-    Args:
-        generated_sequences:
-            The generated sequences from the model. The outer-most list is the
-            batch dimension, the inner-most list is the sequence dimension,
-            consisting of token IDs.
-        tokenizer:
-            The tokenizer used to generate the tokens.
-        dataset_config:
-            The configuration of the dataset.
-
-    Returns:
-        The candidate labels with the smallest edit distance to the predicted labels.
-    """
-    if importlib.util.find_spec("Levenshtein") is None:
-        raise NeedsExtraInstalled(extra="openai")
-
-    raw_predictions = extract_raw_predictions(
-        generated_sequences=generated_sequences, tokenizer=tokenizer
-    )
-
-    candidate_labels = [
-        dataset_config.prompt_label_mapping[lbl]
-        for lbl in dataset_config.id2label.values()
-    ]
-    new_predicted_labels: list[str] = list()
-    for predicted_label in raw_predictions:
-        edit_distances = [
-            Levenshtein.distance(s1=predicted_label.lower(), s2=candidate_label.lower())
-            for candidate_label in candidate_labels
-        ]
-        closest_label = candidate_labels[np.argmin(edit_distances).item()]
-        new_predicted_labels.append(closest_label)
-    return new_predicted_labels
diff --git a/src/scandeval/speed_benchmark.py b/src/scandeval/speed_benchmark.py
index c4322fddf..a0dc8f5f4 100644
--- a/src/scandeval/speed_benchmark.py
+++ b/src/scandeval/speed_benchmark.py
@@ -1,152 +1,100 @@
 """Benchmarking model inference speed."""
 
+import collections.abc as c
 import logging
-from collections import defaultdict
-from typing import TYPE_CHECKING
+import typing as t
 
 import pyinfer
-import torch
 from tqdm.auto import tqdm
-from transformers.modeling_utils import GenerationConfig, PreTrainedModel
-
+from transformers import AutoTokenizer
+
+from .benchmark_modules import (
+    BenchmarkModule,
+    HuggingFaceEncoderModel,
+    LiteLLMModel,
+    VLLMModel,
+)
+from .data_models import BenchmarkConfig
 from .exceptions import InvalidBenchmark
-from .model_loading import load_model
-from .utils import clear_memory, model_is_generative
-
-if TYPE_CHECKING:
-    from transformers.tokenization_utils import PreTrainedTokenizer
-
-    from .config import BenchmarkConfig, DatasetConfig, ModelConfig
-    from .protocols import GenerativeModel, Tokenizer
-
+from .utils import clear_memory
 
-logger = logging.getLogger(__package__)
+logger = logging.getLogger("scandeval")
 
 
 def benchmark_speed(
-    itr: tqdm,
-    tokenizer: "PreTrainedTokenizer",
-    model: PreTrainedModel,
-    model_config: "ModelConfig",
-    dataset_config: "DatasetConfig",
-    benchmark_config: "BenchmarkConfig",
-) -> dict[str, list[dict[str, float]]]:
+    model: "BenchmarkModule", benchmark_config: "BenchmarkConfig"
+) -> list[dict[str, float]]:
     """Benchmark model inference speed.
 
     Args:
-        itr:
-            tqdm iterator.
-        tokenizer:
-            Tokenizer to use.
         model:
             Model to use.
-        model_config:
-            Model configuration.
-        dataset_config:
-            Dataset configuration.
         benchmark_config:
-            Benchmark configuration.
+            Configuration for the benchmark.
 
     Returns:
         Dictionary of scores.
     """
-    scores: dict[str, list[dict[str, float]]] = defaultdict(list)
-
-    for itr_idx in itr:
-        itr_scores = benchmark_speed_single_iteration(
-            tokenizer=tokenizer,
-            model=model,
-            itr_idx=itr_idx,
-            model_config=model_config,
-            dataset_config=dataset_config,
-            benchmark_config=benchmark_config,
-        )
+    scores: list[dict[str, float]] = list()
+    for idx in tqdm(
+        iterable=range(benchmark_config.num_iterations),
+        desc="Benchmarking",
+        disable=not benchmark_config.progress_bar,
+    ):
+        itr_scores = benchmark_speed_single_iteration(model=model, itr_idx=idx)
         clear_memory()
-
-        if isinstance(itr_scores, Exception):
-            raise InvalidBenchmark(f"Speed benchmark failed with error: {itr_scores!r}")
-        else:
-            scores["test"].append(itr_scores["test"])
-            if benchmark_config.evaluate_train:
-                scores["train"].append(itr_scores["train"])
-            logger.debug(f"Scores for iteration {itr_idx}: {itr_scores}")
-
+        scores.append(itr_scores)
+        logger.debug(f"Scores for iteration {idx}: {itr_scores}")
     return scores
 
 
 def benchmark_speed_single_iteration(
-    tokenizer: "Tokenizer",
-    model: "PreTrainedModel | GenerativeModel",
-    itr_idx: int,
-    model_config: "ModelConfig",
-    dataset_config: "DatasetConfig",
-    benchmark_config: "BenchmarkConfig",
-) -> dict[str, dict[str, float]] | Exception:
+    model: "BenchmarkModule", itr_idx: int
+) -> dict[str, float]:
     """Run a single iteration of the speed benchmark.
 
     Args:
-        tokenizer:
-            The tokenizer to use in the benchmark.
         model:
             The model to use in the benchmark.
         itr_idx:
             The index of the iteration.
-        model_config:
-            The model configuration.
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
 
     Returns:
-        A dictionary containing the scores for the current iteration, with keys `train`
-        and `test`. If an exception is raised, then the exception is returned.
+        A dictionary containing the scores for the current iteration.
     """
-    is_generative = model_is_generative(model=model)
+    gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+    base_doc = "Document which contains roughly 10 tokens. "
+    multiplier = 10 * (1 + itr_idx)
+    doc = base_doc * multiplier
+    short_multiplier = 1.25 * (1 + itr_idx)
+    short_doc = base_doc * round(short_multiplier)
+
+    def generate_messages_predict(doc: str) -> None:
+        model.generate(inputs=dict(messages=[[dict(role="user", content=doc)]]))
+
+    def generate_prompt_predict(doc: str) -> None:
+        model.generate(inputs=dict(prompt=[doc]))
+
+    def encoder_predict(doc: str) -> None:
+        tokenizer = model.get_tokenizer()
+        pytorch_model = model.get_pytorch_module()
+        inputs = {
+            key: tensor.to(pytorch_model.device)
+            for key, tensor in tokenizer(
+                text=[doc], truncation=True, return_tensors="pt"
+            ).items()
+        }
+        pytorch_model(**inputs)
+
+    predict_fn_mapping: dict[t.Type[BenchmarkModule], c.Callable[[str], None]] = {
+        HuggingFaceEncoderModel: encoder_predict,
+        LiteLLMModel: generate_messages_predict,
+        VLLMModel: generate_prompt_predict,
+    }
+    predict = predict_fn_mapping[type(model)]
 
-    scores: dict[str, dict[str, float]] = dict()
     try:
-        # Reinitialise a new model
-        if tokenizer is None or model is None:
-            model, tokenizer = load_model(
-                model_config=model_config,
-                dataset_config=dataset_config,
-                benchmark_config=benchmark_config,
-            )
-
-        def predict(doc: str) -> None:
-            """Function used to benchmark inference speed of the model."""
-            # Raise an error if the tokenizer or model is undefined
-            if tokenizer is None or model is None:
-                raise ValueError("Tokenizer and model must not be None.")
-
-            # Tokenize the document
-            inputs = tokenizer(doc, padding=True, truncation=True, return_tensors="pt")
-
-            inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
-
-            # Run inference with the model
-            with torch.inference_mode():
-                if is_generative:
-                    model.generate(
-                        inputs=inputs["input_ids"],
-                        generation_config=GenerationConfig(
-                            max_new_tokens=1,
-                            pad_token_id=model.config.pad_token_id,
-                            eos_token_id=model.config.eos_token_id,
-                            do_sample=False,
-                        ),
-                    )
-                else:
-                    assert isinstance(model, PreTrainedModel)
-                    model(**inputs)
-
-        base_doc = "Document which contains roughly 10 tokens. "
-        multiplier = 10 * (1 + itr_idx)
-        doc = base_doc * multiplier
-        short_multiplier = 1.25 * (1 + itr_idx)
-        short_doc = base_doc * round(short_multiplier)
-
         # Do a warmup run, as the first run is always slower
         pyinfer.InferenceReport(model=predict, inputs=base_doc, n_seconds=1).run(
             print_report=False
@@ -155,25 +103,22 @@ def predict(doc: str) -> None:
         speed_scores = pyinfer.InferenceReport(
             model=predict, inputs=doc, n_seconds=3
         ).run(print_report=False)
-        num_tokens = len(tokenizer([doc], truncation=True)["input_ids"][0])
-        tokens_per_second = speed_scores["Infer(p/sec)"] * num_tokens
+        num_gpt2_tokens = len(gpt2_tokenizer([doc], truncation=True)["input_ids"][0])
+        gpt2_tokens_per_second = speed_scores["Infer(p/sec)"] * num_gpt2_tokens
 
         speed_scores_short = pyinfer.InferenceReport(
             model=predict, inputs=short_doc, n_seconds=3
         ).run(print_report=False)
-        num_tokens_short = len(tokenizer([short_doc], truncation=True)["input_ids"][0])
-        tokens_per_second_short = speed_scores_short["Infer(p/sec)"] * num_tokens_short
-
-        scores["test"] = dict(
-            test_speed=tokens_per_second, test_speed_short=tokens_per_second_short
+        num_gpt2_tokens_short = len(
+            gpt2_tokenizer([short_doc], truncation=True)["input_ids"][0]
+        )
+        gpt2_tokens_per_second_short = (
+            speed_scores_short["Infer(p/sec)"] * num_gpt2_tokens_short
         )
-        if benchmark_config.evaluate_train:
-            scores["train"] = dict(
-                train_speed=tokens_per_second, train_speed_short=tokens_per_second_short
-            )
-
-        # Return the scores
-        return scores
 
     except (RuntimeError, ValueError, IndexError) as e:
-        return e
+        raise InvalidBenchmark(f"Speed benchmark failed with error: {e!r}")
+
+    return dict(
+        test_speed=gpt2_tokens_per_second, test_speed_short=gpt2_tokens_per_second_short
+    )
diff --git a/src/scandeval/structured_generation_utils.py b/src/scandeval/structured_generation_utils.py
index 4f6d4e4f6..1c9447c1a 100644
--- a/src/scandeval/structured_generation_utils.py
+++ b/src/scandeval/structured_generation_utils.py
@@ -1,18 +1,16 @@
 """Utility functions related to structured generation."""
 
+import collections.abc as c
 import importlib.util
-from typing import TYPE_CHECKING, Any, Callable
+import typing as t
 
 import torch
 from pydantic import BaseModel, conlist, create_model
 
 if importlib.util.find_spec("outlines") is not None:
-    from outlines.integrations.transformers import JSONPrefixAllowedTokens
     from outlines.integrations.vllm import JSONLogitsProcessor
 
-if TYPE_CHECKING:
-    from outlines.integrations.transformers import JSONPrefixAllowedTokens
-    from transformers import PreTrainedTokenizerBase
+if t.TYPE_CHECKING:
     from vllm import LLM
 
 
@@ -26,37 +24,16 @@ def get_ner_schema(ner_tag_names: list[str]) -> type[BaseModel]:
     Returns:
         The schema for the NER answer format.
     """
-    keys_and_their_types: dict[str, Any] = {
+    keys_and_their_types: dict[str, t.Any] = {
         tag_name: (conlist(str, max_length=5), ...) for tag_name in ner_tag_names
     }
     schema = create_model("AnswerFormat", **keys_and_their_types)
     return schema
 
 
-def get_ner_prefix_allowed_tokens_fn(
-    ner_tag_names: list[str], tokenizer: "PreTrainedTokenizerBase"
-) -> "JSONPrefixAllowedTokens":
-    """Get the prefix allowed tokens function for the NER task, used in `transformers`.
-
-    Args:
-        ner_tag_names:
-            The NER tag names.
-        tokenizer:
-            The tokenizer to use for tokenizing the JSON Schema.
-
-    Returns:
-        The prefix allowed tokens function for the NER task.
-    """
-    return JSONPrefixAllowedTokens(
-        schema=get_ner_schema(ner_tag_names=ner_tag_names),
-        tokenizer_or_pipe=tokenizer,
-        whitespace_pattern=r" ?",
-    )
-
-
 def get_ner_logits_processors(
     ner_tag_names: list[str], llm: "LLM"
-) -> list[Callable[[list[int], torch.Tensor], torch.Tensor]]:
+) -> list[c.Callable[[list[int], torch.Tensor], torch.Tensor]]:
     """Get the logits processors for the NER task, used in vLLM.
 
     Args:
diff --git a/src/scandeval/task_utils/__init__.py b/src/scandeval/task_utils/__init__.py
new file mode 100644
index 000000000..aa8d95b46
--- /dev/null
+++ b/src/scandeval/task_utils/__init__.py
@@ -0,0 +1,1 @@
+"""Utility functions related to the different tasks and supertasks."""
diff --git a/src/scandeval/task_utils/question_answering.py b/src/scandeval/task_utils/question_answering.py
new file mode 100644
index 000000000..58ee91418
--- /dev/null
+++ b/src/scandeval/task_utils/question_answering.py
@@ -0,0 +1,668 @@
+"""Utility functions related to the question-answering supertask."""
+
+import logging
+import typing as t
+from collections import defaultdict
+
+import evaluate
+import numpy as np
+from evaluate import EvaluationModule
+from transformers import PreTrainedTokenizer
+from transformers.trainer import Trainer
+
+from ..data_models import BenchmarkConfig, DatasetConfig, GenerativeModelOutput
+from ..utils import (
+    get_special_token_metadata,
+    raise_if_model_output_contains_nan_values,
+)
+
+if t.TYPE_CHECKING:
+    from datasets.arrow_dataset import Dataset
+    from transformers.tokenization_utils_base import BatchEncoding
+
+    from ..types import Labels, Predictions
+
+logger = logging.getLogger("scandeval")
+
+
+class QuestionAnsweringTrainer(Trainer):
+    """Trainer subclass for question answering tasks."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        """Initialize the trainer."""
+        super().__init__(*args, **kwargs)
+
+        # Get the CLS token id for the tokenizer
+        special_token_metadata = get_special_token_metadata(self.tokenizer)
+        self.cls_token_id = special_token_metadata["cls_token_id"]
+
+        # Set the label names
+        self.label_names = ["start_positions", "end_positions"]
+
+    def evaluate(
+        self,
+        eval_dataset: "Dataset | None" = None,
+        orig_eval_dataset: "Dataset | None" = None,
+        ignore_keys: list[str] | None = None,
+        metric_key_prefix: str = "eval",
+    ) -> dict[str, float] | None:
+        """Evaluate the model on the given dataset.
+
+        Args:
+            eval_dataset:
+                The dataset to evaluate on. If None, then use the stored evaluation
+                dataset.
+            orig_eval_dataset:
+                The original evaluation dataset, before any postprocessing. If None,
+                then use the stored original evaluation dataset.
+            ignore_keys:
+                The keys to ignore when computing the metrics.
+            metric_key_prefix:
+                The prefix to use for the metric keys.
+
+        Returns:
+            The metrics computed on the evaluation dataset.
+        """
+        eval_dataloader = self.get_eval_dataloader(eval_dataset)
+
+        # Temporarily disable metric computation, we will do it in the loop here.
+        compute_metrics = self.compute_metrics  # type: ignore[has-type]
+        self.compute_metrics = None
+        eval_loop = (
+            self.prediction_loop
+            if self.args.use_legacy_prediction_loop
+            else self.evaluation_loop
+        )
+        try:
+            output = eval_loop(
+                eval_dataloader,
+                description="Evaluation",
+                prediction_loss_only=True if compute_metrics is None else None,
+                ignore_keys=ignore_keys,
+                metric_key_prefix=metric_key_prefix,
+            )
+        finally:
+            self.compute_metrics = compute_metrics
+
+        if orig_eval_dataset is not None:
+            preds_and_labels = postprocess_predictions_and_labels(
+                predictions=output.predictions,
+                dataset=orig_eval_dataset,
+                prepared_dataset=eval_dataset,
+                cls_token_index=self.cls_token_id,
+            )
+            output.metrics.update(self.compute_metrics(preds_and_labels))
+
+            # Prefix all keys with metric_key_prefix + '_'
+            for key in list(output.metrics.keys()):
+                if not key.startswith(f"{metric_key_prefix}_"):
+                    output.metrics[f"{metric_key_prefix}_{key}"] = output.metrics.pop(
+                        key
+                    )
+
+        # Only the main node log the results by default
+        if self.args.should_log:
+            self.log(output.metrics)
+
+        self.control = self.callback_handler.on_evaluate(
+            self.args,
+            self.state,
+            self.control,  # type: ignore[has-type]
+            output.metrics,
+        )
+        return output.metrics
+
+
+def compute_metrics(
+    model_outputs_and_labels: tuple["Predictions", "Labels"],
+    id2label: dict[int, str],
+    dataset_config: "DatasetConfig",
+    benchmark_config: "BenchmarkConfig",
+) -> dict[str, float]:
+    """Compute the metrics needed for evaluation.
+
+    Args:
+        model_outputs_and_labels:
+            The first sequence contains the model outputs and the second sequence
+            contains the true labels.
+        id2label:
+            Conversion of indices to labels.
+        dataset_config:
+            The configuration of the dataset.
+        benchmark_config:
+            The configuration of the benchmark.
+
+    Returns:
+        A dictionary with the names of the metrics as keys and the metric values as
+        values.
+    """
+    model_outputs, labels = model_outputs_and_labels
+    raise_if_model_output_contains_nan_values(model_output=model_outputs)
+
+    metrics = {
+        metric_cfg.name: (
+            evaluate.load(
+                path=metric_cfg.huggingface_id, cache_dir=benchmark_config.cache_dir
+            )
+            if metric_cfg.huggingface_id != ""
+            else None
+        )
+        for metric_cfg in dataset_config.task.metrics
+    }
+
+    model_output_dtype = np.asarray(model_outputs).dtype
+    if model_output_dtype in [np.float16, np.float32, np.float64]:
+        predictions = np.asarray(model_outputs).argmax(axis=-1)
+    else:
+        predictions = model_outputs
+
+    results: dict[str, float] = dict()
+    for cfg in dataset_config.task.metrics:
+        metric = metrics[cfg.name]
+        assert isinstance(metric, EvaluationModule)
+        score_dict: dict[str, float] | None = metric.compute(
+            predictions=predictions, references=labels, **cfg.compute_kwargs
+        )
+
+        # The metric returns None if we are running on multi-GPU and the current
+        # process is not the main process
+        if score_dict is not None:
+            scores = score_dict[cfg.results_key]
+            if isinstance(scores, list):
+                scores = sum(scores) / len(scores)
+            results[cfg.name] = scores
+
+    return results
+
+
+def extract_labels_from_generation(
+    input_batch: dict[str, list], model_output: "GenerativeModelOutput"
+) -> list[t.Any]:
+    """Extract the predicted labels from the generated output.
+
+    Args:
+        input_batch:
+            The input batch, where the keys are the feature names and the values
+            are lists with the feature values.
+        model_output:
+            The raw generated output of the model.
+
+    Returns:
+        The predicted labels.
+    """
+    raw_predictions = model_output.sequences
+    predictions = [
+        dict(id=id, prediction_text=predicted_answer.lower(), no_answer_probability=0.0)
+        for id, predicted_answer in zip(input_batch["id"], raw_predictions)
+    ]
+    return predictions
+
+
+def prepare_train_examples(
+    examples: "BatchEncoding", tokenizer: "PreTrainedTokenizer"
+) -> "BatchEncoding":
+    """Prepare the features for training.
+
+    Args:
+        examples:
+            The examples to prepare.
+        tokenizer:
+            The tokenizer to use to prepare the examples.
+
+    Returns:
+        The prepared examples.
+    """
+    # Some of the questions have lots of whitespace on the left, which is not useful
+    # and will make the truncation of the context fail (the tokenized question will
+    # take a lots of space). So we remove that left whitespace
+    examples["question"] = [q.lstrip() for q in examples["question"]]
+
+    # Extract special token metadata from the tokenizer
+    special_token_metadata = get_special_token_metadata(tokenizer=tokenizer)
+    has_cls_token = special_token_metadata["has_cls_token"]
+    has_sep_token = special_token_metadata["has_sep_token"]
+    cls_token_id = special_token_metadata["cls_token_id"]
+    cls_token = special_token_metadata["cls_token"]
+    sep_token = special_token_metadata["sep_token"]
+
+    # If the tokenizer is not adding special tokens, then we add them manually
+    if not has_cls_token and not has_sep_token:
+        examples["question"] = [
+            f"{cls_token}{q}{sep_token}" for q in examples["question"]
+        ]
+        examples["context"] = [f"{c}{sep_token}" for c in examples["context"]]
+
+    # Set the stride used during tokenization, when the context is long enough to be
+    # split into several features. Since we are always keeping the question tokens, we
+    # need to make sure that the stride does not exceed the resulting maximum context
+    # length.
+    max_question_tokens = max(len(tokenizer(q).input_ids) for q in examples["question"])
+    num_special_tokens = int(has_cls_token) + int(has_sep_token)
+    stride = tokenizer.model_max_length // 4
+    max_length = tokenizer.model_max_length - stride
+    stride = min(stride, max_length - max_question_tokens - num_special_tokens)
+    max_length = tokenizer.model_max_length - stride
+
+    # Tokenize our examples with truncation and padding, but keep the overflows using a
+    # stride. This results in one example possible giving several features when a
+    # context is long, each of those features having a context that overlaps a bit the
+    # context of the previous feature.
+    tokenized_examples = tokenizer(
+        text=examples["question"],
+        text_pair=examples["context"],
+        truncation="only_second",
+        max_length=max_length,
+        stride=stride,
+        return_overflowing_tokens=True,
+        return_offsets_mapping=True,
+        padding="max_length",
+    )
+
+    # Since one example might give us several features if it has a long context, we
+    # need a map from a feature to its corresponding example. This key gives us just
+    # that
+    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+
+    # The offset mappings will give us a map from token to character position in the
+    # original context. This will help us compute the start_positions and
+    # end_positions.
+    offset_mapping = tokenized_examples.pop("offset_mapping")
+
+    # Initialise the start- and end positions of the answers
+    tokenized_examples["start_positions"] = list()
+    tokenized_examples["end_positions"] = list()
+
+    for i, offsets in enumerate(offset_mapping):
+        # Get the input IDs for the current example
+        input_ids = tokenized_examples.input_ids[i]
+
+        # We will label impossible answers with the index of the CLS token
+        cls_index = input_ids.index(cls_token_id)
+
+        # Grab the sequence corresponding to that example (to know what is the context
+        # and what is the question).
+        sequence_ids = tokenized_examples.sequence_ids(i)
+
+        # Manually ensure that the special tokens are set to None in `sequence_ids`
+        for special_token in tokenizer.special_tokens_map.keys():
+            if hasattr(tokenizer, f"{special_token}_id"):
+                special_token_id = getattr(tokenizer, f"{special_token}_id")
+                if special_token_id is not None:
+                    sequence_ids = [
+                        None if token_id == special_token_id else seq_id
+                        for token_id, seq_id in zip(input_ids, sequence_ids)
+                    ]
+
+        # One example can give several spans, this is the index of the example
+        # containing this span of text.
+        sample_index = sample_mapping[i]
+        answers = examples["answers"][sample_index]
+
+        # If no answers are given, set the cls_index as answer.
+        if len(answers["answer_start"]) == 0:
+            tokenized_examples.start_positions.append(cls_index)
+            tokenized_examples.end_positions.append(cls_index)
+
+        else:
+            # Start/end character index of the answer in the text.
+            start_char = answers["answer_start"][0]
+            end_char = start_char + len(answers["text"][0])
+
+            # Start token index of the current span in the text.
+            token_start_index = 0
+            while sequence_ids[token_start_index] != 1:
+                token_start_index += 1
+
+            # End token index of the current span in the text.
+            token_end_index = len(input_ids) - 1
+            while sequence_ids[token_end_index] != 1:
+                token_end_index -= 1
+
+            # Detect if the answer is out of the span (in which case this feature is
+            # labeled with the CLS index).
+            if not (
+                offsets[token_start_index][0] <= start_char
+                and offsets[token_end_index][1] >= end_char
+            ):
+                tokenized_examples.start_positions.append(cls_index)
+                tokenized_examples.end_positions.append(cls_index)
+
+            # Otherwise move the token_start_index and token_end_index to the two ends
+            # of the answer. Note: we could go after the last offset if the answer is
+            # the last word (edge case).
+            else:
+                while (
+                    token_start_index <= token_end_index
+                    and offsets[token_start_index][0] <= start_char
+                ):
+                    token_start_index += 1
+                token_start_index -= 1
+                tokenized_examples.start_positions.append(token_start_index)
+                while (
+                    token_start_index <= token_end_index
+                    and offsets[token_end_index][1] >= end_char
+                ):
+                    token_end_index -= 1
+                token_end_index += 1
+                tokenized_examples.end_positions.append(token_end_index)
+                assert token_end_index >= token_start_index
+
+    return tokenized_examples
+
+
+def prepare_test_examples(
+    examples: "BatchEncoding", tokenizer: "PreTrainedTokenizer"
+) -> "BatchEncoding":
+    """Prepare test examples.
+
+    Args:
+        examples:
+            Dictionary of test examples.
+        tokenizer:
+            The tokenizer used to preprocess the examples.
+
+    Returns:
+        The prepared test examples.
+    """
+    # Some of the questions have lots of whitespace on the left, which is not useful
+    # and will make the truncation of the context fail (the tokenized question will
+    # take a lots of space). So we remove that left whitespace
+    examples["question"] = [q.lstrip() for q in examples["question"]]
+
+    # Extract special token metadata from the tokenizer
+    special_token_metadata = get_special_token_metadata(tokenizer=tokenizer)
+    has_cls_token = special_token_metadata["has_cls_token"]
+    has_sep_token = special_token_metadata["has_sep_token"]
+    cls_token = special_token_metadata["cls_token"]
+    sep_token = special_token_metadata["sep_token"]
+
+    # If the tokenizer is not adding special tokens, then we add them manually
+    if not has_cls_token and not has_sep_token:
+        examples["question"] = [
+            f"{cls_token}{q}{sep_token}" for q in examples["question"]
+        ]
+        examples["context"] = [f"{c}{sep_token}" for c in examples["context"]]
+
+    # Set the stride used during tokenization, when the context is long enough to be
+    # split into several features. Since we are always keeping the question tokens, we
+    # need to make sure that the stride does not exceed the resulting maximum context
+    # length.
+    max_question_tokens = max(len(tokenizer(q).input_ids) for q in examples["question"])
+    num_special_tokens = int(has_cls_token) + int(has_sep_token)
+    stride = tokenizer.model_max_length // 4
+    max_length = tokenizer.model_max_length - stride
+    stride = min(stride, max_length - max_question_tokens - num_special_tokens)
+    max_length = tokenizer.model_max_length - stride
+
+    # Tokenize our examples with truncation and maybe padding, but keep the overflows
+    # using a stride. This results in one example possible giving several features when
+    # a context is long, each of those features having a context that overlaps a bit
+    # the context of the previous feature.
+    tokenized_examples = tokenizer(
+        text=examples["question"],
+        text_pair=examples["context"],
+        truncation="only_second",
+        max_length=max_length,
+        stride=stride,
+        return_overflowing_tokens=True,
+        return_offsets_mapping=True,
+        padding="max_length",
+    )
+
+    # Since one example might give us several features if it has a long context, we
+    # need a map from a feature to its corresponding example. This key gives us just
+    # that.
+    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+
+    # We keep the id that gave us this feature and we will store the offset mappings.
+    tokenized_examples["id"] = list()
+
+    for i in range(len(tokenized_examples.input_ids)):
+        # Grab the sequence corresponding to that example (to know what is the context
+        # and what is the question).
+        sequence_ids = tokenized_examples.sequence_ids(i)
+        context_index = 1
+
+        # One example can give several spans, this is the index of the example
+        # containing this span of text.
+        sample_index = sample_mapping[i]
+        tokenized_examples.id.append(examples["id"][sample_index])
+
+        # Set to (-1, -1) the offset_mapping that are not part of the context so it's
+        # easy to determine if a token position is part of the context or not.
+        tokenized_examples.offset_mapping[i] = [
+            (o if sequence_ids[k] == context_index else (-1, -1))
+            for k, o in enumerate(tokenized_examples.offset_mapping[i])
+        ]
+
+    return tokenized_examples
+
+
+def postprocess_predictions_and_labels(
+    predictions: list,
+    dataset: "Dataset",
+    prepared_dataset: "Dataset",
+    cls_token_index: int,
+) -> tuple[list[dict], list[dict]]:
+    """Postprocess the predictions and labels, to allow easier metric computation.
+
+    Args:
+        predictions:
+            A pair of (start_logits, end_logits) predictions.
+        dataset:
+            The dataset containing the examples.
+        prepared_dataset:
+            The dataset containing the prepared examples.
+        cls_token_index:
+            The index of the CLS token.
+
+    Returns:
+        The postprocessed predictions and labels.
+    """
+    # Extract the logits from the predictions
+    all_start_logits = predictions[0]
+    all_end_logits = predictions[1]
+
+    # Build a map from an example to its corresponding features, being the blocks of
+    # text from the context that we're feeding into the model. An example can have
+    # multiple features/blocks if it has a long context.
+    id_to_index = {k: i for i, k in enumerate(dataset["id"])}
+    features_per_example = defaultdict(list)
+    for i, feature in enumerate(prepared_dataset):
+        id = feature["id"]
+        example_index = id_to_index[id]
+        features_per_example[example_index].append(i)
+
+    # Loop over all the examples
+    predictions = list()
+    labels = list()
+    for example_index, example in enumerate(dataset):
+        # Extract the best valid answer associated with the current example
+        best_answer = find_best_answer(
+            all_start_logits=all_start_logits,
+            all_end_logits=all_end_logits,
+            prepared_dataset=prepared_dataset,
+            feature_indices=features_per_example[example_index],
+            context=example["context"],
+            max_answer_length=30,
+            num_best_logits=20,
+            min_null_score=0.0,
+            cls_token_index=cls_token_index,
+        )
+
+        # Create the final prediction dictionary, to be added to the list of
+        # predictions
+        prediction = dict(
+            id=example["id"], prediction_text=best_answer, no_answer_probability=0.0
+        )
+
+        # Add the answer to the list of predictions
+        predictions.append(prediction)
+
+        # Create the associated reference dictionary, to be added to the list of
+        # references
+        label = dict(
+            id=example["id"],
+            answers=dict(
+                text=example["answers"]["text"],
+                answer_start=example["answers"]["answer_start"],
+            ),
+        )
+
+        # Add the answer and label to the list of predictions and labels, respectively
+        labels.append(label)
+
+    return predictions, labels
+
+
+def find_best_answer(
+    all_start_logits: np.ndarray,
+    all_end_logits: np.ndarray,
+    prepared_dataset: "Dataset",
+    feature_indices: list[int],
+    context: str,
+    max_answer_length: int,
+    num_best_logits: int,
+    min_null_score: float,
+    cls_token_index: int,
+) -> str:
+    """Find the best answer for a given example.
+
+    Args:
+        all_start_logits:
+            The start logits for all the features.
+        all_end_logits:
+            The end logits for all the features.
+        prepared_dataset:
+            The dataset containing the prepared examples.
+        feature_indices:
+            The indices of the features associated with the current example.
+        context:
+            The context of the example.
+        max_answer_length:
+            The maximum length of the answer.
+        num_best_logits:
+            The number of best logits to consider.
+        min_null_score:
+            The minimum score an answer can have.
+        cls_token_index:
+            The index of the CLS token.
+
+    Returns:
+        The best answer for the example.
+    """
+    # Loop through all the features associated to the current example
+    valid_answers = list()
+    for feature_index in feature_indices:
+        # Get the features associated with the current example
+        features = prepared_dataset[feature_index]
+
+        # Get the predictions of the model for this feature
+        start_logits = all_start_logits[feature_index]
+        end_logits = all_end_logits[feature_index]
+
+        # Update minimum null prediction
+        cls_index = features["input_ids"].index(cls_token_index)
+        feature_null_score = (start_logits[cls_index] + end_logits[cls_index]).item()
+        if min_null_score < feature_null_score:
+            min_null_score = feature_null_score
+
+        # Find the valid answers for the feature
+        valid_answers_for_feature = find_valid_answers(
+            start_logits=start_logits,
+            end_logits=end_logits,
+            offset_mapping=features["offset_mapping"],
+            context=context,
+            max_answer_length=max_answer_length,
+            num_best_logits=num_best_logits,
+            min_null_score=min_null_score,
+        )
+        valid_answers.extend(valid_answers_for_feature)
+
+    # In the very rare edge case we have not a single non-null prediction, we create a
+    # fake prediction to avoid failure
+    if not valid_answers:
+        return ""
+
+    # Otherwise, we select the answer with the largest score as the best answer, and
+    # return it
+    best_answer_dict = sorted(valid_answers, key=lambda x: x["score"], reverse=True)[0]
+    return best_answer_dict["text"]
+
+
+def find_valid_answers(
+    start_logits: np.ndarray,
+    end_logits: np.ndarray,
+    offset_mapping: list[tuple[int, int]],
+    context: str,
+    max_answer_length: int,
+    num_best_logits: int,
+    min_null_score: float,
+) -> list[dict]:
+    """Find the valid answers from the start and end indexes.
+
+    Args:
+        start_logits:
+            The logits for the start of the answer.
+        end_logits:
+            The logits for the end of the answer.
+        offset_mapping:
+            The offset mapping, being a list of pairs of integers for each token index,
+            containing the start and end character index in the original context.
+        context:
+            The context of the example.
+        max_answer_length:
+            The maximum length of the answer.
+        num_best_logits:
+            The number of best logits to consider. Note that this function will run in
+            O(`num_best_logits` ^ 2) time.
+        min_null_score:
+            The minimum score an answer can have.
+
+    Returns:
+        A list of the valid answers, each being a dictionary with keys "text" and
+        "score", the score being the sum of the start and end logits.
+    """
+    # Fetch the top-k predictions for the start- and end token indices
+    start_indexes = np.argsort(start_logits)[-1 : -num_best_logits - 1 : -1].tolist()
+    end_indexes = np.argsort(end_logits)[-1 : -num_best_logits - 1 : -1].tolist()
+
+    # We loop over all combinations of starting and ending indexes for valid answers
+    valid_answers = list()
+    for start_index in start_indexes:
+        for end_index in end_indexes:
+            # If the starting or ending index is out-of-scope, meaning that they are
+            # either out of bounds or correspond to part of the input_ids that are not
+            # in the context, then we skip this index
+            if (
+                start_index >= len(offset_mapping)
+                or end_index >= len(offset_mapping)
+                or tuple(offset_mapping[start_index]) == (-1, -1)
+                or tuple(offset_mapping[end_index]) == (-1, -1)
+            ):
+                continue
+
+            # Do not consider answers with a length that is either negative or greater
+            # than the context length
+            max_val = max_answer_length + start_index - 1
+            if end_index < start_index or end_index > max_val:
+                continue
+
+            # If we got to this point then the answer is valid, so we store the
+            # corresponding start- and end character indices in the original context,
+            # and from these extract the answer
+            start_char = offset_mapping[start_index][0]
+            end_char = offset_mapping[end_index][1]
+            text = context[start_char:end_char]
+
+            # Compute the score of the answer, being the sum of the start and end
+            # logits. Intuitively, this indicates how likely the answer is to be
+            # correct, and allows us to pick the best valid answer.
+            score = start_logits[start_index] + end_logits[end_index]
+
+            # Add the answer to the list of valid answers, if the score is greater
+            # than the minimum null score
+            if score > min_null_score:
+                valid_answers.append(dict(score=score, text=text))
+
+    return valid_answers
diff --git a/src/scandeval/task_utils/sequence_classification.py b/src/scandeval/task_utils/sequence_classification.py
new file mode 100644
index 000000000..e355a6159
--- /dev/null
+++ b/src/scandeval/task_utils/sequence_classification.py
@@ -0,0 +1,232 @@
+"""Utility functions related to the sequence-classification supertask."""
+
+import logging
+import re
+import typing as t
+
+import evaluate
+import Levenshtein
+import numpy as np
+from evaluate import EvaluationModule
+
+from ..data_models import BenchmarkConfig, GenerativeModelOutput
+from ..utils import raise_if_model_output_contains_nan_values
+
+if t.TYPE_CHECKING:
+    from ..data_models import DatasetConfig
+    from ..types import Labels, Predictions
+
+
+logger = logging.getLogger("scandeval")
+
+
+def compute_metrics(
+    model_outputs_and_labels: tuple["Predictions", "Labels"],
+    id2label: dict[int, str],
+    dataset_config: "DatasetConfig",
+    benchmark_config: "BenchmarkConfig",
+) -> dict[str, float]:
+    """Compute the metrics needed for evaluation.
+
+    Args:
+        model_outputs_and_labels:
+            The first sequence contains the model outputs and the second sequence
+            contains the true labels.
+        id2label:
+            Conversion of indices to labels.
+        dataset_config:
+            The configuration of the dataset.
+        benchmark_config:
+            The configuration of the benchmark.
+
+    Returns:
+        A dictionary with the names of the metrics as keys and the metric values as
+        values.
+    """
+    model_outputs, labels = model_outputs_and_labels
+    label2id = {label: idx for idx, label in id2label.items()}
+    raise_if_model_output_contains_nan_values(model_output=model_outputs)
+
+    metrics = {
+        metric_cfg.name: (
+            evaluate.load(
+                path=metric_cfg.huggingface_id, cache_dir=benchmark_config.cache_dir
+            )
+            if metric_cfg.huggingface_id != ""
+            else None
+        )
+        for metric_cfg in dataset_config.task.metrics
+    }
+
+    model_output_dtype = np.asarray(model_outputs).dtype
+    if model_output_dtype in [np.float16, np.float32, np.float64]:
+        predictions = np.asarray(model_outputs).argmax(axis=-1)
+    else:
+        predictions = model_outputs
+
+    prompt_label_to_label_mapping = {
+        prompt_label: label
+        for label, prompt_label in dataset_config.prompt_label_mapping.items()
+    }
+    predictions = [
+        (
+            label2id[prompt_label_to_label_mapping[pred.lower()]]
+            if isinstance(pred, str)
+            else pred
+        )
+        for pred in predictions
+    ]
+
+    label_ids = [
+        label2id[label.lower()] if isinstance(label, str) else label for label in labels
+    ]
+
+    results: dict[str, float] = dict()
+    for cfg in dataset_config.task.metrics:
+        metric = metrics[cfg.name]
+        assert isinstance(metric, EvaluationModule)
+        score_dict: dict[str, float] | None = metric.compute(
+            predictions=predictions, references=label_ids, **cfg.compute_kwargs
+        )
+
+        # The metric returns None if we are running on multi-GPU and the current
+        # process is not the main process
+        if score_dict is not None:
+            scores = score_dict[cfg.results_key]
+            if isinstance(scores, list):
+                scores = sum(scores) / len(scores)
+            results[cfg.name] = scores
+
+    return results
+
+
+def extract_labels_from_generation(
+    input_batch: dict[str, list],
+    model_output: GenerativeModelOutput,
+    dataset_config: "DatasetConfig",
+) -> list[str]:
+    """Extract the predicted labels from the generated output.
+
+    Args:
+        input_batch:
+            The input batch, where the keys are the feature names and the values
+            are lists with the feature values.
+        model_output:
+            The raw generated output of the model.
+        dataset_config:
+            The configuration of the dataset.
+
+    Returns:
+        The predicted labels.
+    """
+    if model_output.scores is not None:
+        return get_closest_logprobs_labels(
+            generation_logprobs=model_output.scores, dataset_config=dataset_config
+        )
+    else:
+        return get_closest_word_edit_labels(
+            generated_sequences=model_output.sequences, dataset_config=dataset_config
+        )
+
+
+def get_closest_logprobs_labels(
+    generation_logprobs: list[list[list[tuple[str, float]]]],
+    dataset_config: "DatasetConfig",
+) -> list[str]:
+    """Get the labels with the highest predicted logprob value.
+
+    In case a candidate label is split into multiple tokens, we only use the first
+    token to compute the logprob value. E.g., if the candidate label "positive" is
+    tokenised as ["pos", "itive"], we only use the logprob value of "pos" to
+    represent the logprob value of the entire label.
+
+    Args:
+        generation_logprobs:
+            The logprobs of the generated tokens, for all samples in the batch. Of shape
+            (batch_size, num_tokens, num_logprobs).
+        dataset_config:
+            The configuration of the dataset.
+
+    Returns:
+        The predicted labels.
+
+    Raises:
+        InvalidBenchmark:
+            If no candidate label can be found for any of the generated labels.
+    """
+    english_labels = list(dataset_config.id2label.values())
+    english2local = dataset_config.prompt_label_mapping
+    candidate_labels = [
+        english2local[lbl].lower() for lbl in english_labels
+    ] + english_labels
+
+    output_labels: list[str] = list()
+    given_warning: bool = False
+    for sample in generation_logprobs:
+        for logprob_list in sample:
+            generated_labels = [
+                re.sub(
+                    pattern=r"^[^a-zæøåüöä]+|[^a-zæøåüöä]+$",
+                    repl="",
+                    string=label.lower(),
+                )
+                for label, _ in logprob_list
+            ]
+            generated_labels = [label for label in generated_labels if label != ""]
+
+            # We want to use the first generated label which starts with a candidate
+            # label, as the output label
+            output_label: str | None = None
+            for generated_label in generated_labels:
+                candidate_output_labels = [
+                    candidate_label
+                    for candidate_label in candidate_labels
+                    if candidate_label.startswith(generated_label)
+                ]
+                if candidate_output_labels:
+                    output_label = candidate_output_labels[0]
+                    break
+
+            if output_label is not None:
+                output_label = english2local.get(output_label, output_label)
+                output_labels.append(output_label)
+                break
+        else:
+            if not given_warning:
+                logger.debug(
+                    "Could not find a candidate label for any of the generated "
+                    f"labels in the sample {sample}. Using {candidate_labels[0]!r} "
+                    "as the output label."
+                )
+                given_warning = True
+            output_labels.append(candidate_labels[0])
+    return output_labels
+
+
+def get_closest_word_edit_labels(
+    generated_sequences: list[str], dataset_config: "DatasetConfig"
+) -> list[str]:
+    """Get the labels with the smallest edit distance to the predicted labels.
+
+    Args:
+        generated_sequences:
+            The generated sequences from the model.
+        dataset_config:
+            The configuration of the dataset.
+
+    Returns:
+        The candidate labels with the smallest edit distance to the predicted labels.
+    """
+    candidate_labels = [
+        dataset_config.prompt_label_mapping[lbl]
+        for lbl in dataset_config.id2label.values()
+    ]
+    new_predicted_labels: list[str] = list()
+    for predicted_label in generated_sequences:
+        edit_distances = [
+            Levenshtein.distance(s1=predicted_label.lower(), s2=candidate_label.lower())
+            for candidate_label in candidate_labels
+        ]
+        closest_label = candidate_labels[np.argmin(edit_distances).item()]
+        new_predicted_labels.append(closest_label)
+    return new_predicted_labels
diff --git a/src/scandeval/task_utils/text_to_text.py b/src/scandeval/task_utils/text_to_text.py
new file mode 100644
index 000000000..657b4c03b
--- /dev/null
+++ b/src/scandeval/task_utils/text_to_text.py
@@ -0,0 +1,153 @@
+"""Utility functions related to the text-to-text supertask."""
+
+import logging
+import typing as t
+
+import evaluate
+import numpy as np
+from evaluate import EvaluationModule
+
+from ..constants import METRIC_ATTRIBUTES_TAKING_UP_MEMORY
+from ..data_models import BenchmarkConfig, DatasetConfig, GenerativeModelOutput
+from ..exceptions import InvalidBenchmark
+from ..utils import (
+    HiddenPrints,
+    clear_memory,
+    raise_if_model_output_contains_nan_values,
+)
+
+if t.TYPE_CHECKING:
+    from ..types import Labels, Predictions
+
+
+logger = logging.getLogger("scandeval")
+
+
+def compute_metrics(
+    model_outputs_and_labels: tuple["Predictions", "Labels"],
+    id2label: dict[int, str],
+    dataset_config: "DatasetConfig",
+    benchmark_config: "BenchmarkConfig",
+) -> dict[str, float]:
+    """Compute the metrics needed for evaluation.
+
+    Args:
+        model_outputs_and_labels:
+            The first sequence contains the model outputs and the second sequence
+            contains the true labels.
+        id2label:
+            Conversion of indices to labels.
+        dataset_config:
+            The configuration of the dataset.
+        benchmark_config:
+            The configuration of the benchmark.
+
+    Returns:
+        A dictionary with the names of the metrics as keys and the metric values as
+        values.
+    """
+    model_outputs, labels = model_outputs_and_labels
+    raise_if_model_output_contains_nan_values(model_output=model_outputs)
+
+    metrics = {
+        metric_cfg.name: (
+            evaluate.load(
+                path=metric_cfg.huggingface_id, cache_dir=benchmark_config.cache_dir
+            )
+            if metric_cfg.huggingface_id != ""
+            else None
+        )
+        for metric_cfg in dataset_config.task.metrics
+    }
+
+    model_output_dtype = np.asarray(model_outputs).dtype
+    output_is_prob = model_output_dtype in [np.float16, np.float32, np.float64]
+    if output_is_prob:
+        predictions = np.asarray(model_outputs).argmax(axis=-1)
+    else:
+        predictions = model_outputs
+
+    results: dict[str, float] = dict()
+    for cfg in dataset_config.task.metrics:
+        metric = metrics[cfg.name]
+        assert isinstance(metric, EvaluationModule)
+
+        # Some metrics can be computed on hardware accelerators. In this case we
+        # start by setting the device to the same device as the model
+        if cfg.compute_kwargs.get("device", None) == "auto":
+            cfg.compute_kwargs["device"] = benchmark_config.device.type
+
+        while True:
+            try:
+                with HiddenPrints():
+                    score_dict: dict[str, float] | None = metric.compute(
+                        predictions=predictions, references=labels, **cfg.compute_kwargs
+                    )
+
+                # Clear the cache of the BERTScorer to avoid memory leaks
+                for attribute in METRIC_ATTRIBUTES_TAKING_UP_MEMORY:
+                    if hasattr(metric, attribute):
+                        delattr(metric, attribute)
+
+                clear_memory()
+                break
+            except Exception as e:
+                # Clear the cache of the BERTScorer to avoid memory leaks
+                if hasattr(metric, "cached_bertscorer"):
+                    del metric.cached_bertscorer
+                    clear_memory()
+
+                oom_error = [
+                    "CUDA out of memory",
+                    "CUDA error",
+                    "MPS backend out of memory",
+                ]
+                if not any(error in str(e) for error in oom_error):
+                    raise InvalidBenchmark(str(e))
+
+                if cfg.compute_kwargs.get("batch_size", 1) > 1:
+                    batch_size = cfg.compute_kwargs["batch_size"]
+                    cfg.compute_kwargs["batch_size"] = batch_size // 2
+                    logger.debug(
+                        "Out of memory error occurred during the computation of "
+                        f"the metric {cfg.pretty_name}. Reducing the batch size to "
+                        f"{cfg.compute_kwargs['batch_size']}."
+                    )
+                elif cfg.compute_kwargs.get("device", "cpu") != "cpu":
+                    cfg.compute_kwargs["batch_size"] = 32
+                    cfg.compute_kwargs["device"] = "cpu"
+                    logger.debug(
+                        "Out of memory error occurred during the computation of "
+                        f"the metric {cfg.pretty_name}. Moving the computation to "
+                        "the CPU."
+                    )
+                else:
+                    raise InvalidBenchmark(str(e))
+
+        # The metric returns None if we are running on multi-GPU and the current
+        # process is not the main process
+        if score_dict is not None:
+            scores = score_dict[cfg.results_key]
+            if isinstance(scores, list):
+                scores = sum(scores) / len(scores)
+            results[cfg.name] = scores
+
+    return results
+
+
+def extract_labels_from_generation(
+    input_batch: dict[str, list], model_output: "GenerativeModelOutput"
+) -> list[t.Any]:
+    """Extract the predicted labels from the generated output.
+
+    Args:
+        input_batch:
+            The input batch, where the keys are the feature names and the values
+            are lists with the feature values.
+        model_output:
+            The raw generated output of the model.
+
+    Returns:
+        The predicted labels.
+    """
+    return model_output.sequences
diff --git a/src/scandeval/task_utils/token_classification.py b/src/scandeval/task_utils/token_classification.py
new file mode 100644
index 000000000..08690bbb1
--- /dev/null
+++ b/src/scandeval/task_utils/token_classification.py
@@ -0,0 +1,467 @@
+"""Utility functions related to the token-classification supertask."""
+
+import importlib.util
+import logging
+import re
+import typing as t
+from copy import deepcopy
+
+import evaluate
+import numpy as np
+from evaluate import EvaluationModule
+from transformers import PreTrainedTokenizer
+
+from ..data_models import BenchmarkConfig, DatasetConfig, GenerativeModelOutput
+from ..exceptions import InvalidBenchmark, NeedsExtraInstalled
+from ..utils import raise_if_model_output_contains_nan_values
+
+if t.TYPE_CHECKING:
+    from transformers import BatchEncoding
+
+    from ..types import Labels, Predictions
+
+if importlib.util.find_spec("demjson3") is not None:
+    import demjson3
+
+
+logger = logging.getLogger("scandeval")
+
+
+def compute_metrics(
+    model_outputs_and_labels: tuple["Predictions", "Labels"],
+    id2label: dict[int, str],
+    has_misc_tags: bool,
+    dataset_config: "DatasetConfig",
+    benchmark_config: "BenchmarkConfig",
+) -> dict[str, float]:
+    """Compute the metrics needed for evaluation.
+
+    Args:
+        model_outputs_and_labels:
+            The first array contains the probability predictions and the second
+            array contains the true labels.
+        id2label:
+            Conversion of indices to labels.
+        has_misc_tags:
+            Whether the dataset has MISC tags.
+        dataset_config:
+            The configuration of the dataset.
+        benchmark_config:
+            The configuration of the benchmark.
+
+    Returns:
+        A dictionary with the names of the metrics as keys and the metric values as
+        values.
+    """
+    model_outputs, labels = model_outputs_and_labels
+    raise_if_model_output_contains_nan_values(model_output=model_outputs)
+
+    metrics = {
+        metric_cfg.name: (
+            evaluate.load(
+                path=metric_cfg.huggingface_id, cache_dir=benchmark_config.cache_dir
+            )
+            if metric_cfg.huggingface_id != ""
+            else None
+        )
+        for metric_cfg in dataset_config.task.metrics
+    }
+
+    predictions: list[list[str]]
+    if not isinstance(model_outputs[0][0], str):
+        raw_predictions: list[list[int]] = np.argmax(model_outputs, axis=-1).tolist()
+
+        # Remove ignored index (special tokens)
+        predictions = [
+            [
+                id2label[pred_id]
+                for pred_id, lbl_id in zip(pred, label)
+                if lbl_id != -100
+            ]
+            for pred, label in zip(raw_predictions, labels)
+        ]
+        labels = [
+            [
+                (
+                    id2label[int(lbl_id)]
+                    if isinstance(lbl_id, int) or isinstance(lbl_id, np.int_)
+                    else lbl_id
+                )
+                for lbl_id in label
+                if lbl_id != -100
+            ]
+            for label in labels
+        ]
+
+    else:
+        predictions = model_outputs  # type: ignore[assignment]
+
+    # Replace predicted tag with either MISC or O tags if they are not part of the
+    # dataset
+    labels_without_misc = {
+        label
+        for label in dataset_config.id2label.values()
+        if label not in {"b-misc", "i-misc"}
+    }
+    ner_tag: str
+    for i, prediction_list in enumerate(predictions):
+        for j, ner_tag in enumerate(prediction_list):
+            if ner_tag not in labels_without_misc:
+                if has_misc_tags and ner_tag[:2] == "b-":
+                    predictions[i][j] = "b-misc"
+                elif has_misc_tags and ner_tag[:2] == "i-":
+                    predictions[i][j] = "i-misc"
+                else:
+                    predictions[i][j] = "o"
+
+    # Remove MISC labels from predictions
+    predictions_no_misc = deepcopy(predictions)
+    for i, prediction_list in enumerate(predictions_no_misc):
+        for j, ner_tag in enumerate(prediction_list):
+            if ner_tag[-4:] == "misc":
+                predictions_no_misc[i][j] = "o"
+
+    # Remove MISC labels from labels
+    labels_no_misc: list[list[str]] = deepcopy(labels)  # type: ignore[arg-type]
+    for i, label_list in enumerate(labels_no_misc):
+        for j, ner_tag in enumerate(label_list):
+            if (
+                isinstance(ner_tag, str)
+                and len(ner_tag) >= 4
+                and ner_tag[-4:] == "misc"
+            ):
+                labels_no_misc[i][j] = "o"
+
+    # Compute the metrics
+    # We manually set the F1 metric to be 100% if both the labels and the models
+    # have no NER tags in them, since this causes an error with the `compute`
+    # method otherwise
+    predictions_all_zero = all(
+        all(ner_tag == "o" for ner_tag in prediction_list)
+        for prediction_list in predictions
+    )
+    labels_all_zero = all(
+        all(ner_tag == "o" for ner_tag in label_list) for label_list in labels
+    )
+    if predictions_all_zero and labels_all_zero:
+        results = dict(overall_f1=1.0)
+    else:
+        metric = metrics["micro_f1"]
+        assert isinstance(metric, EvaluationModule)
+        results = metric.compute(predictions=predictions, references=labels)
+
+    # Compute the metrics without MISC tags
+    # We manually set the F1 metric to be 100% if both the labels and the models
+    # have no NER tags in them, since this causes an error with the `compute`
+    # method otherwise
+    predictions_no_misc_all_zero = all(
+        all(ner_tag == "o" for ner_tag in prediction_list)
+        for prediction_list in predictions_no_misc
+    )
+    labels_no_misc_all_zero = all(
+        all(ner_tag == "o" for ner_tag in label_list) for label_list in labels_no_misc
+    )
+    if predictions_no_misc_all_zero and labels_no_misc_all_zero:
+        results_no_misc = dict(overall_f1=1.0)
+    else:
+        metric = metrics["micro_f1_no_misc"]
+        assert isinstance(metric, EvaluationModule)
+        results_no_misc = metric.compute(
+            predictions=predictions_no_misc, references=labels_no_misc
+        )
+
+    # Raise error if the metrics are invalid
+    if results is None or results_no_misc is None:
+        raise InvalidBenchmark("The predictions and labels are not of the same length.")
+
+    return dict(
+        micro_f1_no_misc=results_no_misc["overall_f1"], micro_f1=results["overall_f1"]
+    )
+
+
+def extract_labels_from_generation(
+    input_batch: dict[str, list],
+    model_output: "GenerativeModelOutput",
+    dataset_config: "DatasetConfig",
+) -> list[t.Any]:
+    """Extract the predicted labels from the generated output.
+
+    Args:
+        input_batch:
+            The input batch, where the keys are the feature names and the values
+            are lists with the feature values.
+        model_output:
+            The raw generated output of the model.
+        dataset_config:
+            The configuration of the dataset.
+
+    Returns:
+        The predicted labels.
+    """
+    if importlib.util.find_spec("demjson3") is None:
+        raise NeedsExtraInstalled(extra="generative")
+
+    raw_predictions = model_output.sequences
+
+    # Attempt to extract the JSON dictionary from the predictions
+    json_regex = r"\{.+?\}"
+    json_matches = [
+        re.search(pattern=json_regex, string=raw_prediction, flags=re.DOTALL)
+        or raw_prediction
+        for raw_prediction in raw_predictions
+    ]
+    raw_predictions = [
+        json_match.group() if isinstance(json_match, re.Match) else json_match
+        for json_match in json_matches
+    ]
+
+    tokens = input_batch["tokens"]
+    predicted_labels: list[list[str]] = [["o"] * len(token_ids) for token_ids in tokens]
+    for idx, raw_prediction in enumerate(raw_predictions):
+        try:
+            json_output = demjson3.decode(txt=raw_prediction)
+            if not isinstance(json_output, dict):
+                logger.debug(
+                    "The model output is not a JSON dictionary, so cannot parse "
+                    f"it. Skipping. Here is the output: {raw_prediction}"
+                )
+                continue
+            elif not all(isinstance(key, str) for key in json_output.keys()):
+                logger.debug(
+                    "The model output is not a JSON dictionary with string keys, "
+                    "so cannot parse it. Skipping. Here is the output: "
+                    f"{raw_prediction}"
+                )
+                continue
+            elif not all(isinstance(value, list) for value in json_output.values()):
+                logger.debug(
+                    "The model output is not a JSON dictionary with list values, "
+                    "so cannot parse it. Skipping. Here is the output: "
+                    f"{raw_prediction}"
+                )
+                continue
+            prediction_dict: dict[str, list[str]] = json_output
+        except demjson3.JSONDecodeError:
+            logger.debug(
+                "The model output is not valid JSON, so cannot parse it. Skipping. "
+                f"Here is the output: {raw_prediction!r}"
+            )
+            continue
+
+        prompt_label_mapping = dataset_config.prompt_label_mapping
+        for prompt_tag_name, named_entities in prediction_dict.items():
+            try:
+                tag_name = [
+                    tag[2:]
+                    for tag, prompt_tag in prompt_label_mapping.items()
+                    if prompt_tag == prompt_tag_name
+                ][0]
+            except IndexError:
+                logger.debug(
+                    "The model produced an invalid prompt tag name, "
+                    f"{prompt_tag_name}. Skipping."
+                )
+                continue
+
+            named_entities = [str(named_entity) for named_entity in named_entities]
+            for named_entity in named_entities:
+                for ne_idx, named_entity_word in enumerate(named_entity.split()):
+                    for token_idx, token in enumerate(tokens[idx]):
+                        if named_entity_word in token:
+                            if ne_idx == 0:
+                                predicted_labels[idx][token_idx] = f"b-{tag_name}"
+                            elif (
+                                predicted_labels[idx][token_idx] == "o"
+                                and predicted_labels[idx][token_idx - 1][2:] == tag_name
+                            ):
+                                predicted_labels[idx][token_idx] = f"i-{tag_name}"
+    return predicted_labels
+
+
+def tokenize_and_align_labels(
+    examples: dict, tokenizer: "PreTrainedTokenizer", label2id: dict[str, int]
+) -> "BatchEncoding":
+    """Tokenise all texts and align the labels with them.
+
+    Args:
+        examples:
+            The examples to be tokenised.
+        tokenizer:
+            A pretrained tokenizer.
+        label2id:
+            A dictionary that converts NER tags to IDs.
+
+    Returns:
+        A dictionary containing the tokenized data as well as labels.
+    """
+    # Tokenize the texts. We use the `is_split_into_words` argument here because
+    # the texts in our dataset are lists of words (with a label for each word)
+    tokenized_inputs = tokenizer(
+        examples["tokens"], is_split_into_words=True, truncation=True, padding=True
+    )
+
+    # Extract a mapping between all the tokens and their corresponding word. If the
+    # tokenizer is of a "fast" variant then this can be accessed through the
+    # `word_ids` method. Otherwise, we have to extract it manually.
+    all_labels: list[list[int]] = list()
+    labels: list[str]
+    word_ids: list[int | None]
+    for i, labels in enumerate(examples["labels"]):
+        # Try to get the word IDs from the tokenizer
+        try:
+            word_ids = tokenized_inputs.word_ids(batch_index=i)
+
+        # If the tokenizer is not of a "fast" variant, we have to extract the word
+        # IDs manually
+        except ValueError:
+            # Get the list of words in the document
+            words: list[str] = examples["tokens"][i]
+
+            # Get the list of token IDs in the document
+            tok_ids: list[int] = tokenized_inputs.input_ids[i]
+
+            # Decode the token IDs
+            tokens = tokenizer.convert_ids_to_tokens(tok_ids)
+            assert isinstance(tokens, list)
+
+            # Remove prefixes from the tokens
+            prefixes_to_remove = ["▁", "##"]
+            for tok_idx, tok in enumerate(tokens):
+                if tok:
+                    for prefix in prefixes_to_remove:
+                        if tok.startswith(prefix):
+                            tokens[tok_idx] = tok[len(prefix) :]
+
+            # Replace UNK tokens with the correct word
+            tokens = handle_unk_tokens(tokenizer=tokenizer, tokens=tokens, words=words)
+
+            # Get list of special tokens. Some tokenizers do not record these
+            # properly, which is why we convert the values to their indices and
+            # then back to strings
+            sp_toks = [
+                tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids(sp_tok))
+                for sp_tok in tokenizer.special_tokens_map.values()
+            ]
+
+            # Replace special tokens with `None`
+            tokens_with_none = [None if tok in sp_toks else tok for tok in tokens]
+
+            # Get the alignment between the words and the tokens, on a character
+            # level
+            word_idxs = [
+                word_idx for word_idx, word in enumerate(words) for _ in str(word)
+            ]
+            token_idxs = [
+                tok_idx
+                for tok_idx, tok_or_none in enumerate(tokens_with_none)
+                for _ in str(tok_or_none)
+                if tok_or_none is not None
+            ]
+            alignment = list(zip(word_idxs, token_idxs))
+
+            # Raise error if there are not as many characters in the words as in
+            # the tokens. This can be due to the use of a different prefix.
+            if len(word_idxs) != len(token_idxs):
+                raise InvalidBenchmark(
+                    "The tokens could not be aligned with the words during manual "
+                    "word-token alignment. It seems that the tokenizer is neither "
+                    "of the fast variant nor of a SentencePiece/WordPiece variant."
+                )
+
+            # Get the aligned word IDs
+            word_ids = list()
+            for tok_idx, tok_or_none in enumerate(tokens_with_none):
+                if tok_or_none is None or tok_or_none == "":
+                    word_ids.append(None)
+                else:
+                    word_idx = [
+                        word_idx
+                        for word_idx, token_idx in alignment
+                        if token_idx == tok_idx
+                    ][0]
+                    word_ids.append(word_idx)
+
+        previous_word_idx: int | None = None
+        label_ids: list[int] = list()
+        for word_id in word_ids:
+            # Special tokens have a word id that is None. We set the label to -100
+            # so they are automatically ignored in the loss function
+            if word_id is None:
+                label_ids.append(-100)
+
+            # We set the label for the first token of each word
+            elif word_id != previous_word_idx:
+                label = labels[word_id]
+                try:
+                    label_id = label2id[label.lower()]
+                except KeyError:
+                    msg = f"The label {label} was not found in the model's config."
+                    raise InvalidBenchmark(msg)
+                label_ids.append(label_id)
+
+            # For the other tokens in a word, we set the label to -100
+            else:
+                label_ids.append(-100)
+
+            previous_word_idx = word_id
+
+        all_labels.append(label_ids)
+    tokenized_inputs["labels"] = all_labels
+    return tokenized_inputs
+
+
+def handle_unk_tokens(
+    tokenizer: "PreTrainedTokenizer", tokens: list[str], words: list[str]
+) -> list[str]:
+    """Replace unknown tokens in the tokens with the corresponding word.
+
+    Args:
+        tokenizer:
+            The tokenizer used to tokenize the words.
+        tokens:
+            The list of tokens.
+        words:
+            The list of words.
+
+    Returns:
+        The list of tokens with unknown tokens replaced by the corresponding word.
+    """
+    # Locate the token indices of the unknown tokens
+    token_unk_idxs = [i for i, tok in enumerate(tokens) if tok == tokenizer.unk_token]
+
+    # Locate the word indices of the words which contain an unknown token
+    word_unk_idxs = [
+        i
+        for i, word in enumerate(words)
+        if tokenizer.unk_token
+        in tokenizer.convert_ids_to_tokens(
+            tokenizer.encode(word, add_special_tokens=False)
+        )
+    ]
+
+    # Iterate over the token index and word index pairs
+    for tok_idx, word_idx in zip(token_unk_idxs, word_unk_idxs):
+        # Fetch the word
+        word = words[word_idx]
+
+        # Tokenize the word, which is now a list containing at least one UNK token
+        tokens_with_unk = tokenizer.convert_ids_to_tokens(
+            tokenizer.encode(word, add_special_tokens=False)
+        )
+
+        # Iterate over the tokens in the word
+        for possible_unk_token in tokens_with_unk:
+            # If the token is not an UNK token then we remove the first occurence
+            # of the content of this token from the word. The result of the `word`
+            # variable will be the content of the UNK token.
+            # NOTE: This is a bit hacky and not bulletproof. For instance, if the
+            # word is "1925-1950" and the tokenizer splits it into ["[UNK]", "-",
+            # "19", "50"], then the result will be 2519 instead of 1925. This
+            # happens almost never, however, so we can live with it.
+            if possible_unk_token != tokenizer.unk_token:
+                word = word.replace(possible_unk_token, "", 1)
+
+        # Replace the token with the word
+        tokens[tok_idx] = word
+
+    return tokens
diff --git a/src/scandeval/tasks.py b/src/scandeval/tasks.py
index 2b48931d2..159c32204 100644
--- a/src/scandeval/tasks.py
+++ b/src/scandeval/tasks.py
@@ -1,6 +1,6 @@
 """All benchmarks tasks used in ScandEval."""
 
-from .config import MetricConfig, Task
+from .data_models import MetricConfig, Task
 
 
 def get_all_tasks() -> dict[str, Task]:
diff --git a/src/scandeval/text_to_text.py b/src/scandeval/text_to_text.py
deleted file mode 100644
index b5be2b22a..000000000
--- a/src/scandeval/text_to_text.py
+++ /dev/null
@@ -1,321 +0,0 @@
-"""Text-to-text generation benchmark dataset."""
-
-import logging
-import random
-import re
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-from evaluate import EvaluationModule
-from transformers.data.data_collator import DataCollatorWithPadding
-
-from .benchmark_dataset import BenchmarkDataset
-from .exceptions import InvalidBenchmark
-from .generation import extract_raw_predictions
-from .utils import (
-    METRIC_ATTRIBUTES_TAKING_UP_MEMORY,
-    HiddenPrints,
-    clear_memory,
-    convert_prompt_to_instruction,
-    raise_if_model_output_contains_nan_values,
-)
-
-if TYPE_CHECKING:
-    from datasets.arrow_dataset import Dataset
-    from transformers import BatchEncoding, PreTrainedModel
-    from transformers.utils import ModelOutput
-
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Labels, Predictions
-
-
-logger = logging.getLogger(__package__)
-
-
-class TextToText(BenchmarkDataset):
-    """Text-to-text benchmark dataset.
-
-    Args:
-        dataset_config:
-            The dataset configuration.
-        benchmark_config:
-            The benchmark configuration.
-
-    Attributes:
-        dataset_config:
-            The configuration of the dataset.
-        benchmark_config:
-            The configuration of the benchmark.
-    """
-
-    def _preprocess_data(self, dataset: "Dataset", **kwargs) -> "Dataset":
-        """Preprocess a dataset by tokenizing and aligning the labels.
-
-        Args:
-            dataset:
-                The dataset to preprocess.
-            kwargs:
-                Extra keyword arguments containing objects used in preprocessing the
-                dataset.
-
-        Returns:
-            The preprocessed dataset.
-        """
-        tokenizer: "Tokenizer" = kwargs["tokenizer"]
-
-        def tokenise(examples: dict) -> "BatchEncoding":
-            return tokenizer(text=examples["text"], truncation=True, padding=False)
-
-        tokenised = dataset.map(
-            tokenise, batched=True, load_from_cache_file=False, keep_in_memory=True
-        )
-
-        return tokenised
-
-    def _load_data_collator(
-        self,
-        tokenizer: "Tokenizer | None" = None,
-        model: "PreTrainedModel | GenerativeModel | None" = None,
-    ):
-        """Load the data collator used to prepare samples during finetuning.
-
-        Args:
-            tokenizer:
-                A pretrained tokenizer. Can be None if the tokenizer is not used in the
-                initialisation of the data collator. Defaults to None.
-            model:
-                A pretrained model. Can be None if the model is not used. Defaults to
-                None.
-
-        Returns:
-            The data collator.
-        """
-        return DataCollatorWithPadding(tokenizer, padding="longest")
-
-    def _compute_metrics(
-        self,
-        model_outputs_and_labels: tuple["Predictions", "Labels"],
-        id2label: dict[int, str],
-    ) -> dict[str, float]:
-        """Compute the metrics needed for evaluation.
-
-        Args:
-            model_outputs_and_labels:
-                The first sequence contains the model outputs and the second sequence
-                contains the true labels.
-            id2label:
-                Conversion of indices to labels.
-
-        Returns:
-            A dictionary with the names of the metrics as keys and the metric values as
-            values.
-        """
-        model_outputs, labels = model_outputs_and_labels
-
-        raise_if_model_output_contains_nan_values(model_output=model_outputs)
-
-        model_output_dtype = np.asarray(model_outputs).dtype
-        output_is_prob = model_output_dtype in [np.float16, np.float32, np.float64]
-        if output_is_prob:
-            predictions = np.asarray(model_outputs).argmax(axis=-1)
-        else:
-            predictions = model_outputs
-
-        results: dict[str, float] = dict()
-        for cfg in self.dataset_config.task.metrics:
-            metric = self._metrics[cfg.name]
-            assert isinstance(metric, EvaluationModule)
-
-            # Some metrics can be computed on hardware accelerators. In this case we
-            # start by setting the device to the same device as the model
-            if cfg.compute_kwargs.get("device", None) == "auto":
-                cfg.compute_kwargs["device"] = self.benchmark_config.device.type
-
-            while True:
-                try:
-                    with HiddenPrints():
-                        score_dict: dict[str, float] | None = metric.compute(
-                            predictions=predictions,
-                            references=labels,
-                            **cfg.compute_kwargs,
-                        )
-
-                    # Clear the cache of the BERTScorer to avoid memory leaks
-                    for attribute in METRIC_ATTRIBUTES_TAKING_UP_MEMORY:
-                        if hasattr(metric, attribute):
-                            delattr(metric, attribute)
-
-                    clear_memory()
-                    break
-                except Exception as e:
-                    # Clear the cache of the BERTScorer to avoid memory leaks
-                    if hasattr(metric, "cached_bertscorer"):
-                        del metric.cached_bertscorer
-                        clear_memory()
-
-                    oom_error = [
-                        "CUDA out of memory",
-                        "CUDA error",
-                        "MPS backend out of memory",
-                    ]
-                    if not any(error in str(e) for error in oom_error):
-                        raise InvalidBenchmark(str(e))
-
-                    if cfg.compute_kwargs.get("batch_size", 1) > 1:
-                        batch_size = cfg.compute_kwargs["batch_size"]
-                        cfg.compute_kwargs["batch_size"] = batch_size // 2
-                        logger.debug(
-                            "Out of memory error occurred during the computation of "
-                            f"the metric {cfg.pretty_name}. Reducing the batch size to "
-                            f"{cfg.compute_kwargs['batch_size']}."
-                        )
-                    elif cfg.compute_kwargs.get("device", "cpu") != "cpu":
-                        cfg.compute_kwargs["batch_size"] = 32
-                        cfg.compute_kwargs["device"] = "cpu"
-                        logger.debug(
-                            "Out of memory error occurred during the computation of "
-                            f"the metric {cfg.pretty_name}. Moving the computation to "
-                            "the CPU."
-                        )
-                    else:
-                        raise InvalidBenchmark(str(e))
-
-            # The metric returns None if we are running on multi-GPU and the current
-            # process is not the main process
-            if score_dict is not None:
-                scores = score_dict[cfg.results_key]
-                if isinstance(scores, list):
-                    scores = sum(scores) / len(scores)
-                results[cfg.name] = scores
-
-        return results
-
-    def _extract_few_shot_examples(
-        self, train_dataset: "Dataset", random_seed: int
-    ) -> list[dict[str, Any]]:
-        """Extract few-shot examples from the training dataset.
-
-        Args:
-            train_dataset:
-                The training dataset.
-            random_seed:
-                The random seed to use when extracting the few-shot examples.
-
-        Returns:
-            The few-shot examples.
-        """
-        shuffled_train = train_dataset.shuffle(seed=random_seed)
-        num_few_shots = self.dataset_config.num_few_shot_examples
-        few_shot_examples: list[dict[str, Any]] = list()
-
-        # We pick the few-shot examples one at a time rather than all at once since
-        # we're working with a bootstrapped training dataset, meaning that it will have
-        # duplicates. This ensures that we don't have any duplicates in the few-shot
-        # examples
-        while len(few_shot_examples) < num_few_shots:
-            example = shuffled_train.select(range(1))[0]
-            few_shot_examples.append(example)
-            shuffled_train = shuffled_train.filter(
-                lambda x: x["text"] != example["text"]
-            )
-
-        random.seed(random_seed)
-        random.shuffle(few_shot_examples)
-        return few_shot_examples
-
-    def _apply_few_shot_prompt(
-        self, examples: dict, few_shot_examples: list[dict], tokenizer: "Tokenizer"
-    ) -> dict:
-        """Apply a few-shot prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            few_shot_examples:
-                The examples to be included in the few-shot prompt.
-            tokenizer:
-                The tokenizer to use to encode the few-shot prompt.
-
-        Returns:
-            The examples with the few-shot prompt applied.
-        """
-        # Build the few-shot part of the prompt
-        few_shot_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=example["text"].replace("\n", " ").strip(),
-                target_text=example["target_text"].replace("\n", " ").strip(),
-            )
-            for example in few_shot_examples
-        ]
-        prompt_prefix = ""
-        if self.dataset_config.prompt_prefix:
-            prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"
-        few_shot_prompt = prompt_prefix + "\n\n".join(few_shot_prompts)
-
-        # Add the texts from the examples to the prompts. We remove newlines from the
-        # examples as they have the special function to separate the few-shot examples
-        # from one another
-        new_prompts = [
-            self.dataset_config.prompt_template.format(
-                text=text.replace("\n", " ").strip(), target_text=""
-            ).strip()
-            for text in examples["text"]
-        ]
-
-        final_prompts = [
-            few_shot_prompt + "\n\n" + new_prompt for new_prompt in new_prompts
-        ]
-
-        examples["text"] = final_prompts
-
-        return examples
-
-    def _apply_instruction_prompt(self, examples: dict, tokenizer: "Tokenizer") -> dict:
-        """Apply an instruction prompt to the examples.
-
-        Args:
-            examples:
-                The examples to apply the prompt to.
-            tokenizer:
-                The tokenizer to use to encode the instruction prompt.
-
-        Returns:
-            The examples with the instruction prompt applied.
-        """
-        prompts = [
-            self.dataset_config.instruction_prompt.format(
-                text=re.sub(r"\n+", "\n", text).strip()
-            )
-            for text in examples["text"]
-        ]
-        prompts = [
-            convert_prompt_to_instruction(prompt=prompt, tokenizer=tokenizer)
-            for prompt in prompts
-        ]
-        examples["text"] = prompts
-        return examples
-
-    def _extract_labels_from_generation(
-        self,
-        input_batch: dict[str, list],
-        model_output: "ModelOutput",
-        tokenizer: "Tokenizer",
-    ) -> list[Any]:
-        """Extract the predicted labels from the generated output.
-
-        Args:
-            input_batch:
-                The input batch, where the keys are the feature names and the values
-                are lists with the feature values.
-            model_output:
-                The raw generated output of the model.
-            tokenizer:
-                The tokenizer used together with the model.
-
-        Returns:
-            The predicted labels.
-        """
-        raw_predictions = extract_raw_predictions(
-            generated_sequences=model_output.sequences, tokenizer=tokenizer
-        )
-        return raw_predictions
diff --git a/src/scandeval/types.py b/src/scandeval/types.py
index 987a02aca..c21177f19 100644
--- a/src/scandeval/types.py
+++ b/src/scandeval/types.py
@@ -1,15 +1,63 @@
 """Types used throughout the project."""
 
-from typing import Any, TypeGuard
+import typing as t
 
-import numpy as np
+from numpy.typing import NDArray
 
-ScoreDict = dict[str, dict[str, float] | dict[str, list[dict[str, float]]]]
-Predictions = np.ndarray | list[str] | list[list[str]]
-Labels = np.ndarray | list[str] | list[list[str]]
+if t.TYPE_CHECKING:
+    from .data_models import GenerativeModelOutput
 
 
-def is_list_of_int(x: Any) -> TypeGuard[list[int]]:
+ScoreDict = dict[str, dict[str, float] | list[dict[str, float]]]
+Predictions = NDArray | list[str] | list[list[str]]
+Labels = NDArray | list[str] | list[list[str]]
+
+
+class ComputeMetricsFunction(t.Protocol):
+    """A function used to compute the metrics."""
+
+    def __call__(
+        self,
+        model_outputs_and_labels: tuple[
+            NDArray | list[str] | list[list[str]], NDArray | list[str] | list[list[str]]
+        ],
+        id2label: dict[int, str],
+    ) -> dict[str, float]:
+        """Compute the metrics.
+
+        Args:
+            model_outputs_and_labels:
+                The model outputs and labels.
+            id2label:
+                The mapping from label IDs to labels.
+
+        Returns:
+            The computed metrics.
+        """
+        ...
+
+
+class ExtractLabelsFunction(t.Protocol):
+    """A function used to extract the labels from the generated output."""
+
+    def __call__(
+        self, input_batch: dict[str, list], model_output: "GenerativeModelOutput"
+    ) -> list[str]:
+        """Extract the labels from the generated output.
+
+        Args:
+            input_batch:
+                The input batch.
+            model_output:
+                The model output.
+
+        Returns:
+            The extracted labels.
+        """
+        ...
+
+
+def is_list_of_int(x: t.Any) -> t.TypeGuard[list[int]]:
     """Check if an object is a list of integers.
 
     Args:
@@ -22,7 +70,7 @@ def is_list_of_int(x: Any) -> TypeGuard[list[int]]:
     return isinstance(x, list) and all(isinstance(i, int) for i in x)
 
 
-def is_list_of_list_of_int(x: Any) -> TypeGuard[list[list[int]]]:
+def is_list_of_list_of_int(x: t.Any) -> t.TypeGuard[list[list[int]]]:
     """Check if an object is a list of list of integers.
 
     Args:
@@ -39,7 +87,7 @@ def is_list_of_list_of_int(x: Any) -> TypeGuard[list[list[int]]]:
     )
 
 
-def is_list_of_str(x: Any) -> TypeGuard[list[str]]:
+def is_list_of_str(x: t.Any) -> t.TypeGuard[list[str]]:
     """Check if an object is a list of integers.
 
     Args:
diff --git a/src/scandeval/utils.py b/src/scandeval/utils.py
index 00b96b8c2..874205ece 100644
--- a/src/scandeval/utils.py
+++ b/src/scandeval/utils.py
@@ -8,67 +8,31 @@
 import random
 import re
 import sys
+import typing as t
 import warnings
-from collections import defaultdict
-from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Type
 
+import litellm
 import numpy as np
 import pkg_resources
 import requests
 import torch
 from datasets.utils import disable_progress_bar
-from huggingface_hub import HfApi
 from requests.exceptions import RequestException
-from transformers import GenerationConfig
+from transformers import PreTrainedTokenizer
 from transformers import logging as tf_logging
 
 from .enums import Framework
-from .exceptions import InvalidModel, NaNValueInModelOutput
-from .languages import DA, NB, NN, NO, SV, get_all_languages
-from .openai_models import OpenAIModel, OpenAITokenizer
-
-if TYPE_CHECKING:
-    from huggingface_hub.hf_api import ModelInfo
-    from transformers import PreTrainedModel
-
-    from .config import Language
-    from .protocols import GenerativeModel, Tokenizer
-    from .types import Predictions
-
-logger = logging.getLogger(__package__)
-
+from .exceptions import NaNValueInModelOutput
 
 if importlib.util.find_spec("ray") is not None:
     import ray
 
-
-# This is used as input to generative models; it cannot be a special token
-DUMMY_FILL_VALUE = 100
-
-
-GENERATIVE_MODEL_TASKS = [
-    "text-generation",
-    "conversational",
-    # "text2text-generation",  # TODO: Add this when we support it
-]
-
-
-GENERATIVE_DATASET_TASKS = [
-    "knowledge",
-    "common-sense-reasoning",
-    "multiple-choice-reading-comprehension",
-]
-
-
-GENERATIVE_DATASET_SUPERTASKS = ["text-to-text", "text-modelling"]
-
-
-SUPERTASKS_USING_LOGPROBS = ["sequence-classification"]
+if t.TYPE_CHECKING:
+    from .types import Predictions
 
 
-METRIC_ATTRIBUTES_TAKING_UP_MEMORY = ["cached_bertscorer"]
+logger = logging.getLogger("scandeval")
 
 
 def create_model_cache_dir(cache_dir: str, model_id: str) -> str:
@@ -91,7 +55,8 @@ def create_model_cache_dir(cache_dir: str, model_id: str) -> str:
 
 def clear_memory():
     """Clears the memory of unused items."""
-    gc.collect()
+    for gc_generation in range(3):
+        gc.collect(generation=gc_generation)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     if torch.backends.mps.is_available():
@@ -182,6 +147,7 @@ def block_terminal_output():
     logging.getLogger("ray._private.worker").setLevel(logging.CRITICAL)
     logging.getLogger("matplotlib.font_manager").setLevel(logging.CRITICAL)
     logging.getLogger("accelerate").setLevel(logging.CRITICAL)
+    logging.getLogger("LiteLLM").setLevel(logging.CRITICAL)
 
     # This suppresses vLLM logging
     os.environ["LOG_LEVEL"] = "CRITICAL"
@@ -198,10 +164,13 @@ def block_terminal_output():
     tf_logging.set_verbosity(logging.CRITICAL)
     logging.getLogger("transformers.trainer").setLevel(logging.CRITICAL)
 
+    # Disable logging from `litellm`
+    litellm.suppress_debug_info = True
+
 
 def get_class_by_name(
     class_name: str | list[str], module_name: str | None = None
-) -> Type | None:
+) -> t.Type | None:
     """Get a class by its name.
 
     Args:
@@ -218,36 +187,30 @@ def get_class_by_name(
     Returns:
         The class. If the class is not found, None is returned.
     """
-    # Ensure that `class_name` is a sequence
     if isinstance(class_name, str):
         class_name = [class_name]
 
-    # Loop over the class names
     error_messages = list()
     for name in class_name:
-        # Get the snake_case and PascalCase version of the class name
         name_snake = name.replace("-", "_")
         name_pascal = kebab_to_pascal(kebab_string=name)
 
-        # Import the module
-        try:
-            if not module_name:
-                module_name = f"scandeval.{name_snake}"
-            module = importlib.import_module(name=module_name)
-        except ModuleNotFoundError as e:
-            error_messages.append(str(e))
-            module_name = None
-            continue
-
-        # Get the class from the module
-        try:
-            class_: Type = getattr(module, name_pascal)
-        except AttributeError:
-            module_name = None
-            continue
-
-        # Return the class
-        return class_
+        module_names = (
+            [module_name]
+            if module_name is not None
+            else [
+                f"scandeval.{name_snake}",
+                f"scandeval.benchmark_datasets.{name_snake}",
+                f"scandeval.benchmark_modules.{name_snake}",
+            ]
+        )
+        for m_name in module_names:
+            try:
+                module = importlib.import_module(name=m_name)
+                class_: t.Type = getattr(module, name_pascal)
+                return class_
+            except (ModuleNotFoundError, AttributeError) as e:
+                error_messages.append(str(e))
 
     if error_messages:
         errors = "\n- " + "\n- ".join(error_messages)
@@ -286,7 +249,7 @@ def internet_connection_available() -> bool:
         return False
 
 
-def get_special_token_metadata(tokenizer: "Tokenizer") -> dict:
+def get_special_token_metadata(tokenizer: "PreTrainedTokenizer") -> dict:
     """Get the special token metadata for a tokenizer.
 
     Args:
@@ -338,203 +301,6 @@ def get_special_token_metadata(tokenizer: "Tokenizer") -> dict:
     )
 
 
-def get_huggingface_model_lists(
-    languages: list["Language"] | None, token: bool | str | None
-) -> dict[str, list[str]]:
-    """Fetches up-to-date model lists from the Hugging Face Hub.
-
-    Args:
-        languages:
-            The language codes of the language to consider. If None then the models
-            will not be filtered on language.
-        token:
-            The authentication token for the Hugging Face Hub. If a boolean value is
-            specified then the token will be fetched from the Hugging Face CLI, where
-            the user has logged in through `huggingface-cli login`. If a string is
-            specified then it will be used as the token.
-
-    Returns:
-        The keys are filterings of the list, which includes all language codes,
-        including 'multilingual', as well as 'all'. The values are lists of model IDs.
-    """
-    # Get list of all languages
-    all_languages = list(get_all_languages().values())
-
-    # If no languages are specified, then include all languages
-    language_list = all_languages if languages is None else languages
-
-    # Form string of languages
-    if len(language_list) == 1:
-        language_string = language_list[0].name
-    else:
-        language_list = sorted(language_list, key=lambda x: x.name)
-        if {lang.code for lang in language_list} == {
-            lang.code for lang in all_languages
-        }:
-            language_string = "all"
-        else:
-            # Remove generic 'Norwegian' from the list of languages if both 'Bokmål'
-            # and 'Nynorsk' already exist in the list
-            if all([lang in language_list for lang in [NO, NB, NN]]):
-                language_list = [lang for lang in language_list if lang != NO]
-
-            language_string = (
-                f"{', '.join(lang.name for lang in language_list[:-1])} and "
-                f"{language_list[-1].name}"
-            )
-
-    # Log fetching message
-    logger.info(f"Fetching list of {language_string} models from the Hugging Face Hub.")
-
-    # Initialise the API
-    api: HfApi = HfApi()
-
-    # Initialise model lists
-    model_lists = defaultdict(list)
-
-    # Do not iterate over all the languages if we are not filtering on language
-    language_itr: list["Language | None"]
-    if {lang.code for lang in language_list} == {lang.code for lang in all_languages}:
-        language_itr = [None]
-    else:
-        language_itr = deepcopy(language_list)  # type: ignore[arg-type]
-
-    for language in language_itr:
-        # Extract the language code
-        language_str: str | None
-        if language is not None:
-            language_str = language.code
-        else:
-            language_str = None
-
-        # Fetch the model list
-        models: list["ModelInfo"] = list(
-            api.list_models(language=language_str, token=token)
-        )
-
-        # Filter the models to only keep the ones with the specified language
-        models = [
-            model
-            for model in models
-            if (language is None or language.code in model.tags)
-        ]
-
-        # Only keep the models which are not finetuned
-        models = [
-            model
-            for model in models
-            if model.pipeline_tag is None
-            or model.pipeline_tag
-            in {"fill-mask", "sentence-similarity", "feature-extraction"}
-            | set(GENERATIVE_MODEL_TASKS)
-        ]
-
-        # Extract the model IDs
-        model_ids: list[str] = [model.modelId for model in models if model.modelId]
-
-        # Remove models that have "finetuned" in their name
-        model_ids = [
-            model_id for model_id in model_ids if "finetuned" not in model_id.lower()
-        ]
-
-        # Store the model IDs
-        model_lists["all"].extend(model_ids)
-        if language is not None:
-            model_lists[language.code].extend(model_ids)
-
-    # Add multilingual models manually
-    multi_models = [
-        "google-bert/bert-base-multilingual-cased",
-        "google-bert/bert-base-multilingual-uncased",
-        "distilbert-base-multilingual-cased",
-        "distilbert/cardiffnlp/twitter-xlm-roberta-base",
-        "microsoft/infoxlm-base",
-        "microsoft/infoxlm-large",
-        "microsoft/xlm-align-base",
-        "microsoft/mdeberta-v3-base",
-        "setu4993/LaBSE",
-        "sentence-transformers/distilbert-multilingual-nli-stsb-quora-ranking",
-        "sentence-transformers/distiluse-base-multilingual-cased",
-        "sentence-transformers/distiluse-base-multilingual-cased-v1",
-        "sentence-transformers/distiluse-base-multilingual-cased-v2",
-        "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-        "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-        "sentence-transformers/paraphrase-xlm-r-multilingual-v1",
-        "sentence-transformers/quora-distilbert-multilingual",
-        "sentence-transformers/stsb-xlm-r-multilingual",
-        "sentence-transformers/use-cmlm-multilingual",
-        "studio-ousia/mluke-base",
-        "studio-ousia/mluke-large",
-        "FacebookAI/xlm-roberta-base",
-        "FacebookAI/xlm-roberta-large",
-        "dbmdz/bert-tiny-historic-multilingual-cased",
-        "dbmdz/bert-mini-historic-multilingual-cased",
-        "dbmdz/bert-base-historic-multilingual-cased",
-        "dbmdz/bert-medium-historic-multilingual-cased",
-    ]
-    model_lists["multilingual"] = multi_models
-    model_lists["all"].extend(multi_models)
-
-    # Add fresh models
-    fresh_models = ["fresh-xlm-roberta-base", "fresh-electra-small"]
-    model_lists["fresh"].extend(fresh_models)
-    model_lists["all"].extend(fresh_models)
-
-    # Add some multilingual Danish models manually that have not marked 'da' as their
-    # language
-    if DA in language_itr:
-        multi_da_models: list[str] = [
-            "Geotrend/bert-base-en-da-cased",
-            "Geotrend/bert-base-25lang-cased",
-            "Geotrend/bert-base-en-fr-de-no-da-cased",
-            "Geotrend/distilbert-base-en-da-cased",
-            "Geotrend/distilbert-base-25lang-cased",
-            "Geotrend/distilbert-base-en-fr-de-no-da-cased",
-        ]
-        model_lists["da"].extend(multi_da_models)
-        model_lists["all"].extend(multi_da_models)
-
-    # Add some multilingual Swedish models manually that have not marked 'sv' as their
-    # language
-    if SV in language_itr:
-        multi_sv_models: list[str] = []
-        model_lists["sv"].extend(multi_sv_models)
-        model_lists["all"].extend(multi_sv_models)
-
-    # Add some multilingual Norwegian models manually that have not marked 'no', 'nb'
-    # or 'nn' as their language
-    if any(lang in language_itr for lang in [NO, NB, NN]):
-        multi_no_models: list[str] = [
-            "Geotrend/bert-base-en-no-cased",
-            "Geotrend/bert-base-25lang-cased",
-            "Geotrend/bert-base-en-fr-de-no-da-cased",
-            "Geotrend/distilbert-base-en-no-cased",
-            "Geotrend/distilbert-base-25lang-cased",
-            "Geotrend/distilbert-base-en-fr-de-no-da-cased",
-        ]
-        model_lists["no"].extend(multi_no_models)
-        model_lists["all"].extend(multi_no_models)
-
-    # Remove duplicates from the lists
-    for lang, model_list in model_lists.items():
-        model_lists[lang] = list(set(model_list))
-
-    # Remove banned models
-    BANNED_MODELS = [
-        r"TransQuest/siamesetransquest-da.*",
-        r"M-CLIP/.*",
-        r".*/.*CTRL.*",  # TEMP
-    ]
-    for lang, model_list in model_lists.items():
-        model_lists[lang] = [
-            model
-            for model in model_list
-            if not any(re.search(regex, model) is not None for regex in BANNED_MODELS)
-        ]
-
-    return dict(model_lists)
-
-
 class HiddenPrints:
     """Context manager which removes all terminal output."""
 
@@ -553,53 +319,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         sys.stderr = self._original_stderr
 
 
-def model_is_generative(model: "PreTrainedModel | GenerativeModel") -> bool:
-    """Check if a model is generative or not.
-
-    Args:
-        model:
-            The model to check.
-
-    Returns:
-        Whether the model is generative or not.
-    """
-    known_generative_models = ["VLLMModel", "OpenAIModel"]
-    if any(model.__class__.__name__ == name for name in known_generative_models):
-        return True
-
-    try:
-        dummy_inputs = torch.tensor(
-            [[DUMMY_FILL_VALUE]], device=model.device, dtype=torch.long
-        )
-        generation_config = GenerationConfig(
-            max_new_tokens=1,
-            pad_token_id=model.config.pad_token_id,
-            eos_token_id=model.config.eos_token_id,
-        )
-        try:
-            model.generate(
-                inputs=dummy_inputs,
-                attention_mask=torch.ones_like(dummy_inputs),
-                generation_config=generation_config,
-            )
-        except ValueError as e:
-            if "attention_mask" not in str(e):
-                raise e
-            model.generate(inputs=dummy_inputs, generation_config=generation_config)
-        return True
-    except (NotImplementedError, TypeError) as e:
-        if "PYTORCH_ENABLE_MPS_FALLBACK" in str(e):
-            raise InvalidModel(
-                "The benchmark failed because the environment variable "
-                "`PYTORCH_ENABLE_MPS_FALLBACK` is not set. Please set this "
-                "environment variable to `1` and try again."
-            )
-        logger.debug(
-            f"The model was found not to be generative, as an error occurred: {e}"
-        )
-        return False
-
-
 def raise_if_model_output_contains_nan_values(model_output: "Predictions") -> None:
     """Raise an exception if the model output contains NaN values.
 
@@ -623,7 +342,7 @@ def raise_if_model_output_contains_nan_values(model_output: "Predictions") -> No
 
 
 def should_prompts_be_stripped(
-    labels_to_be_generated: list[str], tokenizer: "Tokenizer"
+    labels_to_be_generated: list[str], tokenizer: "PreTrainedTokenizer"
 ) -> bool:
     """Determine if we should strip the prompts for few-shot evaluation.
 
@@ -662,7 +381,7 @@ def should_prompts_be_stripped(
 
 
 def should_prefix_space_be_added_to_labels(
-    labels_to_be_generated: list[str], tokenizer: "Tokenizer"
+    labels_to_be_generated: list[str], tokenizer: "PreTrainedTokenizer"
 ) -> bool:
     """Determine if we should add a prefix space to the labels.
 
@@ -678,11 +397,6 @@ def should_prefix_space_be_added_to_labels(
     Returns:
         Whether we should add a prefix space to the labels.
     """
-    # We don't add a prefix space to OpenAI models, since they output strings directly,
-    # and we always strip these for token ID consistency
-    if isinstance(tokenizer, OpenAITokenizer):
-        return False
-
     if not should_prompts_be_stripped(
         labels_to_be_generated=labels_to_be_generated, tokenizer=tokenizer
     ):
@@ -707,7 +421,7 @@ def should_prefix_space_be_added_to_labels(
     return add_prefix_space
 
 
-def get_end_of_chat_token_ids(tokenizer: "Tokenizer") -> list[int] | None:
+def get_end_of_chat_token_ids(tokenizer: "PreTrainedTokenizer") -> list[int] | None:
     """Get the end token ID for chat models.
 
     This is only relevant for tokenizers with a chat template.
@@ -727,7 +441,7 @@ def get_end_of_chat_token_ids(tokenizer: "Tokenizer") -> list[int] | None:
     if tokenizer.chat_template is None:
         return None
 
-    user_message: dict[Literal["role", "content"], str] = dict()
+    user_message: dict[t.Literal["role", "content"], str] = dict()
     user_message["role"] = "user"
     user_message["content"] = "X"
     token_ids = tokenizer.apply_chat_template(conversation=[user_message])
@@ -749,7 +463,7 @@ def get_end_of_chat_token_ids(tokenizer: "Tokenizer") -> list[int] | None:
     return end_of_chat_tokens
 
 
-def convert_prompt_to_instruction(prompt: str, tokenizer: "Tokenizer") -> str:
+def convert_prompt_to_instruction(prompt: str, tokenizer: "PreTrainedTokenizer") -> str:
     """Make an instruction prompt conform to a chat template.
 
     Note that it is expected that the prompt has the following format:
@@ -766,13 +480,9 @@ def convert_prompt_to_instruction(prompt: str, tokenizer: "Tokenizer") -> str:
     if tokenizer.chat_template is None:
         return prompt
 
-    user_message: dict[Literal["role"] | Literal["content"], str] = dict()
-    user_message["role"] = "user"
-    user_message["content"] = prompt
-
     instruction_prompt = tokenizer.apply_chat_template(
-        conversation=[user_message],
-        chat_template=tokenizer.chat_template,
+        conversation=[dict(role="user", content=prompt)],
+        chat_template=tokenizer.get_chat_template(),
         add_generation_prompt=True,
         tokenize=False,
     )
@@ -812,76 +522,3 @@ def unscramble(scrambled_text: str) -> str:
     inverse_permutation = np.argsort(permutation)
     unscrambled = "".join(scrambled_text[i] for i in inverse_permutation)
     return unscrambled
-
-
-def get_model_max_length(
-    model: "PreTrainedModel | GenerativeModel", tokenizer: "Tokenizer | None" = None
-) -> int:
-    """Get the maximum context length of a model.
-
-    Args:
-        model:
-            The model.
-        tokenizer:
-            The tokenizer, or None if the tokenizer is not available.
-
-    Returns:
-        The maximum context length.
-    """
-    all_max_lengths: list[int] = list()
-
-    if tokenizer is not None:
-        # Add the registered max length of the tokenizer
-        if hasattr(tokenizer, "model_max_length") and tokenizer.model_max_length < int(
-            1e30
-        ):
-            all_max_lengths.append(tokenizer.model_max_length)
-
-        # Add the max length derived from the model's input sizes
-        if hasattr(tokenizer, "max_model_input_sizes"):
-            all_max_lengths.extend(
-                [
-                    size
-                    for size in tokenizer.max_model_input_sizes.values()
-                    if size is not None
-                ]
-            )
-
-    # Add max length candidates from the model's configuration
-    candidate_config_max_lengths = [
-        "max_position_embeddings",
-        "model_max_length",
-        "max_sequence_length",
-        "sliding_window",
-        "sliding_window_size",
-    ]
-    for candidate_config_max_length in candidate_config_max_lengths:
-        if (
-            hasattr(model.config, candidate_config_max_length)
-            and (value := getattr(model.config, candidate_config_max_length))
-            is not None
-        ):
-            all_max_lengths.append(value)
-
-    # To avoid models having artificially low max lengths, we remove any max lengths
-    # that are less than 128
-    all_max_lengths = [
-        max_length for max_length in all_max_lengths if max_length >= 128
-    ]
-
-    if len(list(all_max_lengths)) > 0:
-        model_max_length = min(list(all_max_lengths))
-
-        # If the model is an OpenAI chat model then we add on 7 extra tokens, as
-        # these are part of the chat prompt and was removed from the sequence
-        # length
-        if (
-            model_max_length >= 0
-            and isinstance(model, OpenAIModel)
-            and model.is_chat_model
-        ):
-            model_max_length += 7
-    else:
-        model_max_length = -1
-
-    return model_max_length
diff --git a/src/scandeval/vllm_models.py b/src/scandeval/vllm_models.py
deleted file mode 100644
index b14921068..000000000
--- a/src/scandeval/vllm_models.py
+++ /dev/null
@@ -1,407 +0,0 @@
-"""A wrapper for vLLM models."""
-
-import importlib.util
-import logging
-import math
-import sys
-from pathlib import Path
-from types import MethodType
-from typing import TYPE_CHECKING
-
-import torch
-from huggingface_hub import snapshot_download
-from tqdm import tqdm
-from transformers import GenerationConfig
-from transformers.utils import ModelOutput
-
-from .exceptions import NeedsExtraInstalled
-from .protocols import Tokenizer
-from .utils import clear_memory, get_end_of_chat_token_ids
-
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedTokenizerBase
-    from vllm import LLM, RequestOutput
-    from vllm.lora.request import LoRARequest
-
-    from .config import ModelConfig
-
-if importlib.util.find_spec("ray") is not None:
-    import ray
-
-if importlib.util.find_spec("vllm") is not None:
-    from vllm import LLM, SamplingParams
-    from vllm.lora.request import LoRARequest
-
-    try:
-        from vllm.model_executor.parallel_utils.parallel_state import (
-            destroy_model_parallel,
-        )
-    except ImportError:
-        from vllm.distributed.parallel_state import destroy_model_parallel
-
-
-logger = logging.getLogger(__package__)
-
-
-class VLLMModel:
-    """A wrapper for vLLM models."""
-
-    def __init__(
-        self,
-        model_config: "ModelConfig",
-        hf_model_config: "PretrainedConfig",
-        model_cache_dir: "str | Path",
-        trust_remote_code: bool,
-        tokenizer: "PreTrainedTokenizerBase | None" = None,
-        adapter_base_model_id: str | None = None,
-    ) -> None:
-        """Initialize a vLLM model.
-
-        Args:
-            model_config:
-                A model configuration.
-            hf_model_config:
-                A Hugging Face model configuration.
-            model_cache_dir:
-                The directory to cache the model in.
-            trust_remote_code:
-                Whether to trust remote code, e.g., from Hugging Face.
-            tokenizer:
-                A Hugging Face tokenizer. If None, the tokenizer will need to be
-                loaded separately.
-            adapter_base_model_id:
-                The base model ID of the adapter model, if the model is an adapter.
-                None otherwise.
-        """
-        self.model_config = model_config
-        self.config = hf_model_config
-        self.model_cache_dir = model_cache_dir
-        self.trust_remote_code = trust_remote_code
-        self.device = torch.device("cuda")
-        self.tokenizer = tokenizer
-        self.adapter_base_model_id = adapter_base_model_id
-        self.adapter_path: str | None = None
-        self.lora_request: "LoRARequest | None" = None
-
-        self.max_model_len = 5_000
-        potential_max_model_length_config_names = [
-            "max_position_embeddings",
-            "max_sequence_length",
-            "model_max_length",
-            "sliding_window",
-            "sliding_window_size",
-            "n_positions",
-        ]
-        for config_name in potential_max_model_length_config_names:
-            if hasattr(hf_model_config, config_name):
-                model_len = getattr(hf_model_config, config_name)
-                if model_len is not None:
-                    self.max_model_len = min(self.max_model_len, model_len)
-
-        quantization = None
-        if hasattr(self.config, "quantization_config"):
-            quantization = self.config.quantization_config.get("quant_method", None)
-
-        # The quantised models require extra dependencies
-        if quantization == "gptq" and (
-            importlib.util.find_spec("auto_gptq") is None
-            or importlib.util.find_spec("optimum") is None
-        ):
-            raise NeedsExtraInstalled(extra="quantization")
-        if quantization == "awq" and importlib.util.find_spec("awq") is None:
-            raise NeedsExtraInstalled(extra="quantization")
-
-        dtype: str | torch.dtype = "auto"
-        if quantization is not None and self.config.torch_dtype != torch.float16:
-            logger.info(
-                "You are loading a quantized model with dtype "
-                f"{self.config.torch_dtype}, which vLLM does not support. Setting "
-                "dtype to float16 instead."
-            )
-            dtype = torch.float16
-
-        if self.adapter_base_model_id is not None:
-            download_dir = str(Path(self.model_cache_dir) / "base_model")
-        else:
-            download_dir = str(self.model_cache_dir)
-
-        vllm_kwargs = dict(
-            model=self.adapter_base_model_id or self.model_config.model_id,
-            tokenizer=self.adapter_base_model_id or self.model_config.model_id,
-            gpu_memory_utilization=0.95,
-            max_model_len=self.max_model_len,
-            download_dir=download_dir,
-            trust_remote_code=self.trust_remote_code,
-            revision=self.model_config.revision,
-            seed=4242,
-            distributed_executor_backend="ray",
-            tensor_parallel_size=torch.cuda.device_count(),
-            disable_custom_all_reduce=True,
-            quantization=quantization,
-            dtype=dtype,
-            enforce_eager=True,
-            max_logprobs=10,
-            # TEMP: Prefix caching isn't supported with sliding window in vLLM yet, so
-            # we disable it for now
-            enable_prefix_caching=False,
-            enable_lora=self.adapter_base_model_id is not None,
-            max_lora_rank=256,
-        )
-
-        self._model = self._initialise(vllm_kwargs=vllm_kwargs)
-
-    def _initialise(self, vllm_kwargs: dict) -> "LLM":
-        """Initialise the vLLM model.
-
-        Args:
-            vllm_kwargs:
-                The keyword arguments to pass to the vLLM
-
-        Returns:
-            The initialised vLLM model.
-        """
-        clear_vllm()
-        model = LLM(**vllm_kwargs)
-        model._run_engine = MethodType(_run_engine_with_fixed_progress_bars, model)
-
-        if self.adapter_base_model_id is not None:
-            self.adapter_path = snapshot_download(
-                repo_id=self.model_config.model_id, cache_dir=Path(self.model_cache_dir)
-            )
-            self.lora_request = LoRARequest(
-                lora_name="adapter", lora_int_id=1, lora_path=self.adapter_path
-            )
-
-        return model
-
-    def __del__(self) -> None:
-        """Clear the GPU memory used by the model, and remove the model itself."""
-        if hasattr(self, "_model"):
-            del self._model
-        del self
-        clear_vllm()
-
-    def generate(
-        self,
-        inputs: torch.Tensor,
-        generation_config: GenerationConfig | None = None,
-        **generation_kwargs,
-    ) -> torch.Tensor | torch.LongTensor | ModelOutput:
-        """Generate sequences using the model.
-
-        Args:
-            inputs:
-                The input batch of sequences to generate from.
-            generation_config:
-                The generation config to use for generation.
-            **generation_kwargs:
-                Additional generation kwargs to pass to the model.
-
-        Returns:
-            The generated sequences.
-        """
-        if self.tokenizer is None:
-            raise ValueError("Tokenizer must be loaded to generate sequences.")
-
-        if generation_config is None:
-            generation_config = GenerationConfig(**generation_kwargs)
-        else:
-            for key, value in generation_kwargs.items():
-                setattr(generation_config, key, value)
-
-        # Define which tokens to use as stopping criteria. We want to use the padding
-        # token, end-of-sentence token, and a double newline (since these separate the
-        # few-shot examples in the input)
-        stop_tokens: list[str] = ["\n\n"]
-        if self.tokenizer.pad_token_id is not None:
-            stop_tokens.append(self.tokenizer.pad_token)
-        if self.tokenizer.eos_token_id is not None:
-            stop_tokens.append(self.tokenizer.eos_token)
-            if self.tokenizer.pad_token_id is None:
-                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-        if (
-            self.tokenizer.bos_token_id is not None
-            and self.tokenizer.pad_token_id is None
-        ):
-            self.tokenizer.pad_token_id = self.tokenizer.bos_token_id
-            self.tokenizer.pad_token = self.tokenizer.bos_token
-        assert self.tokenizer.pad_token_id is not None
-
-        # Add end of chat token as a stopping token, if it exists
-        end_of_chat_token_ids = get_end_of_chat_token_ids(tokenizer=self.tokenizer)
-        if end_of_chat_token_ids is not None:
-            end_of_chat_token = self.tokenizer.decode(end_of_chat_token_ids).strip()
-            if end_of_chat_token:
-                stop_tokens.append(end_of_chat_token)
-
-        # Define the parameters used for vLLM generation
-        max_tokens: int = generation_config.max_new_tokens or 1
-        temperature = (
-            0.0 if not generation_config.do_sample else generation_config.temperature
-        )
-        sampling_params = SamplingParams(
-            # What to output
-            max_tokens=max_tokens,
-            logprobs=10 if generation_config.output_scores else None,
-            n=generation_config.num_return_sequences,
-            # How to sample
-            temperature=temperature,
-            top_p=generation_config.top_p,
-            top_k=generation_config.top_k,
-            stop=[stop_token for stop_token in stop_tokens if stop_token],
-            repetition_penalty=generation_config.repetition_penalty,
-            frequency_penalty=generation_config.repetition_penalty - 1.0,
-            logits_processors=generation_kwargs.get("logits_processors"),
-        )
-
-        # The inputs are tokenised, so we decode them to get the original text, which
-        # is the input to the vLLM model
-        prompts = self.tokenizer.batch_decode(
-            sequences=inputs, skip_special_tokens=True
-        )
-
-        # If any of the prompts are empty then we need to replace them with a BOS token
-        # so that the vLLM model can generate from them
-        if any(len(prompt) == 0 for prompt in prompts):
-            logger.debug("Found empty prompts, replacing with BOS token.")
-            prompts = [
-                prompt if len(prompt) > 0 else self.tokenizer.bos_token
-                for prompt in prompts
-            ]
-
-        # Generate sequences using vLLM
-        input_is_a_test = len(prompts) == 1 and len(set(prompts[0])) == 1
-        raw_outputs = self._model.generate(
-            prompts=prompts,
-            sampling_params=sampling_params,
-            use_tqdm=(not input_is_a_test),
-            lora_request=self.lora_request,
-        )
-
-        # Collect the generated sequences into a single tensor of shape
-        # (batch_size, generated_sequence_length)
-        output = torch.nn.utils.rnn.pad_sequence(
-            sequences=[
-                torch.LongTensor(output.outputs[0].token_ids) for output in raw_outputs
-            ],
-            batch_first=True,
-            padding_value=float(self.tokenizer.pad_token_id),
-        )
-
-        if generation_config.return_dict_in_generate:
-            # Add logprobs scores to the output
-            if generation_config.output_scores:
-                # Create a list with placeholder logprobs for every token generated.
-                # Each tensor in the list will be of shape (batch_size, vocab_size)
-                batch_size = len(raw_outputs)
-                vocab_size = len(self.tokenizer.get_vocab())
-                max_seq_len = max(
-                    len(raw_output.outputs[0].logprobs) for raw_output in raw_outputs
-                )
-                scores = [
-                    torch.full(size=(batch_size, vocab_size), fill_value=-math.inf)
-                    for _ in range(max_seq_len)
-                ]
-
-                # Fill in the logprobs for each generated token. The logprobs from the
-                # vLLM output only contain the logprobs for the top-k tokens, so we
-                # only fill in these and leave the rest at ~0% probability
-                for sample_idx, raw_output in enumerate(raw_outputs):
-                    assert raw_output.outputs[0].logprobs is not None
-                    seq_len = len(raw_output.outputs[0].logprobs)
-                    for gen_token_idx in range(seq_len):
-                        logprobs_dict = raw_output.outputs[0].logprobs[gen_token_idx]
-                        for token_idx, logprob_obj in logprobs_dict.items():
-                            logprob = logprob_obj.logprob
-                            scores[gen_token_idx][sample_idx, token_idx] = logprob
-
-                output = ModelOutput(dict(sequences=output, scores=tuple(scores)))
-            else:
-                output = ModelOutput(dict(sequences=output))
-
-        return output
-
-    def __call__(
-        self,
-        inputs: torch.Tensor,
-        generation_config: GenerationConfig | None = None,
-        **generation_kwargs,
-    ) -> torch.Tensor | torch.LongTensor | ModelOutput:
-        """Generate sequences using the model.
-
-        Args:
-            inputs:
-                The input batch of sequences to generate from.
-            generation_config:
-                The generation config to use for generation.
-            **generation_kwargs:
-                Additional generation kwargs to pass to the model.
-
-        Returns:
-            The generated sequences.
-        """
-        return self.generate(
-            inputs=inputs, generation_config=generation_config, **generation_kwargs
-        )
-
-    def set_tokenizer(self, tokenizer: "Tokenizer") -> None:
-        """Set the tokenizer to use for generation.
-
-        Args:
-            tokenizer:
-                The tokenizer to use for generation.
-        """
-        self.tokenizer = tokenizer
-        self._model.set_tokenizer(tokenizer)
-
-    def to(self, _: torch.device) -> None:
-        """Dummy method to make the model compatible with the benchmarking script."""
-        pass
-
-    def eval(self) -> None:
-        """Dummy method to make the model compatible with the benchmarking script."""
-        pass
-
-    def children(self) -> list:
-        """Dummy method to make the model compatible with the benchmarking script."""
-        return []
-
-
-def _run_engine_with_fixed_progress_bars(
-    self: "LLM", use_tqdm: bool
-) -> list["RequestOutput"]:
-    if use_tqdm:
-        num_requests = self.llm_engine.get_num_unfinished_requests()
-        pbar = tqdm(
-            total=num_requests, leave=False, disable=hasattr(sys, "_called_from_test")
-        )
-
-    # Run the engine.
-    outputs: list["RequestOutput"] = list()
-    while self.llm_engine.has_unfinished_requests():
-        step_outputs = self.llm_engine.step()
-        for output in step_outputs:
-            if output.finished:
-                outputs.append(output)
-                if use_tqdm:
-                    pbar.update(1)
-
-    if use_tqdm:
-        pbar.close()
-
-    # Sort the outputs by request ID. This is necessary because some requests may be
-    # finished earlier than its previous requests.
-    outputs = sorted(outputs, key=lambda x: int(x.request_id))
-
-    return outputs
-
-
-def clear_vllm() -> None:
-    """Clear the GPU memory used by the vLLM model, enabling re-initialisation."""
-    try:
-        destroy_model_parallel()
-    except ImportError:
-        pass
-    clear_memory()
-    ray.shutdown()
diff --git a/src/scripts/create_norglm_multiqa.py b/src/scripts/create_norglm_multiqa.py
index e96d67ca0..a757ea416 100644
--- a/src/scripts/create_norglm_multiqa.py
+++ b/src/scripts/create_norglm_multiqa.py
@@ -61,7 +61,8 @@ def qa_to_list(row: str) -> list:
 
     df["question_answer"] = df["question_answer"].apply(qa_to_list)
 
-    # split question_answer list [(question, answer) ...] column to question and answer with same context
+    # split question_answer list [(question, answer) ...] column to question and answer
+    # with same context
     df = df.explode("question_answer")
     df[["question", "answer"]] = pd.DataFrame(
         df.question_answer.tolist(), index=df.index
@@ -93,7 +94,8 @@ def rephrase_answer(question: str, answer: str, context: str) -> str:
                 The context.
 
         Returns:
-            The rephrased answer (if the answer is already in the context, it is returned as is).
+            The rephrased answer (if the answer is already in the context, it is
+            returned as is).
         """
         answer = answer[:-1] if answer.endswith(".") else answer
 
@@ -106,11 +108,14 @@ def rephrase_answer(question: str, answer: str, context: str) -> str:
                 {
                     "role": "user",
                     "content": (
-                        f"Given the following context: '{context}', and the '{question}' please rephrase the answer, so that it matches exactly "
-                        f"with a phrase from the context in a case-insensitive way. E.g. it must be possible to find the rephrased answer "
-                        f"in the context without any modifications. "
-                        f"The rephrased answer should be as concise as possible, preferable not more than 7 words, and ideally 3 or less words."
-                        f"Now rephrase this answer: '{answer}'"
+                        f"Given the following context: '{context}', and the "
+                        f"'{question}' please rephrase the answer, so that it matches "
+                        "exactly with a phrase from the context in a case-insensitive "
+                        "way. E.g. it must be possible to find the rephrased answer "
+                        "in the context without any modifications. The rephrased "
+                        "answer should be as concise as possible, preferable not "
+                        "more than 7 words, and ideally 3 or less words. Now "
+                        f"rephrase this answer: '{answer}'"
                     ),
                 }
             ],
@@ -120,7 +125,7 @@ def rephrase_answer(question: str, answer: str, context: str) -> str:
         )
 
         rephrased_answer = chat_completion.choices[0].message.content
-        return rephrased_answer
+        return rephrased_answer or answer
 
     # Use the original answer if it could be found in the context
     df_orig_bool = df[["answer", "context"]].apply(
diff --git a/uv.lock b/uv.lock
index 55ead8ee1..dc9ea390c 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1113,7 +1113,7 @@ name = "importlib-metadata"
 version = "8.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "zipp", marker = "platform_system != 'Darwin'" },
+    { name = "zipp" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/cd/12/33e59336dca5be0c398a7482335911a33aa0e20776128f038019f1a95f1b/importlib_metadata-8.5.0.tar.gz", hash = "sha256:71522656f0abace1d072b9e5481a48f07c138e00f079c38c8f883823f9c26bd7", size = 55304 }
 wheels = [
@@ -1479,6 +1479,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f3/98/c119974fdce4808afdf3622230759c871bc4c73287cf34b338db2be936b8/levenshtein-0.26.1-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:9e70d7ee157a9b698c73014f6e2b160830e7d2d64d2e342fefc3079af3c356fc", size = 95854 },
 ]
 
+[[package]]
+name = "litellm"
+version = "1.53.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "aiohttp" },
+    { name = "click" },
+    { name = "importlib-metadata" },
+    { name = "jinja2" },
+    { name = "jsonschema" },
+    { name = "openai" },
+    { name = "pydantic" },
+    { name = "python-dotenv" },
+    { name = "requests" },
+    { name = "tiktoken" },
+    { name = "tokenizers" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/70/fc/fb491dc60307cd1a8826050fbdeeef84e560000a3a2593430887f0b44598/litellm-1.53.1.tar.gz", hash = "sha256:20bb2f98a4613c4c7b86e9186df84c27d7166eca65a7b58584c1e061a29f7519", size = 6170999 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d5/a6/8a3b430d89b9935ef6d1e88e565f631038d3124788e6effa3c0c9ad4756e/litellm-1.53.1-py3-none-any.whl", hash = "sha256:7a0b480c0bd4f9236910a426d8d8f46a517af4daae2a61d57e42bff4b4daf9b0", size = 6418676 },
+]
+
 [[package]]
 name = "llvmlite"
 version = "0.43.0"
@@ -1924,6 +1946,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/14/87/30323ad2e52f56262019a4493fe5f5e71067c5561ce7e2f9c75de520f5e8/ml_dtypes-0.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:cb5cc7b25acabd384f75bbd78892d0c724943f3e2e1986254665a1aa10982e07", size = 213195 },
 ]
 
+[[package]]
+name = "more-itertools"
+version = "10.5.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/51/78/65922308c4248e0eb08ebcbe67c95d48615cc6f27854b6f2e57143e9178f/more-itertools-10.5.0.tar.gz", hash = "sha256:5482bfef7849c25dc3c6dd53a6173ae4795da2a41a80faea6700d9f5846c5da6", size = 121020 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/48/7e/3a64597054a70f7c86eb0a7d4fc315b8c1ab932f64883a297bdffeb5f967/more_itertools-10.5.0-py3-none-any.whl", hash = "sha256:037b0d3203ce90cca8ab1defbbdac29d5f993fc20131f3664dc8d6acfa872aef", size = 60952 },
+]
+
 [[package]]
 name = "mpmath"
 version = "1.3.0"
@@ -3726,19 +3757,25 @@ version = "13.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
+    { name = "bert-score" },
     { name = "click" },
     { name = "datasets" },
     { name = "evaluate" },
     { name = "huggingface-hub" },
+    { name = "levenshtein" },
+    { name = "litellm" },
+    { name = "more-itertools" },
     { name = "numpy" },
     { name = "pandas" },
     { name = "protobuf" },
     { name = "pydantic" },
     { name = "pyinfer" },
     { name = "python-dotenv" },
+    { name = "rouge-score" },
     { name = "sacremoses" },
     { name = "sentencepiece" },
     { name = "seqeval" },
+    { name = "tenacity" },
     { name = "termcolor" },
     { name = "torch" },
     { name = "transformers" },
@@ -3746,57 +3783,37 @@ dependencies = [
 
 [package.optional-dependencies]
 all = [
-    { name = "bert-score" },
     { name = "bitsandbytes", marker = "platform_system == 'Linux'" },
     { name = "demjson3" },
     { name = "flax" },
     { name = "gradio" },
     { name = "jax" },
     { name = "jaxlib" },
-    { name = "levenshtein" },
-    { name = "openai" },
     { name = "outlines" },
-    { name = "rouge-score" },
-    { name = "tiktoken" },
     { name = "vllm", marker = "platform_system == 'Linux'" },
 ]
 cpu-all = [
-    { name = "bert-score" },
     { name = "demjson3" },
     { name = "flax" },
     { name = "gradio" },
     { name = "jax" },
     { name = "jaxlib" },
-    { name = "levenshtein" },
-    { name = "openai" },
     { name = "outlines" },
-    { name = "rouge-score" },
-    { name = "tiktoken" },
 ]
 generative = [
-    { name = "bert-score" },
     { name = "bitsandbytes", marker = "platform_system == 'Linux'" },
     { name = "demjson3" },
     { name = "outlines" },
-    { name = "rouge-score" },
     { name = "vllm", marker = "platform_system == 'Linux'" },
 ]
 human-evaluation = [
     { name = "gradio" },
-    { name = "levenshtein" },
 ]
 jax = [
     { name = "flax" },
     { name = "jax" },
     { name = "jaxlib" },
 ]
-openai = [
-    { name = "bert-score" },
-    { name = "levenshtein" },
-    { name = "openai" },
-    { name = "rouge-score" },
-    { name = "tiktoken" },
-]
 
 [package.dev-dependencies]
 dev = [
@@ -3822,10 +3839,7 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "accelerate", specifier = ">=0.34.2" },
-    { name = "bert-score", marker = "extra == 'all'", specifier = ">=0.3.13" },
-    { name = "bert-score", marker = "extra == 'cpu-all'", specifier = ">=0.3.13" },
-    { name = "bert-score", marker = "extra == 'generative'", specifier = ">=0.3.13" },
-    { name = "bert-score", marker = "extra == 'openai'", specifier = ">=0.3.13" },
+    { name = "bert-score", specifier = ">=0.3.13" },
     { name = "bitsandbytes", marker = "platform_system == 'Linux' and extra == 'all'", specifier = ">=0.43.1" },
     { name = "bitsandbytes", marker = "platform_system == 'Linux' and extra == 'generative'", specifier = ">=0.43.1" },
     { name = "click", specifier = ">=8.1.3" },
@@ -3847,14 +3861,10 @@ requires-dist = [
     { name = "jaxlib", marker = "extra == 'all'", specifier = ">=0.4.24" },
     { name = "jaxlib", marker = "extra == 'cpu-all'", specifier = ">=0.4.24" },
     { name = "jaxlib", marker = "extra == 'jax'", specifier = ">=0.4.24" },
-    { name = "levenshtein", marker = "extra == 'all'", specifier = ">=0.24.0" },
-    { name = "levenshtein", marker = "extra == 'cpu-all'", specifier = ">=0.24.0" },
-    { name = "levenshtein", marker = "extra == 'human-evaluation'", specifier = ">=0.24.0" },
-    { name = "levenshtein", marker = "extra == 'openai'", specifier = ">=0.24.0" },
+    { name = "levenshtein", specifier = ">=0.24.0" },
+    { name = "litellm", specifier = ">=1.53.1" },
+    { name = "more-itertools", specifier = ">=10.5.0" },
     { name = "numpy", specifier = ">=1.23.0,<2.0.0" },
-    { name = "openai", marker = "extra == 'all'", specifier = ">=1.51.2" },
-    { name = "openai", marker = "extra == 'cpu-all'", specifier = ">=1.51.2" },
-    { name = "openai", marker = "extra == 'openai'", specifier = ">=1.51.2" },
     { name = "outlines", marker = "extra == 'all'", specifier = ">=0.0.44,<0.1" },
     { name = "outlines", marker = "extra == 'cpu-all'", specifier = ">=0.0.44,<0.1" },
     { name = "outlines", marker = "extra == 'generative'", specifier = ">=0.0.44,<0.1" },
@@ -3863,17 +3873,12 @@ requires-dist = [
     { name = "pydantic", specifier = ">=2.6.0" },
     { name = "pyinfer", specifier = ">=0.0.3" },
     { name = "python-dotenv", specifier = ">=1.0.1" },
-    { name = "rouge-score", marker = "extra == 'all'", specifier = ">=0.1.2" },
-    { name = "rouge-score", marker = "extra == 'cpu-all'", specifier = ">=0.1.2" },
-    { name = "rouge-score", marker = "extra == 'generative'", specifier = ">=0.1.2" },
-    { name = "rouge-score", marker = "extra == 'openai'", specifier = ">=0.1.2" },
+    { name = "rouge-score", specifier = ">=0.1.2" },
     { name = "sacremoses", specifier = ">=0.1.1" },
     { name = "sentencepiece", specifier = ">=0.1.96" },
     { name = "seqeval", specifier = ">=1.2.2" },
+    { name = "tenacity", specifier = ">=9.0.0" },
     { name = "termcolor", specifier = ">=2.0.0" },
-    { name = "tiktoken", marker = "extra == 'all'", specifier = ">=0.7.0" },
-    { name = "tiktoken", marker = "extra == 'cpu-all'", specifier = ">=0.7.0" },
-    { name = "tiktoken", marker = "extra == 'openai'", specifier = ">=0.7.0" },
     { name = "torch", specifier = ">=2.3.0" },
     { name = "transformers", specifier = ">=4.45.0" },
     { name = "vllm", marker = "platform_system == 'Linux' and extra == 'all'", specifier = ">=0.6.3" },
@@ -4159,6 +4164,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252 },
 ]
 
+[[package]]
+name = "tenacity"
+version = "9.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/94/91fccdb4b8110642462e653d5dcb27e7b674742ad68efd146367da7bdb10/tenacity-9.0.0.tar.gz", hash = "sha256:807f37ca97d62aa361264d497b0e31e92b8027044942bfa756160d908320d73b", size = 47421 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b6/cb/b86984bed139586d01532a587464b5805f12e397594f19f931c4c2fbfa61/tenacity-9.0.0-py3-none-any.whl", hash = "sha256:93de0c98785b27fcf659856aa9f54bfbd399e29969b0621bc7f762bd441b4539", size = 28169 },
+]
+
 [[package]]
 name = "tensorstore"
 version = "0.1.69"


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 8cedda593..243cc9cbb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -7,24 +7,11 @@
 import pytest
 import torch
 
-from scandeval.config import (
-    BenchmarkConfig,
-    DatasetConfig,
-    Language,
-    MetricConfig,
-    ModelConfig,
-    Task,
-)
-from scandeval.dataset_configs import (
-    ANGRY_TWEETS_CONFIG,
-    MMLU_CONFIG,
-    get_all_dataset_configs,
-)
+from scandeval.data_models import BenchmarkConfig, MetricConfig, ModelConfig, Task
+from scandeval.dataset_configs import SPEED_CONFIG, get_all_dataset_configs
 from scandeval.enums import Framework, ModelType
-from scandeval.model_config import get_model_config
-from scandeval.model_loading import load_model
-from scandeval.protocols import GenerativeModel, Tokenizer
-from scandeval.tasks import SPEED
+from scandeval.languages import DA, get_all_languages
+from scandeval.tasks import SENT, SPEED, get_all_tasks
 
 
 def pytest_configure() -> None:
@@ -37,11 +24,22 @@ def pytest_unconfigure() -> None:
     delattr(sys, "_called_from_test")
 
 
+ACTIVE_LANGUAGES = {
+    language_code: language
+    for language_code, language in get_all_languages().items()
+    if any(
+        language in cfg.languages
+        for cfg in get_all_dataset_configs().values()
+        if cfg != SPEED_CONFIG
+    )
+}
+
+
 @pytest.fixture(scope="session")
 def auth() -> Generator[str | bool, None, None]:
     """Yields the authentication token to the Hugging Face Hub."""
     # Get the authentication token to the Hugging Face Hub
-    auth = os.environ.get("HUGGINGFACE_HUB_TOKEN", True)
+    auth = os.environ.get("HUGGINGFACE_API_KEY", True)
 
     # Ensure that the token does not contain quotes or whitespace
     if isinstance(auth, str):
@@ -63,25 +61,18 @@ def device() -> Generator[torch.device, None, None]:
 
 
 @pytest.fixture(scope="session")
-def benchmark_config(
-    language, task, auth, device
-) -> Generator[BenchmarkConfig, None, None]:
+def benchmark_config(auth, device) -> Generator[BenchmarkConfig, None, None]:
     """Yields a benchmark configuration used in tests."""
     yield BenchmarkConfig(
-        model_languages=[language],
-        dataset_languages=[language],
-        tasks=[task],
+        model_languages=[DA],
+        dataset_languages=[DA],
+        tasks=[SENT],
         datasets=list(get_all_dataset_configs().keys()),
         framework=None,
-        batch_size=32,
+        batch_size=1,
         raise_errors=False,
         cache_dir=".scandeval_cache",
-        evaluate_train=False,
-        token=auth,
-        openai_api_key=None,
-        azure_openai_api_key=None,
-        azure_openai_endpoint=None,
-        azure_openai_api_version=None,
+        api_key=auth,
         force=False,
         progress_bar=False,
         save_results=True,
@@ -91,9 +82,9 @@ def benchmark_config(
         load_in_4bit=None,
         use_flash_attention=False,
         clear_model_cache=False,
-        only_validation_split=False,
+        evaluate_test_split=False,
         few_shot=True,
-        num_iterations=10,
+        num_iterations=1,
         debug=False,
         run_with_cli=True,
     )
@@ -110,20 +101,28 @@ def metric_config() -> Generator[MetricConfig, None, None]:
     )
 
 
-@pytest.fixture(scope="session")
-def task() -> Generator[Task, None, None]:
+@pytest.fixture(
+    scope="session",
+    params=[task for task in get_all_tasks().values() if task != SPEED],
+    ids=[name for name, task in get_all_tasks().items() if task != SPEED],
+)
+def task(request) -> Generator[Task, None, None]:
     """Yields a dataset task used in tests."""
-    yield SPEED
+    yield request.param
 
 
-@pytest.fixture(scope="session")
-def language():
+@pytest.fixture(
+    scope="session",
+    params=list(ACTIVE_LANGUAGES.values()),
+    ids=list(ACTIVE_LANGUAGES.keys()),
+)
+def language(request):
     """Yields a language used in tests."""
-    yield Language(code="language_code", name="Language name")
+    yield request.param
 
 
 @pytest.fixture(scope="session")
-def model_id() -> Generator[str, None, None]:
+def encoder_model_id() -> Generator[str, None, None]:
     """Yields a model ID used in tests."""
     yield "jonfd/electra-small-nordic"
 
@@ -135,23 +134,15 @@ def generative_model_id() -> Generator[str, None, None]:
 
 
 @pytest.fixture(scope="session")
-def generative_model_and_tokenizer(
-    benchmark_config, generative_model_id
-) -> Generator[tuple[GenerativeModel | None, Tokenizer | None], None, None]:
-    """Yields a generative model ID used in tests."""
-    if os.getenv("USE_VLLM", "0") == "1":
-        model_config = get_model_config(
-            model_id=generative_model_id, benchmark_config=benchmark_config
-        )
-        model, tokenizer = load_model(
-            model_config=model_config,
-            dataset_config=ANGRY_TWEETS_CONFIG,
-            benchmark_config=benchmark_config,
-        )
-    else:
-        model = None
-        tokenizer = None
-    yield model, tokenizer
+def openai_model_id() -> Generator[str, None, None]:
+    """Yields an OpenAI model ID used in tests."""
+    yield "gpt-4o-mini"
+
+
+@pytest.fixture(scope="session")
+def anthropic_model_id() -> Generator[str, None, None]:
+    """Yields an Anthropic model ID used in tests."""
+    yield "claude-3-5-haiku-20241022"
 
 
 @pytest.fixture(scope="session")
@@ -167,51 +158,15 @@ def gptq_generative_model_id() -> Generator[str, None, None]:
 
 
 @pytest.fixture(scope="session")
-def dataset_config(language, task) -> Generator[DatasetConfig, None, None]:
-    """Yields a dataset configuration used in tests."""
-    yield DatasetConfig(
-        name="dataset_name",
-        pretty_name="Dataset name",
-        huggingface_id="dataset_id",
-        task=task,
-        languages=[language],
-        prompt_prefix="",
-        prompt_template="{text}\n{label}",
-        instruction_prompt="",
-        num_few_shot_examples=0,
-        max_generated_tokens=1,
-    )
-
-
-@pytest.fixture(scope="session")
-def model_config(language) -> Generator[ModelConfig, None, None]:
+def model_config() -> Generator[ModelConfig, None, None]:
     """Yields a model configuration used in tests."""
     yield ModelConfig(
         model_id="model_id",
         revision="revision",
         framework=Framework.PYTORCH,
         task="task",
-        languages=[language],
+        languages=[DA],
         model_type=ModelType.FRESH,
         model_cache_dir="cache_dir",
         adapter_base_model_id=None,
     )
-
-
-@pytest.fixture(scope="session")
-def generative_dataset_config() -> Generator[DatasetConfig, None, None]:
-    """Yields a generative dataset configuration used in tests."""
-    yield MMLU_CONFIG
-
-
-@pytest.fixture(scope="session")
-def all_dataset_configs() -> Generator[list[DatasetConfig], None, None]:
-    """Yields all dataset configurations used in tests."""
-    if os.getenv("TEST_ALL_DATASETS", "0") == "1":
-        yield list(get_all_dataset_configs().values())
-    else:
-        yield [
-            dataset_config
-            for dataset_config in get_all_dataset_configs().values()
-            if not dataset_config.unofficial
-        ]
diff --git a/tests/test_benchmark_config_factory.py b/tests/test_benchmark_config_factory.py
index 5d7068b9f..34deb453b 100644
--- a/tests/test_benchmark_config_factory.py
+++ b/tests/test_benchmark_config_factory.py
@@ -11,6 +11,7 @@
     prepare_languages,
     prepare_tasks_and_datasets,
 )
+from scandeval.dataset_configs import get_all_dataset_configs
 from scandeval.enums import Device
 from scandeval.exceptions import InvalidBenchmark
 from scandeval.languages import DA, EN, NB, NN, NO, get_all_languages
@@ -18,18 +19,18 @@
 
 
 @pytest.fixture(scope="module")
-def all_official_dataset_names(all_dataset_configs) -> Generator[list[str], None, None]:
+def all_official_dataset_names() -> Generator[list[str], None, None]:
     """Fixture for all linguistic acceptability dataset configurations."""
-    yield [cfg.name for cfg in all_dataset_configs if not cfg.unofficial]
+    yield [cfg.name for cfg in get_all_dataset_configs().values() if not cfg.unofficial]
 
 
 @pytest.fixture(scope="module")
-def all_official_la_dataset_names(
-    all_dataset_configs,
-) -> Generator[list[str], None, None]:
+def all_official_la_dataset_names() -> Generator[list[str], None, None]:
     """Fixture for all linguistic acceptability dataset configurations."""
     yield [
-        cfg.name for cfg in all_dataset_configs if LA == cfg.task and not cfg.unofficial
+        cfg.name
+        for cfg in get_all_dataset_configs().values()
+        if LA == cfg.task and not cfg.unofficial
     ]
 
 
diff --git a/tests/test_benchmark_dataset.py b/tests/test_benchmark_dataset.py
deleted file mode 100644
index a6d900fc8..000000000
--- a/tests/test_benchmark_dataset.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `benchmark_dataset` module."""
diff --git a/tests/test_benchmark_modules/__init__.py b/tests/test_benchmark_modules/__init__.py
new file mode 100644
index 000000000..1e54c565d
--- /dev/null
+++ b/tests/test_benchmark_modules/__init__.py
@@ -0,0 +1,1 @@
+"""Tests for the `benchmark_modules` subpackage."""
diff --git a/tests/test_benchmark_modules/test_base.py b/tests/test_benchmark_modules/test_base.py
new file mode 100644
index 000000000..6f5bb7545
--- /dev/null
+++ b/tests/test_benchmark_modules/test_base.py
@@ -0,0 +1,1 @@
+"""Unit tests for the `base` module."""
diff --git a/tests/test_benchmark_modules/test_fresh.py b/tests/test_benchmark_modules/test_fresh.py
new file mode 100644
index 000000000..76837cb6b
--- /dev/null
+++ b/tests/test_benchmark_modules/test_fresh.py
@@ -0,0 +1,1 @@
+"""Unit tests for the `fresh` module."""
diff --git a/tests/test_benchmark_modules/test_hf.py b/tests/test_benchmark_modules/test_hf.py
new file mode 100644
index 000000000..328e9efb3
--- /dev/null
+++ b/tests/test_benchmark_modules/test_hf.py
@@ -0,0 +1,35 @@
+"""Unit tests for the `hf` module."""
+
+import pytest
+import torch
+
+from scandeval.benchmark_modules.hf import get_torch_dtype
+
+
+@pytest.mark.parametrize(
+    argnames=["test_device", "torch_dtype_is_set", "bf16_available", "expected"],
+    argvalues=[
+        ("cpu", True, True, torch.float32),
+        ("cpu", True, False, torch.float32),
+        ("cpu", False, True, torch.float32),
+        ("cpu", False, False, torch.float32),
+        ("mps", True, True, torch.float32),
+        ("mps", True, False, torch.float32),
+        ("mps", False, True, torch.float32),
+        ("mps", False, False, torch.float32),
+        ("cuda", True, True, "auto"),
+        ("cuda", True, False, "auto"),
+        ("cuda", False, True, torch.bfloat16),
+        ("cuda", False, False, torch.float16),
+    ],
+)
+def test_get_torch_dtype(test_device, torch_dtype_is_set, bf16_available, expected):
+    """Test that the torch dtype is set correctly."""
+    assert (
+        get_torch_dtype(
+            device=torch.device(test_device),
+            torch_dtype_is_set=torch_dtype_is_set,
+            bf16_available=bf16_available,
+        )
+        == expected
+    )
diff --git a/tests/test_benchmark_modules/test_litellm.py b/tests/test_benchmark_modules/test_litellm.py
new file mode 100644
index 000000000..8a4f52428
--- /dev/null
+++ b/tests/test_benchmark_modules/test_litellm.py
@@ -0,0 +1,1 @@
+"""Unit tests for the `litellm` module."""
diff --git a/tests/test_benchmark_modules/test_vllm.py b/tests/test_benchmark_modules/test_vllm.py
new file mode 100644
index 000000000..4bfbafc42
--- /dev/null
+++ b/tests/test_benchmark_modules/test_vllm.py
@@ -0,0 +1,1 @@
+"""Unit tests for the `vllm` module."""
diff --git a/tests/test_benchmarker.py b/tests/test_benchmarker.py
index e33dbceca..40c0eefc5 100644
--- a/tests/test_benchmarker.py
+++ b/tests/test_benchmarker.py
@@ -1,183 +1,99 @@
 """Unit tests for the `benchmarker` module."""
 
-import json
+import logging
+import os
+import time
+from collections.abc import Generator
 from pathlib import Path
-from typing import Generator, TypedDict
 
 import pytest
+import torch
 
-from scandeval import __version__
 from scandeval.benchmarker import (
     Benchmarker,
     BenchmarkResult,
+    adjust_logging_level,
+    clear_model_cache_fn,
     model_has_been_benchmarked,
+    prepare_dataset_configs,
 )
-from scandeval.types import ScoreDict
+from scandeval.dataset_configs import ANGRY_TWEETS_CONFIG, DANSK_CONFIG
+from scandeval.exceptions import HuggingFaceHubDown
 
 
-class DataKwargs(TypedDict):
-    """Helper dict with keyword arguments for `BenchmarkResult`, to avoid redundancy."""
+@pytest.fixture(scope="module")
+def benchmarker() -> Generator[Benchmarker, None, None]:
+    """A `Benchmarker` instance."""
+    yield Benchmarker(progress_bar=False, save_results=False, num_iterations=1)
 
-    num_model_parameters: int
-    max_sequence_length: int
-    vocabulary_size: int
-    dataset_languages: list[str]
-    task: str
-    results: ScoreDict
 
+def test_benchmark_results_is_a_list(benchmarker) -> None:
+    """Test that the `benchmark_results` property is a list."""
+    assert isinstance(benchmarker.benchmark_results, list)
 
-DATA_KWARGS = DataKwargs(
-    num_model_parameters=100,
-    max_sequence_length=100,
-    vocabulary_size=100,
-    dataset_languages=["da"],
-    task="task",
-    results=dict(),
-)
-
-
-def test_benchmarker_initialisation():
-    """Test that the `Benchmarker` class can be initialised."""
-    Benchmarker()
 
+def test_benchmark_encoder(benchmarker, task, language, encoder_model_id):
+    """Test that an encoder model can be benchmarked."""
+    for _ in range(10):
+        try:
+            benchmark_result = benchmarker.benchmark(
+                model=encoder_model_id, task=task.name, language=language.code
+            )
+            break
+        except HuggingFaceHubDown:
+            time.sleep(5)
+    else:
+        raise HuggingFaceHubDown()
+    assert isinstance(benchmark_result, list)
+    assert all(isinstance(result, BenchmarkResult) for result in benchmark_result)
 
-class TestBenchmarkResult:
-    """Tests related to the `BenchmarkResult` class."""
 
-    @pytest.fixture(scope="class")
-    def benchmark_result(self) -> Generator[BenchmarkResult, None, None]:
-        """Fixture for a `BenchmarkResult` object."""
-        yield BenchmarkResult(
-            dataset="dataset",
-            model="model",
-            generative=False,
-            few_shot=True,
-            validation_split=False,
-            **DATA_KWARGS,
-        )
+@pytest.mark.skipif(
+    condition=not torch.cuda.is_available(), reason="CUDA is not available."
+)
+def test_benchmark_generative(benchmarker, task, language, generative_model_id):
+    """Test that a generative model can be benchmarked."""
+    from scandeval.benchmark_modules.vllm import clear_vllm
 
-    @pytest.fixture(scope="class")
-    def results_path(self) -> Generator[Path, None, None]:
-        """Fixture for a `Path` object to a results file."""
-        results_path = Path(".scandeval_cache/test_results.jsonl")
-        results_path.parent.mkdir(parents=True, exist_ok=True)
-        yield results_path
+    for _ in range(10):
+        clear_vllm()
+        try:
+            benchmark_result = benchmarker.benchmark(
+                model=generative_model_id, task=task.name, language=language.code
+            )
+            break
+        except HuggingFaceHubDown:
+            time.sleep(5)
+    else:
+        raise HuggingFaceHubDown()
+    assert isinstance(benchmark_result, list)
+    assert all(isinstance(result, BenchmarkResult) for result in benchmark_result)
 
-    def test_benchmark_result_parameters(self, benchmark_result):
-        """Test that the `BenchmarkResult` parameters are correct."""
-        assert benchmark_result.dataset == "dataset"
-        assert benchmark_result.model == "model"
-        assert benchmark_result.generative is False
-        assert benchmark_result.few_shot is True
-        assert benchmark_result.validation_split is False
-        assert benchmark_result.num_model_parameters == 100
-        assert benchmark_result.max_sequence_length == 100
-        assert benchmark_result.vocabulary_size == 100
-        assert benchmark_result.dataset_languages == ["da"]
-        assert benchmark_result.task == "task"
-        assert benchmark_result.results == dict()
-        assert benchmark_result.scandeval_version == __version__
 
-    @pytest.mark.parametrize(
-        argnames=["config", "expected"],
-        argvalues=[
-            (
-                dict(
-                    dataset="dataset",
-                    model="model",
-                    few_shot=True,
-                    validation_split=False,
-                    **DATA_KWARGS,
-                ),
-                BenchmarkResult(
-                    dataset="dataset",
-                    model="model",
-                    generative=False,
-                    few_shot=True,
-                    validation_split=False,
-                    **DATA_KWARGS,
-                ),
-            ),
-            (
-                dict(
-                    dataset="dataset",
-                    model="model (few-shot)",
-                    validation_split=False,
-                    **DATA_KWARGS,
-                ),
-                BenchmarkResult(
-                    dataset="dataset",
-                    model="model",
-                    generative=True,
-                    few_shot=True,
-                    validation_split=False,
-                    **DATA_KWARGS,
-                ),
-            ),
-            (
-                dict(
-                    dataset="dataset", model="model (val)", few_shot=True, **DATA_KWARGS
-                ),
-                BenchmarkResult(
-                    dataset="dataset",
-                    model="model",
-                    generative=False,
-                    few_shot=True,
-                    validation_split=True,
-                    **DATA_KWARGS,
-                ),
-            ),
-            (
-                dict(dataset="dataset", model="model (few-shot, val)", **DATA_KWARGS),
-                BenchmarkResult(
-                    dataset="dataset",
-                    model="model",
-                    generative=True,
-                    few_shot=True,
-                    validation_split=True,
-                    **DATA_KWARGS,
-                ),
-            ),
-        ],
-        ids=[
-            "normal case",
-            "few-shot model name",
-            "validation split model name",
-            "few-shot and validation split model name",
-        ],
+@pytest.mark.skipif(
+    condition=os.getenv("OPENAI_API_KEY") is None,
+    reason="OpenAI API key is not available.",
+)
+def test_benchmark_openai(benchmarker, task, language, openai_model_id):
+    """Test that an OpenAI model can be benchmarked."""
+    benchmark_result = benchmarker.benchmark(
+        model=openai_model_id, task=task.name, language=language.code
     )
-    def test_from_dict(self, config, expected):
-        """Test that `BenchmarkResult.from_dict` works as expected."""
-        assert BenchmarkResult.from_dict(config) == expected
-
-    def test_append_to_results(self, benchmark_result, results_path):
-        """Test that `BenchmarkResult.append_to_results` works as expected."""
-        results_path.unlink(missing_ok=True)
-        results_path.touch(exist_ok=True)
+    assert isinstance(benchmark_result, list)
+    assert all(isinstance(result, BenchmarkResult) for result in benchmark_result)
 
-        benchmark_result.append_to_results(results_path=results_path)
-        json_str = json.dumps(
-            dict(
-                dataset=benchmark_result.dataset,
-                task=benchmark_result.task,
-                dataset_languages=benchmark_result.dataset_languages,
-                model=benchmark_result.model,
-                results=benchmark_result.results,
-                num_model_parameters=benchmark_result.num_model_parameters,
-                max_sequence_length=benchmark_result.max_sequence_length,
-                vocabulary_size=benchmark_result.vocabulary_size,
-                generative=benchmark_result.generative,
-                few_shot=benchmark_result.few_shot,
-                validation_split=benchmark_result.validation_split,
-                scandeval_version=benchmark_result.scandeval_version,
-            )
-        )
-        assert results_path.read_text() == f"\n{json_str}"
-
-        benchmark_result.append_to_results(results_path=results_path)
-        assert results_path.read_text() == f"\n{json_str}\n{json_str}"
 
-        results_path.unlink(missing_ok=True)
+@pytest.mark.skipif(
+    condition=os.getenv("ANTHROPIC_API_KEY") is None,
+    reason="Anthropic API key is not available.",
+)
+def test_benchmark_anthropic(benchmarker, task, language):
+    """Test that an Anthropic model can be benchmarked."""
+    benchmark_result = benchmarker.benchmark(
+        model="anthropic/anthropictext", task=task.name, language=language.code
+    )
+    assert isinstance(benchmark_result, list)
+    assert all(isinstance(result, BenchmarkResult) for result in benchmark_result)
 
 
 @pytest.mark.parametrize(
@@ -203,7 +119,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             True,
@@ -220,7 +141,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             False,
@@ -237,7 +163,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=True,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             False,
@@ -254,7 +185,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=True,
                     few_shot=True,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             True,
@@ -271,7 +207,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             True,
@@ -288,7 +229,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             False,
@@ -305,7 +251,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=True,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 )
             ],
             True,
@@ -322,7 +273,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 ),
                 BenchmarkResult(
                     model="model",
@@ -330,7 +286,12 @@ def test_append_to_results(self, benchmark_result, results_path):
                     generative=False,
                     few_shot=False,
                     validation_split=False,
-                    **DATA_KWARGS,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
                 ),
             ],
             True,
@@ -365,3 +326,48 @@ def test_model_has_been_benchmarked(
         benchmark_results=benchmark_results,
     )
     assert benchmarked == expected
+
+
+@pytest.mark.parametrize(
+    argnames=["verbose", "expected_logging_level"],
+    argvalues=[(False, logging.INFO), (True, logging.DEBUG)],
+)
+def test_adjust_logging_level(verbose, expected_logging_level):
+    """Test that the logging level is adjusted correctly."""
+    logging_level = adjust_logging_level(verbose=verbose, ignore_testing=True)
+    assert logging_level == expected_logging_level
+
+
+class TestClearCacheFn:
+    """Tests related to the `clear_cache_fn` function."""
+
+    def test_clear_non_existing_cache(self):
+        """Test that no errors are thrown when clearing a non-existing cache."""
+        clear_model_cache_fn(cache_dir="does-not-exist")
+
+    def test_clear_existing_cache(self):
+        """Test that a cache can be cleared."""
+        cache_dir = Path(".test_scandeval_cache")
+        model_cache_dir = cache_dir / "model_cache"
+        example_model_dir = model_cache_dir / "example_model"
+        dir_to_be_deleted = example_model_dir / "dir_to_be_deleted"
+
+        dir_to_be_deleted.mkdir(parents=True, exist_ok=True)
+        assert dir_to_be_deleted.exists()
+
+        clear_model_cache_fn(cache_dir=cache_dir.as_posix())
+        assert not dir_to_be_deleted.exists()
+        assert example_model_dir.exists()
+
+
+@pytest.mark.parametrize(
+    argnames=["dataset_names", "dataset_configs"],
+    argvalues=[
+        ([], []),
+        (["angry-tweets"], [ANGRY_TWEETS_CONFIG]),
+        (["angry-tweets", "dansk"], [ANGRY_TWEETS_CONFIG, DANSK_CONFIG]),
+    ],
+)
+def test_prepare_dataset_configs(dataset_names, dataset_configs):
+    """Test that the `prepare_dataset_configs` function works as expected."""
+    assert prepare_dataset_configs(dataset_names=dataset_names) == dataset_configs
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 2a443eef4..1154a7616 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -12,7 +12,7 @@
 @pytest.fixture(scope="module")
 def params() -> Generator[dict[str | None, ParamType], None, None]:
     """Yields a dictionary of the CLI parameters."""
-    ctx = benchmark.make_context(None, list())
+    ctx = benchmark.make_context(info_name="testing", args=["--model", "test-model"])
     yield {p.name: p.type for p in benchmark.get_params(ctx)}
 
 
@@ -26,19 +26,12 @@ def test_cli_param_names(params):
         "dataset_language",
         "dataset",
         "batch_size",
-        "evaluate_train",
         "progress_bar",
         "raise_errors",
         "verbose",
         "save_results",
         "cache_dir",
-        "token",
-        "use_token",
-        "openai_api_key",
-        "prefer_azure",
-        "azure_openai_api_key",
-        "azure_openai_endpoint",
-        "azure_openai_api_version",
+        "api_key",
         "force",
         "framework",
         "device",
@@ -46,7 +39,7 @@ def test_cli_param_names(params):
         "load_in_4bit",
         "use_flash_attention",
         "clear_model_cache",
-        "only_validation_split",
+        "evaluate_test_split",
         "few_shot",
         "num_iterations",
         "debug",
@@ -63,19 +56,12 @@ def test_cli_param_types(params):
     assert isinstance(params["dataset_language"], Choice)
     assert isinstance(params["task"], Choice)
     assert isinstance(params["batch_size"], Choice)
-    assert params["evaluate_train"] == BOOL
     assert params["progress_bar"] == BOOL
     assert params["raise_errors"] == BOOL
     assert params["verbose"] == BOOL
     assert params["save_results"] == BOOL
     assert params["cache_dir"] == STRING
-    assert params["token"] == STRING
-    assert params["use_token"] == BOOL
-    assert params["openai_api_key"] == STRING
-    assert params["prefer_azure"] == BOOL
-    assert params["azure_openai_api_key"] == STRING
-    assert params["azure_openai_endpoint"] == STRING
-    assert params["azure_openai_api_version"] == STRING
+    assert params["api_key"] == STRING
     assert params["force"] == BOOL
     assert isinstance(params["framework"], Choice)
     assert isinstance(params["device"], Choice)
@@ -83,7 +69,7 @@ def test_cli_param_types(params):
     assert params["load_in_4bit"] == BOOL
     assert params["use_flash_attention"] == BOOL
     assert params["clear_model_cache"] == BOOL
-    assert params["only_validation_split"] == BOOL
+    assert params["evaluate_test_split"] == BOOL
     assert params["few_shot"] == BOOL
     assert params["num_iterations"] == INT
     assert params["debug"] == BOOL
diff --git a/tests/test_config.py b/tests/test_config.py
deleted file mode 100644
index 645648d97..000000000
--- a/tests/test_config.py
+++ /dev/null
@@ -1,160 +0,0 @@
-"""Unit tests for the `config` module."""
-
-import pytest
-
-from scandeval.config import (
-    BenchmarkConfig,
-    DatasetConfig,
-    Language,
-    MetricConfig,
-    ModelConfig,
-    Task,
-)
-from scandeval.enums import Framework, ModelType
-
-
-class TestMetricConfig:
-    """Unit tests for the `MetricConfig` class."""
-
-    def test_metric_config_is_object(self, metric_config):
-        """Test that the metric config is a `MetricConfig` object."""
-        assert isinstance(metric_config, MetricConfig)
-
-    def test_attributes_correspond_to_arguments(self, metric_config):
-        """Test that the metric config attributes correspond to the arguments."""
-        assert metric_config.name == "metric_name"
-        assert metric_config.pretty_name == "Metric name"
-        assert metric_config.huggingface_id == "metric_id"
-        assert metric_config.results_key == "metric_key"
-
-    def test_default_value_of_compute_kwargs(self, metric_config):
-        """Test that the default value of `compute_kwargs` is an empty dictionary."""
-        assert metric_config.compute_kwargs == dict()
-
-    @pytest.mark.parametrize(
-        "inputs,expected",
-        [
-            (0.5, (50.0, "50.00%")),
-            (0.123456, (12.3456, "12.35%")),
-            (0.0, (0.0, "0.00%")),
-            (1.0, (100.0, "100.00%")),
-            (0.999999, (99.9999, "100.00%")),
-            (2.0, (200.0, "200.00%")),
-            (-1.0, (-100.0, "-100.00%")),
-        ],
-    )
-    def test_default_value_of_postprocessing_fn(self, metric_config, inputs, expected):
-        """Test that the default value of `postprocessing_fn` is correct."""
-        assert metric_config.postprocessing_fn(inputs) == expected
-
-
-class TestTask:
-    """Unit tests for the `Task` class."""
-
-    def test_task_is_object(self, task):
-        """Test that the dataset task is a `Task` object."""
-        assert isinstance(task, Task)
-
-    def test_attributes_correspond_to_arguments(self, task):
-        """Test that the dataset task attributes correspond to the arguments."""
-        assert task.name == "speed"
-        assert task.supertask == "sequence-classification"
-        assert task.labels == []
-
-
-class TestLanguage:
-    """Unit tests for the `Language` class."""
-
-    def test_language_is_object(self, language):
-        """Test that the language is a `Language` object."""
-        assert isinstance(language, Language)
-
-    def test_attributes_correspond_to_arguments(self, language):
-        """Test that the language attributes correspond to the arguments."""
-        assert language.code == "language_code"
-        assert language.name == "Language name"
-
-
-class TestBenchmarkConfig:
-    """Unit tests for the `BenchmarkConfig` class."""
-
-    def test_benchmark_config_is_object(self, benchmark_config):
-        """Test that the benchmark config is a `BenchmarkConfig` object."""
-        assert isinstance(benchmark_config, BenchmarkConfig)
-
-    def test_attributes_correspond_to_arguments(
-        self, benchmark_config, language, task, auth, device
-    ):
-        """Test that the benchmark config attributes correspond to the arguments."""
-        assert benchmark_config.model_languages == [language]
-        assert benchmark_config.dataset_languages == [language]
-        assert benchmark_config.tasks == [task]
-        assert benchmark_config.framework is None
-        assert benchmark_config.batch_size == 32
-        assert benchmark_config.raise_errors is False
-        assert benchmark_config.cache_dir == ".scandeval_cache"
-        assert benchmark_config.evaluate_train is False
-        assert benchmark_config.token is auth
-        assert benchmark_config.openai_api_key is None
-        assert benchmark_config.progress_bar is False
-        assert benchmark_config.save_results is True
-        assert benchmark_config.device == device
-        assert benchmark_config.verbose is False
-        assert benchmark_config.trust_remote_code is True
-        assert benchmark_config.load_in_4bit is None
-        assert benchmark_config.use_flash_attention is False
-        assert benchmark_config.clear_model_cache is False
-        assert benchmark_config.only_validation_split is False
-        assert benchmark_config.few_shot is True
-
-
-class TestDatasetConfig:
-    """Unit tests for the `DatasetConfig` class."""
-
-    def test_dataset_config_is_object(self, dataset_config):
-        """Test that the dataset config is a `DatasetConfig` object."""
-        assert isinstance(dataset_config, DatasetConfig)
-
-    def test_attributes_correspond_to_arguments(self, dataset_config, language, task):
-        """Test that the dataset config attributes correspond to the arguments."""
-        assert dataset_config.name == "dataset_name"
-        assert dataset_config.pretty_name == "Dataset name"
-        assert dataset_config.huggingface_id == "dataset_id"
-        assert dataset_config.task == task
-        assert dataset_config.languages == [language]
-        assert dataset_config.prompt_template == "{text}\n{label}"
-        assert dataset_config.max_generated_tokens == 1
-
-    def test_id2label(self, dataset_config):
-        """Test that the `id2label` attribute is correct."""
-        assert dataset_config.id2label == dict()
-
-    def test_label2id(self, dataset_config):
-        """Test that the `label2id` attribute is correct."""
-        assert dataset_config.label2id == dict()
-
-    def test_num_labels(self, dataset_config):
-        """Test that the `num_labels` attribute is correct."""
-        assert dataset_config.num_labels == 0
-
-    def test_default_value_of_prompt_label_mapping(self, dataset_config):
-        """Test that the default value of `prompt_label_mapping` is an empty dictionary."""
-        assert dataset_config.prompt_label_mapping == dict()
-
-
-class TestModelConfig:
-    """Unit tests for the `ModelConfig` class."""
-
-    def test_model_config_is_object(self, model_config):
-        """Test that the model config is a `ModelConfig` object."""
-        assert isinstance(model_config, ModelConfig)
-
-    def test_attributes_correspond_to_arguments(self, model_config, language):
-        """Test that the model config attributes correspond to the arguments."""
-        assert model_config.model_id == "model_id"
-        assert model_config.revision == "revision"
-        assert model_config.framework == Framework.PYTORCH
-        assert model_config.task == "task"
-        assert model_config.languages == [language]
-        assert model_config.model_type == ModelType.FRESH
-        assert model_config.model_cache_dir == "cache_dir"
diff --git a/tests/test_constants.py b/tests/test_constants.py
new file mode 100644
index 000000000..d6655218b
--- /dev/null
+++ b/tests/test_constants.py
@@ -0,0 +1,12 @@
+"""Tests for the `constants` module."""
+
+from scandeval import constants
+
+
+def test_all_objects_in_constants_are_constants():
+    """Test that all objects in the `constants` module are constants."""
+    for name in dir(constants):
+        if name.startswith("__"):
+            continue
+        assert name.isupper()
+        assert isinstance(getattr(constants, name), (int, str, list, dict))
diff --git a/tests/test_data_loading.py b/tests/test_data_loading.py
new file mode 100644
index 000000000..797c8ca17
--- /dev/null
+++ b/tests/test_data_loading.py
@@ -0,0 +1,44 @@
+"""Tests for the `data_loading` module."""
+
+from collections.abc import Generator
+
+import pytest
+from datasets import DatasetDict
+from numpy.random import default_rng
+
+from scandeval.data_loading import load_data
+from scandeval.dataset_configs import ANGRY_TWEETS_CONFIG
+
+
+class TestLoadData:
+    """Tests for the `load_data` function."""
+
+    @pytest.fixture(scope="class")
+    def dataset(self, benchmark_config) -> Generator[list[DatasetDict], None, None]:
+        """A loaded dataset."""
+        yield load_data(
+            rng=default_rng(seed=4242),
+            dataset_config=ANGRY_TWEETS_CONFIG,
+            benchmark_config=benchmark_config,
+        )
+
+    def test_load_data_is_list_of_dataset_dicts(self, dataset):
+        """Test that the `load_data` function returns a list of `DatasetDict`."""
+        assert isinstance(dataset, list)
+        assert all(isinstance(d, DatasetDict) for d in dataset)
+
+    def test_split_names_are_correct(self, dataset):
+        """Test that the split names are correct."""
+        assert all(set(d.keys()) == {"train", "val", "test"} for d in dataset)
+
+    def test_number_of_iterations_is_correct(self, dataset, benchmark_config):
+        """Test that the number of iterations is correct."""
+        assert len(dataset) == benchmark_config.num_iterations
+
+    def test_no_empty_examples(self, dataset):
+        """Test that there are no empty examples in the dataset."""
+        for dataset in dataset:
+            for split in dataset.values():
+                for feature in ["text", "tokens"]:
+                    if feature in split.features:
+                        assert all(len(x) > 0 for x in split[feature])
diff --git a/tests/test_data_models.py b/tests/test_data_models.py
new file mode 100644
index 000000000..53e072ebc
--- /dev/null
+++ b/tests/test_data_models.py
@@ -0,0 +1,263 @@
+"""Unit tests for the `data_models` module."""
+
+import inspect
+import json
+from collections.abc import Generator
+from pathlib import Path
+
+import pytest
+
+from scandeval import __version__, data_models, enums
+from scandeval.data_models import BenchmarkResult, MetricConfig
+
+
+def test_all_classes_are_dataclasses_or_pydantic_models():
+    """Test that all classes in `data_models` are dataclasses or Pydantic models."""
+    all_classes = [
+        getattr(data_models, obj_name)
+        for obj_name in dir(data_models)
+        if not obj_name.startswith("_")
+        and inspect.isclass(object=getattr(data_models, obj_name))
+        and not hasattr(enums, obj_name)
+        and obj_name != "ScoreDict"
+    ]
+    for obj in all_classes:
+        obj_is_dataclass = hasattr(obj, "__dataclass_fields__")
+        obj_is_pydantic_model = hasattr(obj, "model_fields")
+        assert (
+            obj_is_dataclass or obj_is_pydantic_model
+        ), f"{obj.__name__} is not a dataclass or Pydantic model. "
+
+
+class TestMetricConfig:
+    """Unit tests for the `MetricConfig` class."""
+
+    def test_metric_config_is_object(self, metric_config):
+        """Test that the metric config is a `MetricConfig` object."""
+        assert isinstance(metric_config, MetricConfig)
+
+    def test_attributes_correspond_to_arguments(self, metric_config):
+        """Test that the metric config attributes correspond to the arguments."""
+        assert metric_config.name == "metric_name"
+        assert metric_config.pretty_name == "Metric name"
+        assert metric_config.huggingface_id == "metric_id"
+        assert metric_config.results_key == "metric_key"
+
+    def test_default_value_of_compute_kwargs(self, metric_config):
+        """Test that the default value of `compute_kwargs` is an empty dictionary."""
+        assert metric_config.compute_kwargs == dict()
+
+    @pytest.mark.parametrize(
+        "inputs,expected",
+        [
+            (0.5, (50.0, "50.00%")),
+            (0.123456, (12.3456, "12.35%")),
+            (0.0, (0.0, "0.00%")),
+            (1.0, (100.0, "100.00%")),
+            (0.999999, (99.9999, "100.00%")),
+            (2.0, (200.0, "200.00%")),
+            (-1.0, (-100.0, "-100.00%")),
+        ],
+    )
+    def test_default_value_of_postprocessing_fn(self, metric_config, inputs, expected):
+        """Test that the default value of `postprocessing_fn` is correct."""
+        assert metric_config.postprocessing_fn(inputs) == expected
+
+
+# TODO
+@pytest.mark.skip("Not implemented")
+class TestBenchmarkConfigParams:
+    """Unit tests for the `BenchmarkConfigParams` class."""
+
+    pass
+
+
+class TestBenchmarkResult:
+    """Tests related to the `BenchmarkResult` class."""
+
+    @pytest.fixture(scope="class")
+    def benchmark_result(self) -> Generator[BenchmarkResult, None, None]:
+        """Fixture for a `BenchmarkResult` object."""
+        yield BenchmarkResult(
+            dataset="dataset",
+            model="model",
+            generative=False,
+            few_shot=True,
+            validation_split=False,
+            num_model_parameters=100,
+            max_sequence_length=100,
+            vocabulary_size=100,
+            dataset_languages=["da"],
+            task="task",
+            results=dict(),
+        )
+
+    @pytest.fixture(scope="class")
+    def results_path(self) -> Generator[Path, None, None]:
+        """Fixture for a `Path` object to a results file."""
+        results_path = Path(".scandeval_cache/test_results.jsonl")
+        results_path.parent.mkdir(parents=True, exist_ok=True)
+        yield results_path
+
+    def test_benchmark_result_parameters(self, benchmark_result):
+        """Test that the `BenchmarkResult` parameters are correct."""
+        assert benchmark_result.dataset == "dataset"
+        assert benchmark_result.model == "model"
+        assert benchmark_result.generative is False
+        assert benchmark_result.few_shot is True
+        assert benchmark_result.validation_split is False
+        assert benchmark_result.num_model_parameters == 100
+        assert benchmark_result.max_sequence_length == 100
+        assert benchmark_result.vocabulary_size == 100
+        assert benchmark_result.dataset_languages == ["da"]
+        assert benchmark_result.task == "task"
+        assert benchmark_result.results == dict()
+        assert benchmark_result.scandeval_version == __version__
+
+    @pytest.mark.parametrize(
+        argnames=["config", "expected"],
+        argvalues=[
+            (
+                dict(
+                    dataset="dataset",
+                    model="model",
+                    few_shot=True,
+                    validation_split=False,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+                BenchmarkResult(
+                    dataset="dataset",
+                    model="model",
+                    generative=False,
+                    few_shot=True,
+                    validation_split=False,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+            ),
+            (
+                dict(
+                    dataset="dataset",
+                    model="model (few-shot)",
+                    validation_split=False,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+                BenchmarkResult(
+                    dataset="dataset",
+                    model="model",
+                    generative=True,
+                    few_shot=True,
+                    validation_split=False,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+            ),
+            (
+                dict(
+                    dataset="dataset",
+                    model="model (val)",
+                    few_shot=True,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+                BenchmarkResult(
+                    dataset="dataset",
+                    model="model",
+                    generative=False,
+                    few_shot=True,
+                    validation_split=True,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+            ),
+            (
+                dict(
+                    dataset="dataset",
+                    model="model (few-shot, val)",
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+                BenchmarkResult(
+                    dataset="dataset",
+                    model="model",
+                    generative=True,
+                    few_shot=True,
+                    validation_split=True,
+                    num_model_parameters=100,
+                    max_sequence_length=100,
+                    vocabulary_size=100,
+                    dataset_languages=["da"],
+                    task="task",
+                    results=dict(),
+                ),
+            ),
+        ],
+        ids=[
+            "normal case",
+            "few-shot model name",
+            "validation split model name",
+            "few-shot and validation split model name",
+        ],
+    )
+    def test_from_dict(self, config, expected):
+        """Test that `BenchmarkResult.from_dict` works as expected."""
+        assert BenchmarkResult.from_dict(config) == expected
+
+    def test_append_to_results(self, benchmark_result, results_path):
+        """Test that `BenchmarkResult.append_to_results` works as expected."""
+        results_path.unlink(missing_ok=True)
+        results_path.touch(exist_ok=True)
+
+        benchmark_result.append_to_results(results_path=results_path)
+        json_str = json.dumps(
+            dict(
+                dataset=benchmark_result.dataset,
+                task=benchmark_result.task,
+                dataset_languages=benchmark_result.dataset_languages,
+                model=benchmark_result.model,
+                results=benchmark_result.results,
+                num_model_parameters=benchmark_result.num_model_parameters,
+                max_sequence_length=benchmark_result.max_sequence_length,
+                vocabulary_size=benchmark_result.vocabulary_size,
+                generative=benchmark_result.generative,
+                few_shot=benchmark_result.few_shot,
+                validation_split=benchmark_result.validation_split,
+                scandeval_version=benchmark_result.scandeval_version,
+            )
+        )
+        assert results_path.read_text() == f"\n{json_str}"
+
+        benchmark_result.append_to_results(results_path=results_path)
+        assert results_path.read_text() == f"\n{json_str}\n{json_str}"
+
+        results_path.unlink(missing_ok=True)
diff --git a/tests/test_dataset_configs.py b/tests/test_dataset_configs.py
index 4383c9f77..0d7874aeb 100644
--- a/tests/test_dataset_configs.py
+++ b/tests/test_dataset_configs.py
@@ -4,7 +4,7 @@
 
 import pytest
 
-from scandeval.config import DatasetConfig
+from scandeval.data_models import DatasetConfig
 from scandeval.dataset_configs import get_all_dataset_configs, get_dataset_config
 
 
diff --git a/tests/test_dataset_factory.py b/tests/test_dataset_factory.py
deleted file mode 100644
index 2ab90ca40..000000000
--- a/tests/test_dataset_factory.py
+++ /dev/null
@@ -1,101 +0,0 @@
-"""Unit tests for the `dataset_factory` module."""
-
-from copy import deepcopy
-from typing import Generator
-
-import pytest
-
-from scandeval.dataset_factory import DatasetFactory
-from scandeval.named_entity_recognition import NamedEntityRecognition
-from scandeval.question_answering import QuestionAnswering
-from scandeval.sequence_classification import SequenceClassification
-from scandeval.tasks import COMMON_SENSE, KNOW, LA, MCRC, NER, RC, SENT, SUMM
-from scandeval.text_to_text import TextToText
-
-
-@pytest.fixture(scope="module")
-def dataset_factory(benchmark_config) -> Generator[DatasetFactory, None, None]:
-    """Yields a dataset factory."""
-    yield DatasetFactory(benchmark_config)
-
-
-def test_attributes_correspond_to_arguments(dataset_factory, benchmark_config):
-    """Test that the attributes correspond to the arguments."""
-    assert dataset_factory.benchmark_config == benchmark_config
-
-
-def test_configs_are_preserved(dataset_factory, benchmark_config, dataset_config):
-    """Test that the configs are preserved."""
-    dataset = dataset_factory.build_dataset(dataset=dataset_config)
-    assert dataset.benchmark_config == benchmark_config
-    assert dataset.dataset_config == dataset_config
-
-
-def test_build_sent_dataset(dataset_factory, dataset_config):
-    """Test that SENT datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = SENT
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, SequenceClassification)
-
-
-def test_build_la_dataset(dataset_factory, dataset_config):
-    """Test that LA datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = LA
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, SequenceClassification)
-
-
-def test_build_ner_dataset(dataset_factory, dataset_config):
-    """Test that NER datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = NER
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, NamedEntityRecognition)
-
-
-def test_build_qa_dataset(dataset_factory, dataset_config):
-    """Test that RC datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = RC
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, QuestionAnswering)
-
-
-def test_build_summ_dataset(dataset_factory, dataset_config):
-    """Test that SUMM datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = SUMM
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, TextToText)
-
-
-def test_build_know_dataset(dataset_factory, dataset_config):
-    """Test that KNOW datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = KNOW
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, SequenceClassification)
-
-
-def test_build_common_sense_dataset(dataset_factory, dataset_config):
-    """Test that COMMON_SENSE datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = COMMON_SENSE
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, SequenceClassification)
-
-
-def test_build_mcrc_dataset(dataset_factory, dataset_config):
-    """Test that MCRC datasets are built correctly."""
-    cfg = deepcopy(dataset_config)
-    cfg.task = MCRC
-    dataset = dataset_factory.build_dataset(dataset=cfg)
-    assert isinstance(dataset, SequenceClassification)
-
-
-@pytest.mark.skip(reason="Text modelling datasets are not yet implemented.")
-def test_build_text_modelling_dataset(dataset_factory, dataset_config):
-    """Test that TEXT_MODELLING datasets are built correctly."""
-    pass
diff --git a/tests/test_enums.py b/tests/test_enums.py
index de3c849b7..4a3be42a4 100644
--- a/tests/test_enums.py
+++ b/tests/test_enums.py
@@ -1,1 +1,19 @@
-"""Unit tests for the `enums` module."""
+"""Tests for the `enums` module."""
+
+import enum
+import inspect
+
+from scandeval import enums
+
+
+def test_all_classes_are_enums():
+    """Test that all classes in `enums` are Enums."""
+    all_classes = [
+        getattr(enums, obj_name)
+        for obj_name in dir(enums)
+        if not obj_name.startswith("_")
+        and inspect.isclass(object=getattr(enums, obj_name))
+        and not hasattr(enum, obj_name)
+    ]
+    for obj in all_classes:
+        assert issubclass(obj, enum.Enum)
diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
index f47095825..429a7597b 100644
--- a/tests/test_exceptions.py
+++ b/tests/test_exceptions.py
@@ -1,101 +1,17 @@
 """Unit tests for the `exceptions` module."""
 
-import pytest
+import inspect
 
-from scandeval.exceptions import (
-    FlashAttentionNotInstalled,
-    HuggingFaceHubDown,
-    InvalidBenchmark,
-    InvalidModel,
-    MissingHuggingFaceToken,
-    NaNValueInModelOutput,
-    NeedsAdditionalArgument,
-    NeedsExtraInstalled,
-    NoInternetConnection,
-)
+from scandeval import exceptions
 
 
-def test_invalid_benchmark_is_an_exception():
-    """Test that `InvalidBenchmark` is an exception."""
-    with pytest.raises(InvalidBenchmark):
-        raise InvalidBenchmark()
-
-
-def test_invalid_model_is_an_exception():
-    """Test that `InvalidModel` is an exception."""
-    with pytest.raises(InvalidModel):
-        raise InvalidModel()
-
-
-def test_hugging_face_hub_down_is_an_exception():
-    """Test that `HuggingFaceHubDown` is an exception."""
-    with pytest.raises(HuggingFaceHubDown):
-        raise HuggingFaceHubDown()
-
-
-def test_no_internet_connection_is_an_exception():
-    """Test that `NoInternetConnection` is an exception."""
-    with pytest.raises(NoInternetConnection):
-        raise NoInternetConnection()
-
-
-def test_nan_value_in_model_output():
-    """Test that `NaNValueInModelOutput` is an exception."""
-    with pytest.raises(NaNValueInModelOutput):
-        raise NaNValueInModelOutput()
-
-
-def test_flash_attention_not_installed():
-    """Test that `FlashAttentionNotInstalled` is an exception."""
-    with pytest.raises(FlashAttentionNotInstalled):
-        raise FlashAttentionNotInstalled()
-
-
-@pytest.mark.parametrize(
-    argnames=["extra"], argvalues=[("test",), ("",)], ids=["extra", "empty_extra"]
-)
-def test_needs_extra_installed(extra):
-    """Test that `NeedsExtraInstalled` is an exception."""
-    exception_regex = f".*`{extra}`.*"
-    with pytest.raises(NeedsExtraInstalled, match=exception_regex):
-        raise NeedsExtraInstalled(extra)
-
-
-@pytest.mark.parametrize(
-    argnames=["cli_argument", "script_argument", "run_with_cli"],
-    argvalues=[
-        ("--test-argument", "test_argument", True),
-        ("--test-argument", "test_argument", False),
-        ("--", "", True),
-        ("--", "", False),
-    ],
-    ids=[
-        "argument_with_cli",
-        "argument_without_cli",
-        "empty_argument_with_cli",
-        "empty_argument_without_cli",
-    ],
-)
-def test_needs_additional_argument(cli_argument, script_argument, run_with_cli):
-    """Test that `NeedsExtraInstalled` is an exception."""
-    if run_with_cli:
-        exception_regex = f".*`{cli_argument}`.*`scandeval` command.*"
-    else:
-        exception_regex = f".*`{script_argument}`.*`Benchmarker` class.*"
-    with pytest.raises(NeedsAdditionalArgument, match=exception_regex):
-        raise NeedsAdditionalArgument(
-            cli_argument=cli_argument,
-            script_argument=script_argument,
-            run_with_cli=run_with_cli,
-        )
-
-
-@pytest.mark.parametrize(argnames=["run_with_cli"], argvalues=[(True,), (False,)])
-def test_missing_hugging_face_token(run_with_cli):
-    """Test that `NeedsExtraInstalled` is an exception."""
-    if run_with_cli:
-        exception_regex = ".*`huggingface-cli login`.*"
-    else:
-        exception_regex = ".*`Benchmarker` class.*"
-    with pytest.raises(MissingHuggingFaceToken, match=exception_regex):
-        raise MissingHuggingFaceToken(run_with_cli=run_with_cli)
+def test_all_classes_are_exceptions():
+    """Test that all classes in `exceptions` are exceptions."""
+    all_classes = [
+        getattr(exceptions, obj_name)
+        for obj_name in dir(exceptions)
+        if not obj_name.startswith("_")
+        and inspect.isclass(object=getattr(exceptions, obj_name))
+    ]
+    for obj in all_classes:
+        assert issubclass(obj, Exception), f"Class {obj.__name__} is not an exception."
diff --git a/tests/test_finetuning.py b/tests/test_finetuning.py
index 791df948e..147fd0a2f 100644
--- a/tests/test_finetuning.py
+++ b/tests/test_finetuning.py
@@ -1,1 +1,162 @@
-"""Unit tests for the `finetuning` module."""
+"""Tests for the `finetuning` module."""
+
+import pytest
+import torch
+from transformers import IntervalStrategy, TrainingArguments
+
+from scandeval.enums import DataType
+from scandeval.finetuning import get_training_args
+
+
+# TODO
+def test_finetune():
+    """Test that the `finetune` function works as expected."""
+    pass
+
+
+# TODO
+def test_finetune_single_iteration():
+    """Test that the `finetune_single_iteration` function works as expected."""
+    pass
+
+
+class TestGetTrainingArgs:
+    """Test that the `get_training_args` function works as expected."""
+
+    def test_return_type(self, benchmark_config, model_config):
+        """Test that the return type is correct."""
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=None,
+        )
+        assert isinstance(args, TrainingArguments)
+
+    @pytest.mark.parametrize(argnames=["batch_size"], argvalues=[(8,), (16,), (32,)])
+    def test_both_batch_sizes_are_set(self, benchmark_config, model_config, batch_size):
+        """Test that both training and eval batch sizes are set."""
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=batch_size,
+        )
+        assert args.per_device_train_batch_size == batch_size
+        assert args.per_device_eval_batch_size == batch_size
+
+    @pytest.mark.parametrize(
+        argnames=["batch_size", "expected_gradient_accumulation_steps"],
+        argvalues=[(1, 32), (2, 16), (4, 8), (8, 4), (16, 2), (32, 1)],
+    )
+    def test_gradient_accumulation(
+        self,
+        benchmark_config,
+        model_config,
+        batch_size,
+        expected_gradient_accumulation_steps,
+    ):
+        """Test that the gradient accumulation is correct."""
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=batch_size,
+        )
+        assert args.gradient_accumulation_steps == expected_gradient_accumulation_steps
+
+    def test_batch_size_default_value(self, benchmark_config, model_config):
+        """Test that the default value for the batch size is correct."""
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=None,
+        )
+        assert args.per_device_train_batch_size == benchmark_config.batch_size
+
+    @pytest.mark.parametrize(
+        argnames=["verbose", "expected_logging_strategy"],
+        argvalues=[(True, IntervalStrategy.STEPS), (False, IntervalStrategy.NO)],
+    )
+    def test_logging_strategy(
+        self, benchmark_config, model_config, verbose, expected_logging_strategy
+    ):
+        """Test that the logging strategy is correct."""
+        old_verbose = benchmark_config.verbose
+        benchmark_config.verbose = verbose
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=None,
+        )
+        assert args.logging_strategy == expected_logging_strategy
+        benchmark_config.verbose = old_verbose
+
+    @pytest.mark.parametrize(
+        argnames=["progress_bar", "expected_disable_tqdm"],
+        argvalues=[(True, False), (False, True)],
+    )
+    def test_disable_tqdm(
+        self, benchmark_config, model_config, progress_bar, expected_disable_tqdm
+    ):
+        """Test that the disable tqdm option is correct."""
+        old_progress_bar = benchmark_config.progress_bar
+        benchmark_config.progress_bar = progress_bar
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=None,
+        )
+        assert args.disable_tqdm == expected_disable_tqdm
+        benchmark_config.progress_bar = old_progress_bar
+
+    @pytest.mark.parametrize(
+        argnames=["datatype", "expected_fp16", "expected_bf16"],
+        argvalues=[
+            (DataType.FP32, False, False),
+            (DataType.FP16, True, False),
+            (DataType.BF16, False, True),
+        ],
+    )
+    def test_dtype(
+        self, benchmark_config, model_config, datatype, expected_fp16, expected_bf16
+    ):
+        """Test that the fp16 and bf16 arguments have been correctly set."""
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=datatype,
+            batch_size=None,
+        )
+        assert args.fp16 == expected_fp16
+        assert args.bf16 == expected_bf16
+
+    @pytest.mark.parametrize(
+        argnames=["device_name", "expected_use_cpu"],
+        argvalues=[("cuda", False), ("mps", False), ("cpu", True)],
+    )
+    def test_use_cpu(
+        self, benchmark_config, model_config, device_name, expected_use_cpu
+    ):
+        """Test that the use_cpu argument is correct."""
+        old_device = benchmark_config.device
+        benchmark_config.device = torch.device(device_name)
+        args = get_training_args(
+            benchmark_config=benchmark_config,
+            model_config=model_config,
+            iteration_idx=0,
+            dtype=DataType.FP32,
+            batch_size=None,
+        )
+        assert args.use_cpu == expected_use_cpu
+        benchmark_config.device = old_device
diff --git a/tests/test_generation.py b/tests/test_generation.py
index ceff57a3b..66aef60ee 100644
--- a/tests/test_generation.py
+++ b/tests/test_generation.py
@@ -1,1 +1,19 @@
-"""Unit tests for the `generation` module."""
+"""Tests for the `generation` module."""
+
+
+# TODO
+def test_generate():
+    """Test that the `generate` function works as expected."""
+    pass
+
+
+# TODO
+def test_generate_single_iteration():
+    """Test that the `generate_single_iteration` function works as expected."""
+    pass
+
+
+# TODO
+def test_debug_log():
+    """Test that the `debug_log` function works as expected."""
+    pass
diff --git a/tests/test_human_evaluation.py b/tests/test_human_evaluation.py
new file mode 100644
index 000000000..34e9b10f1
--- /dev/null
+++ b/tests/test_human_evaluation.py
@@ -0,0 +1,8 @@
+"""Tests for the `human_evaluation` module."""
+
+
+# TODO
+class TestHumanEvaluator:
+    """Tests for the `HumanEvaluator` class."""
+
+    pass
diff --git a/tests/test_languages.py b/tests/test_languages.py
index 5340da99a..aed6fc71a 100644
--- a/tests/test_languages.py
+++ b/tests/test_languages.py
@@ -4,7 +4,7 @@
 
 import pytest
 
-from scandeval.config import Language
+from scandeval.data_models import Language
 from scandeval.languages import get_all_languages
 
 
diff --git a/tests/test_model_cache.py b/tests/test_model_cache.py
index 05db4156a..467bf0970 100644
--- a/tests/test_model_cache.py
+++ b/tests/test_model_cache.py
@@ -1,1 +1,46 @@
-"""Unit tests for the `model_cache` module."""
+"""Tests for the `model_cache` module."""
+
+
+# TODO
+class TestModelCache:
+    """Test that the `ModelCache` class works as expected."""
+
+    def test_initialisation(self):
+        """Test that the `ModelCache` class can be instantiated."""
+        pass
+
+    def test_save_and_load(self):
+        """Test that the `ModelCache` class can be saved and loaded."""
+        pass
+
+    def test_getitem(self):
+        """Test that the `ModelCache` class can be accessed like a dictionary."""
+        pass
+
+    def test_setitem(self):
+        """Test that the `ModelCache` class can be set like a dictionary."""
+        pass
+
+    def test_remove(self):
+        """Test that the `ModelCache` class can be removed."""
+        pass
+
+    def test_cached_texts(self):
+        """Test that the `ModelCache` class can return all the cached texts."""
+        pass
+
+    def test_add_to_cache(self):
+        """Test that the `ModelCache` class can add items to the cache."""
+        pass
+
+
+# TODO
+def test_split_dataset_into_cached_and_non_cached():
+    """Test that the `split_dataset_into_cached_and_non_cached` function works."""
+    pass
+
+
+# TODO
+def test_load_cached_model_outputs():
+    """Test that the `load_cached_model_outputs` function works."""
+    pass
diff --git a/tests/test_model_config.py b/tests/test_model_config.py
index e804940ad..49e2765af 100644
--- a/tests/test_model_config.py
+++ b/tests/test_model_config.py
@@ -1,1 +1,36 @@
-"""Unit tests for the `model_config` module."""
+"""Tests for the `model_config` module."""
+
+import pytest
+
+from scandeval.data_models import ModelConfig
+from scandeval.exceptions import InvalidModel
+from scandeval.model_config import get_model_config
+
+
+@pytest.mark.parametrize(
+    argnames=["model_id", "should_raise"],
+    argvalues=[
+        ("Maltehb/aelaectra-danish-electra-small-cased", False),
+        ("openai-community/gpt2", False),
+        ("gpt-4o-mini", False),
+        ("claude-3-5-haiku-20241022", False),
+        ("does-not-exist", True),
+    ],
+    ids=[
+        "encoder-model",
+        "decoder-model",
+        "openai-model",
+        "anthropic-model",
+        "non-existent-model",
+    ],
+)
+def test_get_model_config(benchmark_config, model_id, should_raise) -> None:
+    """Test that the `get_model_config` function works as expected."""
+    if should_raise:
+        with pytest.raises(InvalidModel):
+            get_model_config(model_id=model_id, benchmark_config=benchmark_config)
+    else:
+        model_config = get_model_config(
+            model_id=model_id, benchmark_config=benchmark_config
+        )
+        assert isinstance(model_config, ModelConfig)
diff --git a/tests/test_model_loading.py b/tests/test_model_loading.py
index 0ea67b1d2..4572ace7b 100644
--- a/tests/test_model_loading.py
+++ b/tests/test_model_loading.py
@@ -1,72 +1,47 @@
 """Unit tests for the `model_loading` module."""
 
-import os
-
 import pytest
+import torch
 
-from scandeval.config import ModelConfig
-from scandeval.enums import Framework
-from scandeval.exceptions import InvalidBenchmark, InvalidModel
-from scandeval.languages import DA
+from scandeval.dataset_configs import ANGRY_TWEETS_CONFIG, NORDJYLLAND_NEWS_CONFIG
+from scandeval.exceptions import InvalidBenchmark
 from scandeval.model_config import get_model_config
 from scandeval.model_loading import load_model
 
 
-def test_load_non_generative_model(model_id, dataset_config, benchmark_config):
+def test_load_non_generative_model(encoder_model_id, benchmark_config):
     """Test loading a non-generative model."""
     model_config = get_model_config(
-        model_id=model_id, benchmark_config=benchmark_config
+        model_id=encoder_model_id, benchmark_config=benchmark_config
     )
-    tokenizer, model = load_model(
+    model = load_model(
         model_config=model_config,
-        dataset_config=dataset_config,
+        dataset_config=ANGRY_TWEETS_CONFIG,
         benchmark_config=benchmark_config,
     )
-    assert tokenizer is not None
     assert model is not None
 
 
 @pytest.mark.skipif(
-    condition=os.getenv("USE_VLLM", "0") != "1", reason="Not using VLLM."
+    condition=not torch.cuda.is_available(), reason="CUDA is not available."
 )
-def test_load_generative_model(generative_model_and_tokenizer):
+def test_load_generative_model(generative_model):
     """Test loading a generative model."""
-    model, tokenizer = generative_model_and_tokenizer
-    assert tokenizer is not None
+    model = generative_model
     assert model is not None
 
 
 def test_load_non_generative_model_with_generative_data(
-    model_id, generative_dataset_config, benchmark_config
+    encoder_model_id, benchmark_config
 ):
     """Test loading a non-generative model with generative data."""
     model_config = get_model_config(
-        model_id=model_id, benchmark_config=benchmark_config
+        model_id=encoder_model_id, benchmark_config=benchmark_config
     )
     with pytest.raises(InvalidBenchmark):
         load_model(
             model_config=model_config,
-            dataset_config=generative_dataset_config,
-            benchmark_config=benchmark_config,
-        )
-
-
-def test_load_non_existing_model(dataset_config, benchmark_config):
-    """Test loading a non-existing model."""
-    model_config = ModelConfig(
-        model_id="non-existing-model",
-        revision="revision",
-        framework=Framework.PYTORCH,
-        task="task",
-        languages=[DA],
-        model_type="fresh",
-        model_cache_dir="cache_dir",
-        adapter_base_model_id=None,
-    )
-    with pytest.raises(InvalidModel):
-        load_model(
-            model_config=model_config,
-            dataset_config=dataset_config,
+            dataset_config=NORDJYLLAND_NEWS_CONFIG,
             benchmark_config=benchmark_config,
         )
 
@@ -76,9 +51,7 @@ def test_load_non_existing_model(dataset_config, benchmark_config):
 class TestQuantisedModels:
     """Tests for quantised models."""
 
-    def test_load_awq_model(
-        self, awq_generative_model_id, dataset_config, benchmark_config
-    ):
+    def test_load_awq_model(self, awq_generative_model_id, benchmark_config):
         """Test loading an AWQ quantised model."""
         if benchmark_config.device.type != "cuda":
             pytest.skip("Skipping test because the device is not a GPU.")
@@ -87,13 +60,11 @@ def test_load_awq_model(
         )
         load_model(
             model_config=model_config,
-            dataset_config=dataset_config,
+            dataset_config=ANGRY_TWEETS_CONFIG,
             benchmark_config=benchmark_config,
         )
 
-    def test_load_gptq_model(
-        self, gptq_generative_model_id, dataset_config, benchmark_config
-    ):
+    def test_load_gptq_model(self, gptq_generative_model_id, benchmark_config):
         """Test loading a GPTQ quantised model."""
         if benchmark_config.device.type != "cuda":
             pytest.skip("Skipping test because the device is not a GPU.")
@@ -102,6 +73,6 @@ def test_load_gptq_model(
         )
         load_model(
             model_config=model_config,
-            dataset_config=dataset_config,
+            dataset_config=ANGRY_TWEETS_CONFIG,
             benchmark_config=benchmark_config,
         )
diff --git a/tests/test_model_setups/__init__.py b/tests/test_model_setups/__init__.py
deleted file mode 100644
index cd123e86c..000000000
--- a/tests/test_model_setups/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Tests for the model setups."""
diff --git a/tests/test_model_setups/test_fresh.py b/tests/test_model_setups/test_fresh.py
deleted file mode 100644
index 316068a76..000000000
--- a/tests/test_model_setups/test_fresh.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `model_setups.fresh` module."""
diff --git a/tests/test_model_setups/test_hf.py b/tests/test_model_setups/test_hf.py
deleted file mode 100644
index 3f4eef794..000000000
--- a/tests/test_model_setups/test_hf.py
+++ /dev/null
@@ -1,62 +0,0 @@
-"""Unit tests for the `model_setups.hf` module."""
-
-import copy
-
-import pytest
-import torch
-
-from scandeval.model_setups.hf import HFModelSetup
-
-
-@pytest.mark.parametrize(
-    argnames=["test_device", "model_id", "expected"],
-    argvalues=[
-        ("cuda", "microsoft/mdeberta-v3-base", torch.float16),
-        ("cuda", "jonfd/electra-small-nordic", torch.float16),
-        ("cuda", "mistralai/Mistral-7B-v0.1", "auto"),
-        ("cuda", "microsoft/phi-2", "auto"),
-        ("mps", "microsoft/mdeberta-v3-base", torch.float32),
-        ("mps", "jonfd/electra-small-nordic", torch.float32),
-        ("mps", "mistralai/Mistral-7B-v0.1", torch.float32),
-        ("mps", "microsoft/phi-2", torch.float32),
-        ("cpu", "microsoft/mdeberta-v3-base", torch.float32),
-        ("cpu", "jonfd/electra-small-nordic", torch.float32),
-        ("cpu", "mistralai/Mistral-7B-v0.1", torch.float32),
-        ("cpu", "microsoft/phi-2", torch.float32),
-    ],
-    ids=[
-        "cuda:mdeberta-v3-base",
-        "cuda:electra-small-nordic",
-        "cuda:mistral-7b",
-        "cuda:phi-2",
-        "mps:mdeberta-v3-base",
-        "mps:electra-small-nordic",
-        "mps:mistral-7b",
-        "mps:phi-2",
-        "cpu:mdeberta-v3-base",
-        "cpu:electra-small-nordic",
-        "cpu:mistral-7b",
-        "cpu:phi-2",
-    ],
-)
-def test_torch_dtype_is_set_correctly(
-    benchmark_config, test_device, model_id, expected
-):
-    """Test that the torch dtype is set correctly."""
-    benchmark_config_copy = copy.deepcopy(benchmark_config)
-    benchmark_config_copy.device = torch.device(test_device)
-    model_setup = HFModelSetup(benchmark_config=benchmark_config_copy)
-    hf_model_config = model_setup._load_hf_model_config(
-        model_id=model_id,
-        num_labels=0,
-        id2label=dict(),
-        label2id=dict(),
-        revision="main",
-        model_cache_dir=benchmark_config_copy.cache_dir,
-    )
-
-    bf16_available = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
-    if test_device == "cuda" and expected == torch.float16 and bf16_available:
-        expected = torch.bfloat16
-
-    assert model_setup._get_torch_dtype(config=hf_model_config) == expected
diff --git a/tests/test_model_setups/test_local.py b/tests/test_model_setups/test_local.py
deleted file mode 100644
index b18e255f4..000000000
--- a/tests/test_model_setups/test_local.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `model_setups.local` module."""
diff --git a/tests/test_model_setups/test_openai.py b/tests/test_model_setups/test_openai.py
deleted file mode 100644
index 27cdcef1f..000000000
--- a/tests/test_model_setups/test_openai.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `model_setups.openai` module."""
diff --git a/tests/test_model_setups/test_utils.py b/tests/test_model_setups/test_utils.py
deleted file mode 100644
index 250e45b34..000000000
--- a/tests/test_model_setups/test_utils.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `model_setups.utils` module."""
diff --git a/tests/test_named_entity_recognition.py b/tests/test_named_entity_recognition.py
deleted file mode 100644
index 8d9bc9f08..000000000
--- a/tests/test_named_entity_recognition.py
+++ /dev/null
@@ -1,60 +0,0 @@
-"""Unit tests for the `named_entity_recognition` module."""
-
-import os
-from contextlib import nullcontext as does_not_raise
-from typing import Generator
-
-import pytest
-
-from scandeval.benchmark_dataset import BenchmarkDataset
-from scandeval.dataset_configs import get_all_dataset_configs
-from scandeval.exceptions import InvalidBenchmark
-from scandeval.languages import DA
-from scandeval.named_entity_recognition import NamedEntityRecognition
-from scandeval.tasks import NER
-from scandeval.utils import GENERATIVE_DATASET_TASKS
-
-
-@pytest.fixture(
-    scope="module",
-    params=[
-        dataset_config
-        for dataset_config in get_all_dataset_configs().values()
-        if dataset_config.task == NER
-        and (
-            os.getenv("TEST_ALL_DATASETS", "0") == "1"
-            or (not dataset_config.unofficial and dataset_config.languages == [DA])
-        )
-    ],
-    ids=lambda dataset_config: dataset_config.name,
-)
-def benchmark_dataset(
-    benchmark_config, request
-) -> Generator[BenchmarkDataset, None, None]:
-    """Yields a named entity recognition dataset."""
-    yield NamedEntityRecognition(
-        dataset_config=request.param, benchmark_config=benchmark_config
-    )
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_encoder_benchmarking(benchmark_dataset, model_id):
-    """Test that encoder models can be benchmarked on named entity recognition."""
-    if benchmark_dataset.dataset_config.task.name in GENERATIVE_DATASET_TASKS:
-        with pytest.raises(InvalidBenchmark):
-            benchmark_dataset.benchmark(model_id)
-    else:
-        with does_not_raise():
-            benchmark_dataset.benchmark(model_id)
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_decoder_benchmarking(
-    benchmark_dataset, generative_model_id, generative_model_and_tokenizer
-):
-    """Test that decoder models can be benchmarked on named entity recognition."""
-    model, tokenizer = generative_model_and_tokenizer
-    with does_not_raise():
-        benchmark_dataset.benchmark(
-            model_id=generative_model_id, model=model, tokenizer=tokenizer
-        )
diff --git a/tests/test_openai_models.py b/tests/test_openai_models.py
deleted file mode 100644
index ebb900a04..000000000
--- a/tests/test_openai_models.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `openai_models` module."""
diff --git a/tests/test_protocols.py b/tests/test_protocols.py
deleted file mode 100644
index a44f03bfa..000000000
--- a/tests/test_protocols.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `protocols` module."""
diff --git a/tests/test_question_answering.py b/tests/test_question_answering.py
deleted file mode 100644
index 63746695f..000000000
--- a/tests/test_question_answering.py
+++ /dev/null
@@ -1,97 +0,0 @@
-"""Unit tests for the `question_answering` module."""
-
-import os
-from contextlib import nullcontext as does_not_raise
-from typing import Generator
-
-import pytest
-from transformers import AutoTokenizer
-
-from scandeval.benchmark_dataset import BenchmarkDataset
-from scandeval.dataset_configs import get_all_dataset_configs
-from scandeval.exceptions import InvalidBenchmark
-from scandeval.languages import DA
-from scandeval.question_answering import QuestionAnswering, prepare_train_examples
-from scandeval.tasks import RC
-from scandeval.utils import GENERATIVE_DATASET_TASKS
-
-
-@pytest.fixture(
-    scope="module",
-    params=[
-        dataset_config
-        for dataset_config in get_all_dataset_configs().values()
-        if dataset_config.task == RC
-        and (
-            os.getenv("TEST_ALL_DATASETS", "0") == "1"
-            or (not dataset_config.unofficial and dataset_config.languages == [DA])
-        )
-    ],
-    ids=lambda dataset_config: dataset_config.name,
-)
-def benchmark_dataset(
-    benchmark_config, request
-) -> Generator[BenchmarkDataset, None, None]:
-    """Yields a question answering benchmark dataset."""
-    yield QuestionAnswering(
-        dataset_config=request.param, benchmark_config=benchmark_config
-    )
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_encoder_benchmarking(benchmark_dataset, model_id):
-    """Test that encoder models can be benchmarked on question answering datasets."""
-    if benchmark_dataset.dataset_config.task.name in GENERATIVE_DATASET_TASKS:
-        with pytest.raises(InvalidBenchmark):
-            benchmark_dataset.benchmark(model_id)
-    else:
-        with does_not_raise():
-            benchmark_dataset.benchmark(model_id)
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_decoder_benchmarking(
-    benchmark_dataset, generative_model_id, generative_model_and_tokenizer
-):
-    """Test that decoder models can be benchmarked on question answering datasets."""
-    model, tokenizer = generative_model_and_tokenizer
-    with does_not_raise():
-        benchmark_dataset.benchmark(
-            model_id=generative_model_id, model=model, tokenizer=tokenizer
-        )
-
-
-@pytest.mark.parametrize(
-    argnames="tokenizer_model_id",
-    argvalues=["jonfd/electra-small-nordic", "flax-community/swe-roberta-wiki-oscar"],
-)
-@pytest.mark.parametrize(
-    argnames="examples",
-    argvalues=[
-        dict(
-            question=["Hvad er hovedstaden i Sverige?"],
-            context=["Sveriges hovedstad er Stockholm."],
-            answers=[dict(text=["Stockholm"], answer_start=[22])],
-        ),
-        dict(
-            question=["Hvad er hovedstaden i Sverige?"],
-            context=["Sveriges hovedstad er Stockholm." * 100],
-            answers=[dict(text=["Sverige"], answer_start=[0])],
-        ),
-        dict(
-            question=["Hvad er hovedstaden i Danmark?"],
-            context=["Danmarks hovedstad er København. " * 100],
-            answers=[dict(text=["Da"], answer_start=[0])],
-        ),
-    ],
-)
-def test_prepare_train_examples(examples, tokenizer_model_id):
-    """Test that train examples can be prepared for training."""
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_model_id)
-    if (
-        hasattr(tokenizer, "model_max_length")
-        and tokenizer.model_max_length > 100_000_000
-    ):
-        tokenizer.model_max_length = 512
-    with does_not_raise():
-        prepare_train_examples(examples=examples, tokenizer=tokenizer)
diff --git a/tests/test_question_answering_trainer.py b/tests/test_question_answering_trainer.py
deleted file mode 100644
index 43d013fa6..000000000
--- a/tests/test_question_answering_trainer.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `question_answering_trainer` module."""
diff --git a/tests/test_scores.py b/tests/test_scores.py
index be11a82c6..435b9c9b9 100644
--- a/tests/test_scores.py
+++ b/tests/test_scores.py
@@ -1,6 +1,5 @@
 """Unit tests for the `scores` module."""
 
-from copy import deepcopy
 from typing import Generator
 
 import numpy as np
@@ -11,94 +10,35 @@
 
 
 @pytest.fixture(scope="module")
-def scores(metric_config) -> Generator[dict[str, list[dict[str, float]]], None, None]:
+def scores(metric_config) -> Generator[list[dict[str, float]], None, None]:
     """Yield a dictionary of scores."""
-    yield dict(
-        train=[
-            {f"train_{metric_config.name}": 0.70},
-            {f"train_{metric_config.name}": 0.75},
-            {f"train_{metric_config.name}": 0.80},
-        ],
-        test=[
-            {f"test_{metric_config.name}": 0.50},
-            {f"test_{metric_config.name}": 0.55},
-            {f"test_{metric_config.name}": 0.60},
-        ],
-    )
+    yield [
+        {f"test_{metric_config.name}": 0.50},
+        {f"test_{metric_config.name}": 0.55},
+        {f"test_{metric_config.name}": 0.60},
+    ]
 
 
 class TestAggregateScores:
     """Unit tests for the `aggregate_scores` function."""
 
-    def test_only_train_scores(self, scores, metric_config):
-        """Test that `aggregate_scores` works when only train scores are provided."""
-        # Remove the test scores from the scores
-        scores_only_train = deepcopy(scores)
-        scores_only_train.pop("test")
-
-        # Aggregate scores using the `agg_scores` function
-        agg_scores = aggregate_scores(
-            scores=scores_only_train, metric_config=metric_config
-        )
-
-        # Manually compute the mean and standard error of the scores
-        train_scores = [
-            dct[f"train_{metric_config.name}"] for dct in scores_only_train["train"]
-        ]
-        mean = np.mean(train_scores)
-        se = 1.96 * np.std(train_scores, ddof=1) / np.sqrt(len(train_scores))
-
-        # Assert that `aggregate_scores` computed the same
-        assert agg_scores == dict(train=(mean, se))
-
-    def test_only_test_scores(self, scores, metric_config):
-        """Test that `aggregate_scores` works when only test scores are provided."""
-        # Remove the train scores from the scores
-        scores_only_test = deepcopy(scores)
-        scores_only_test.pop("train")
-
-        # Aggregate scores using the `agg_scores` function
-        agg_scores = aggregate_scores(
-            scores=scores_only_test, metric_config=metric_config
-        )
-
-        # Manually compute the mean and standard error of the scores
-        test_scores = [
-            dct[f"test_{metric_config.name}"] for dct in scores_only_test["test"]
-        ]
-        mean = np.mean(test_scores)
-        se = 1.96 * np.std(test_scores, ddof=1) / np.sqrt(len(test_scores))
-
-        # Assert that `aggregate_scores` computed the same
-        assert agg_scores == dict(test=(mean, se))
-
-    def test_all_scores(self, scores, metric_config):
-        """Test that `aggregate_scores` works when train/test scores are provided."""
+    def test_scores(self, scores, metric_config):
+        """Test that `aggregate_scores` works when scores are provided."""
         # Aggregate scores using the `agg_scores` function
         agg_scores = aggregate_scores(scores=scores, metric_config=metric_config)
 
-        # Manually compute the mean and standard error of the train scores
-        train_scores = [dct[f"train_{metric_config.name}"] for dct in scores["train"]]
-        train_mean = np.mean(train_scores)
-        train_se = 1.96 * np.std(train_scores, ddof=1) / np.sqrt(len(train_scores))
-
         # Manually compute the mean and standard error of the test scores
-        test_scores = [dct[f"test_{metric_config.name}"] for dct in scores["test"]]
+        test_scores = [dct[f"test_{metric_config.name}"] for dct in scores]
         test_mean = np.mean(test_scores)
         test_se = 1.96 * np.std(test_scores, ddof=1) / np.sqrt(len(test_scores))
 
         # Assert that `aggregate_scores` computed the same
-        assert agg_scores == dict(
-            train=(train_mean, train_se), test=(test_mean, test_se)
-        )
+        assert agg_scores == (test_mean, test_se)
 
-    def test_no_scores(self, scores, metric_config):
+    def test_no_scores(self, metric_config):
         """Test that `aggregate_scores` works when no scores are provided."""
-        empty_scores = deepcopy(scores)
-        empty_scores.pop("train")
-        empty_scores.pop("test")
-        agg_scores = aggregate_scores(scores=empty_scores, metric_config=metric_config)
-        assert agg_scores == dict()
+        agg_scores = aggregate_scores(scores=list(), metric_config=metric_config)
+        assert np.isnan(agg_scores).all()
 
 
 class TestLogScores:
@@ -135,8 +75,6 @@ def test_total_scores_keys(self, logged_scores, metric_config):
         assert sorted(logged_scores["total"].keys()) == [
             f"test_{metric_config.name}",
             f"test_{metric_config.name}_se",
-            f"train_{metric_config.name}",
-            f"train_{metric_config.name}_se",
         ]
 
     def test_total_scores_values_are_floats(self, logged_scores):
diff --git a/tests/test_sequence_classification.py b/tests/test_sequence_classification.py
deleted file mode 100644
index 165866898..000000000
--- a/tests/test_sequence_classification.py
+++ /dev/null
@@ -1,60 +0,0 @@
-"""Unit tests for the `sequence_classification` module."""
-
-import os
-from contextlib import nullcontext as does_not_raise
-from typing import Generator
-
-import pytest
-
-from scandeval.benchmark_dataset import BenchmarkDataset
-from scandeval.dataset_configs import get_all_dataset_configs
-from scandeval.exceptions import InvalidBenchmark
-from scandeval.languages import DA
-from scandeval.sequence_classification import SequenceClassification
-from scandeval.tasks import COMMON_SENSE, KNOW, SENT
-from scandeval.utils import GENERATIVE_DATASET_TASKS
-
-
-@pytest.fixture(
-    scope="module",
-    params=[
-        dataset_config
-        for dataset_config in get_all_dataset_configs().values()
-        if dataset_config.task in [SENT, KNOW, COMMON_SENSE]
-        and (
-            os.getenv("TEST_ALL_DATASETS", "0") == "1"
-            or (not dataset_config.unofficial and dataset_config.languages == [DA])
-        )
-    ],
-    ids=lambda dataset_config: dataset_config.name,
-)
-def benchmark_dataset(
-    benchmark_config, request
-) -> Generator[BenchmarkDataset, None, None]:
-    """Yields a sequence classification benchmark dataset."""
-    yield SequenceClassification(
-        dataset_config=request.param, benchmark_config=benchmark_config
-    )
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_encoder_benchmarking(benchmark_dataset, model_id):
-    """Test that the encoder can be benchmarked on sequence classification datasets."""
-    if benchmark_dataset.dataset_config.task.name in GENERATIVE_DATASET_TASKS:
-        with pytest.raises(InvalidBenchmark):
-            benchmark_dataset.benchmark(model_id)
-    else:
-        with does_not_raise():
-            benchmark_dataset.benchmark(model_id)
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_decoder_sequence_benchmarking(
-    benchmark_dataset, generative_model_id, generative_model_and_tokenizer
-):
-    """Test that the decoder can be benchmarked on sequence classification datasets."""
-    model, tokenizer = generative_model_and_tokenizer
-    with does_not_raise():
-        benchmark_dataset.benchmark(
-            model_id=generative_model_id, model=model, tokenizer=tokenizer
-        )
diff --git a/tests/test_speed_benchmark.py b/tests/test_speed_benchmark.py
index 13b1b187e..8211a6c62 100644
--- a/tests/test_speed_benchmark.py
+++ b/tests/test_speed_benchmark.py
@@ -4,31 +4,24 @@
 
 import pytest
 from tqdm.auto import tqdm
-from transformers import PreTrainedModel, PreTrainedTokenizer
-from transformers.models.auto import AutoModelForSequenceClassification, AutoTokenizer
 
-from scandeval.config import ModelConfig
-from scandeval.model_setups import HFModelSetup
+from scandeval.benchmark_modules.base import BenchmarkModule
+from scandeval.benchmark_modules.hf import HuggingFaceEncoderModel
+from scandeval.dataset_configs import SPEED_CONFIG
+from scandeval.model_config import get_model_config
 from scandeval.speed_benchmark import benchmark_speed
 
 
 @pytest.fixture(scope="module")
-def tokenizer(model_id) -> Generator[PreTrainedTokenizer, None, None]:
-    """Yields a tokenizer."""
-    yield AutoTokenizer.from_pretrained(model_id)
-
-
-@pytest.fixture(scope="module")
-def model(model_id) -> Generator[PreTrainedModel, None, None]:
+def model(encoder_model_id, benchmark_config) -> Generator[BenchmarkModule, None, None]:
     """Yields a model."""
-    yield AutoModelForSequenceClassification.from_pretrained(model_id)
-
-
-@pytest.fixture(scope="module")
-def model_config(model_id, benchmark_config) -> Generator[ModelConfig, None, None]:
-    """Yields a model configuration."""
-    model_setup = HFModelSetup(benchmark_config=benchmark_config)
-    yield model_setup.get_model_config(model_id=model_id)
+    yield HuggingFaceEncoderModel(
+        model_config=get_model_config(
+            model_id=encoder_model_id, benchmark_config=benchmark_config
+        ),
+        dataset_config=SPEED_CONFIG,
+        benchmark_config=benchmark_config,
+    )
 
 
 class TestBenchmarkSpeed:
@@ -41,43 +34,25 @@ def itr(self) -> Generator[tqdm, None, None]:
 
     @pytest.fixture(scope="class")
     def scores(
-        self, itr, tokenizer, model, model_config, dataset_config, benchmark_config
-    ) -> Generator[dict[str, list[dict[str, float]]], None, None]:
+        self, model, benchmark_config
+    ) -> Generator[list[dict[str, float]], None, None]:
         """Yields the benchmark speed scores."""
-        yield benchmark_speed(
-            itr=itr,
-            tokenizer=tokenizer,
-            model=model,
-            model_config=model_config,
-            dataset_config=dataset_config,
-            benchmark_config=benchmark_config,
-        )
-
-    def test_scores_is_dict(self, scores):
-        """Tests that the scores are a dict."""
-        assert isinstance(scores, dict)
-
-    def test_scores_keys(self, scores):
-        """Tests that the scores have the correct keys."""
-        assert set(scores.keys()) == {"test"}
+        yield benchmark_speed(model=model, benchmark_config=benchmark_config)
 
-    def test_test_scores_is_list(self, scores):
-        """Tests that the test scores are a list."""
-        assert isinstance(scores["test"], list)
+    def test_scores_is_list(self, scores):
+        """Tests that the scores is a list."""
+        assert isinstance(scores, list)
 
-    def test_test_scores_contain_dicts(self, scores):
-        """Tests that the test scores contain dicts."""
-        assert all(isinstance(x, dict) for x in scores["test"])
+    def test_scores_contain_dicts(self, scores):
+        """Tests that the scores contain dicts."""
+        assert all(isinstance(x, dict) for x in scores)
 
-    def test_test_scores_dicts_keys(self, scores):
-        """Tests that the test scores dicts have the correct keys."""
-        assert all(
-            set(x.keys()) == {"test_speed", "test_speed_short"} for x in scores["test"]
-        )
+    def test_scores_dicts_keys(self, scores):
+        """Tests that the scores dicts have the correct keys."""
+        assert all(set(x.keys()) == {"test_speed", "test_speed_short"} for x in scores)
 
-    def test_test_scores_dicts_values_dtypes(self, scores):
-        """Tests that the test scores dicts have the correct values dtypes."""
+    def test_scores_dicts_values_dtypes(self, scores):
+        """Tests that the scores dicts have the correct values dtypes."""
         assert all(
-            all(isinstance(value, float) for value in x.values())
-            for x in scores["test"]
+            all(isinstance(value, float) for value in x.values()) for x in scores
         )
diff --git a/tests/test_structured_generation_utils.py b/tests/test_structured_generation_utils.py
index d77513402..267748449 100644
--- a/tests/test_structured_generation_utils.py
+++ b/tests/test_structured_generation_utils.py
@@ -1,79 +1,61 @@
 """Unit tests for the `structured_generation_utils` module."""
 
-import json
-from typing import TYPE_CHECKING, Generator
-
 import pytest
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    GenerationConfig,
-    PreTrainedModel,
-    PreTrainedTokenizerBase,
-    StoppingCriteriaList,
-)
-
-from scandeval.generation import StopWordCriteria
-from scandeval.structured_generation_utils import get_ner_prefix_allowed_tokens_fn
-from scandeval.utils import create_model_cache_dir
-
-if TYPE_CHECKING:
-    from outlines.integrations.transformers import JSONPrefixAllowedTokens
-
-
-@pytest.fixture(scope="module")
-def tokenizer(generative_model_id) -> Generator[PreTrainedTokenizerBase, None, None]:
-    """A tokenizer for a generative model."""
-    yield AutoTokenizer.from_pretrained(generative_model_id)
 
-
-@pytest.fixture(scope="module")
-def model(
-    generative_model_id, benchmark_config
-) -> Generator[PreTrainedModel, None, None]:
-    """A generative model."""
-    model_cache_dir = create_model_cache_dir(
-        cache_dir=benchmark_config.cache_dir, model_id=generative_model_id
-    )
-    yield AutoModelForCausalLM.from_pretrained(
-        generative_model_id, cache_dir=model_cache_dir, torch_dtype=None
-    )
-
-
-@pytest.fixture(scope="module")
-def ner_prefix_allowed_tokens_fn(
-    tokenizer,
-) -> Generator["JSONPrefixAllowedTokens", None, None]:
-    """A prefix_allowed_tokens_fn for named entity recognition."""
-    yield get_ner_prefix_allowed_tokens_fn(
-        ner_tag_names=["person", "location"], tokenizer=tokenizer
-    )
+from scandeval.structured_generation_utils import get_ner_schema
 
 
 @pytest.mark.parametrize(
-    argnames="input_str",
+    argnames=["ner_tag_names", "expected"],
     argvalues=[
-        "John works at Novo Nordisk.",
-        "Peter lives in Copenhagen.",
-        "He loves to surf.",
+        (list(), dict(properties=dict(), title="AnswerFormat", type="object")),
+        (
+            ["person"],
+            dict(
+                properties=dict(
+                    person=dict(
+                        items=dict(type="string"),
+                        maxItems=5,
+                        title="Person",
+                        type="array",
+                    )
+                ),
+                required=["person"],
+                title="AnswerFormat",
+                type="object",
+            ),
+        ),
+        (
+            ["person", "location"],
+            dict(
+                properties=dict(
+                    location=dict(
+                        items=dict(type="string"),
+                        maxItems=5,
+                        title="Location",
+                        type="array",
+                    ),
+                    person=dict(
+                        items=dict(type="string"),
+                        maxItems=5,
+                        title="Person",
+                        type="array",
+                    ),
+                ),
+                required=["person", "location"],
+                title="AnswerFormat",
+                type="object",
+            ),
+        ),
     ],
-    ids=["person-no-location", "person-location", "no-person-no-location"],
 )
-def test_prefix_allowed_tokens_fn(
-    tokenizer, model, ner_prefix_allowed_tokens_fn, input_str
-):
-    """Test the `prefix_allowed_tokens_fn` for named entity recognition."""
-    input_ids = tokenizer(input_str, return_tensors="pt").input_ids
-    stopping_criteria = StopWordCriteria(
-        stop_word_id_lists=[tokenizer.encode("}"), [tokenizer.eos_token_id]]
-    )
-    completion = model.generate(
-        input_ids,
-        generation_config=GenerationConfig(do_sample=False, max_new_tokens=32),
-        stopping_criteria=StoppingCriteriaList([stopping_criteria]),
-        prefix_allowed_tokens_fn=ner_prefix_allowed_tokens_fn,
-    )[0][input_ids.shape[-1] :]
-    completion_str = tokenizer.decode(completion, skip_special_tokens=True)
-    parsed_output = json.loads(completion_str)
-    assert isinstance(parsed_output, dict)
-    assert set(list(parsed_output.keys())) == {"person", "location"}
+def test_get_ner_schema(ner_tag_names, expected):
+    """Test that the NER schema can be retrieved."""
+    schema = get_ner_schema(ner_tag_names=ner_tag_names).model_json_schema()
+    assert schema == expected
+
+
+# TODO
+def test_get_ner_logits_processors():
+    """Test that the NER logits processors can be retrieved."""
+    pass
diff --git a/tests/test_task_utils/__init__.py b/tests/test_task_utils/__init__.py
new file mode 100644
index 000000000..03fbd0584
--- /dev/null
+++ b/tests/test_task_utils/__init__.py
@@ -0,0 +1,1 @@
+"""Tests for the `task_utils` subpackage."""
diff --git a/tests/test_task_utils/test_question_answering.py b/tests/test_task_utils/test_question_answering.py
new file mode 100644
index 000000000..3a554315e
--- /dev/null
+++ b/tests/test_task_utils/test_question_answering.py
@@ -0,0 +1,1 @@
+"""Tests for the `question_answering` module."""
diff --git a/tests/test_task_utils/test_sequence_classification.py b/tests/test_task_utils/test_sequence_classification.py
new file mode 100644
index 000000000..efc73961a
--- /dev/null
+++ b/tests/test_task_utils/test_sequence_classification.py
@@ -0,0 +1,1 @@
+"""Tests for the `sequence_classification` module."""
diff --git a/tests/test_task_utils/test_text_to_text.py b/tests/test_task_utils/test_text_to_text.py
new file mode 100644
index 000000000..11f93f5e7
--- /dev/null
+++ b/tests/test_task_utils/test_text_to_text.py
@@ -0,0 +1,1 @@
+"""Tests for the `text_to_text` module."""
diff --git a/tests/test_task_utils/test_token_classification.py b/tests/test_task_utils/test_token_classification.py
new file mode 100644
index 000000000..2ca9c5417
--- /dev/null
+++ b/tests/test_task_utils/test_token_classification.py
@@ -0,0 +1,1 @@
+"""Tests for the `token_classification` module."""
diff --git a/tests/test_dataset_tasks.py b/tests/test_tasks.py
similarity index 97%
rename from tests/test_dataset_tasks.py
rename to tests/test_tasks.py
index 4c0098bc3..61f44e950 100644
--- a/tests/test_dataset_tasks.py
+++ b/tests/test_tasks.py
@@ -4,7 +4,7 @@
 
 import pytest
 
-from scandeval.config import Task
+from scandeval.data_models import Task
 from scandeval.tasks import get_all_tasks
 
 
diff --git a/tests/test_text_to_text.py b/tests/test_text_to_text.py
deleted file mode 100644
index a98b83db6..000000000
--- a/tests/test_text_to_text.py
+++ /dev/null
@@ -1,45 +0,0 @@
-"""Unit tests for the `text_to_text` module."""
-
-import os
-from contextlib import nullcontext as does_not_raise
-from typing import Generator
-
-import pytest
-
-from scandeval.benchmark_dataset import BenchmarkDataset
-from scandeval.dataset_configs import get_all_dataset_configs
-from scandeval.languages import DA
-from scandeval.tasks import SUMM
-from scandeval.text_to_text import TextToText
-
-
-@pytest.fixture(
-    scope="module",
-    params=[
-        dataset_config
-        for dataset_config in get_all_dataset_configs().values()
-        if dataset_config.task == SUMM
-        and (
-            os.getenv("TEST_ALL_DATASETS", "0") == "1"
-            or (not dataset_config.unofficial and dataset_config.languages == [DA])
-        )
-    ],
-    ids=lambda dataset_config: dataset_config.name,
-)
-def benchmark_dataset(
-    benchmark_config, request
-) -> Generator[BenchmarkDataset, None, None]:
-    """Yields a text-to-text benchmark dataset."""
-    yield TextToText(dataset_config=request.param, benchmark_config=benchmark_config)
-
-
-@pytest.mark.skipif(condition=os.getenv("TEST_EVALUATIONS") == "0", reason="Skipped")
-def test_decoder_benchmarking(
-    benchmark_dataset, generative_model_id, generative_model_and_tokenizer
-):
-    """Test that decoder models can be benchmarked on text-to-text tasks."""
-    model, tokenizer = generative_model_and_tokenizer
-    with does_not_raise():
-        benchmark_dataset.benchmark(
-            model_id=generative_model_id, model=model, tokenizer=tokenizer
-        )
diff --git a/tests/test_types.py b/tests/test_types.py
index 9a2c8baa8..032ba5083 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -1,26 +1,50 @@
 """Unit tests for the `types` module."""
 
-import re
-import typing
-from typing import Generator
-
 import pytest
 
-import scandeval.types as types
+from scandeval.types import is_list_of_int, is_list_of_list_of_int, is_list_of_str
+
+
+@pytest.mark.parametrize(
+    argnames=["obj", "expected"],
+    argvalues=[
+        ([1, 2, 3], True),
+        ([1.0, 2.0, 3.0], False),
+        (["a", "b", "c"], False),
+        (["a", 1, 2.0], False),
+        ([], True),
+    ],
+)
+def test_is_list_of_int(obj, expected):
+    """Test the `is_list_of_int` function."""
+    assert is_list_of_int(x=obj) == expected
 
 
-@pytest.fixture(scope="module")
-def module_variable_names() -> Generator[list[str], None, None]:
-    """Yields the module variable names."""
-    yield [
-        var
-        for var in dir(types)
-        if "_" not in var and not hasattr(typing, var) and var != "np"
-    ]
+@pytest.mark.parametrize(
+    argnames=["obj", "expected"],
+    argvalues=[
+        ([1, 2, 3], False),
+        ([[1, 2, 3], [4, 5, 6]], True),
+        ([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], False),
+        ([["a", "b", "c"], ["d", "e", "f"]], False),
+        ([["a", 1, 2.0], [3, 4, 5]], False),
+        ([[], []], True),
+    ],
+)
+def test_is_list_of_list_of_int(obj, expected):
+    """Test the `is_list_of_int` function."""
+    assert is_list_of_list_of_int(x=obj) == expected
 
 
-def test_type_variable_names_are_title_case(module_variable_names) -> None:
-    """Tests that all type variable names are title case."""
-    for var in module_variable_names:
-        assert var[0].isupper()
-        assert re.search(r"[^A-Za-z]", var) is None
+@pytest.mark.parametrize(
+    argnames=["obj", "expected"],
+    argvalues=[
+        (["a", "b", "c"], True),
+        (["a", 1, 2.0], False),
+        ([1, 2, 3], False),
+        ([], True),
+    ],
+)
+def test_is_list_of_str(obj, expected):
+    """Test the `is_list_of_str` function."""
+    assert is_list_of_str(x=obj) == expected
diff --git a/tests/test_vllm_models.py b/tests/test_vllm_models.py
deleted file mode 100644
index 03ee2b578..000000000
--- a/tests/test_vllm_models.py
+++ /dev/null
@@ -1,1 +0,0 @@
-"""Unit tests for the `vllms_models` module."""