euroeval.benchmark_modules.hf

docs module euroeval.benchmark_modules.hf

"""Encoder models from the Hugging Face Hub."""

import collections.abc as c
import logging
import os
import typing as t
from functools import cached_property, partial
from json import JSONDecodeError
from pathlib import Path
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 ModelInfo as HfApiModelInfo
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,
    BatchEncoding,
    DataCollatorForTokenClassification,
    DataCollatorWithPadding,
    PretrainedConfig,
    PreTrainedModel,
    PreTrainedTokenizer,
    Trainer,
)
from transformers.modelcard import TASK_MAPPING
from urllib3.exceptions import RequestError

from ..constants import (
    DUMMY_FILL_VALUE,
    GENERATIVE_PIPELINE_TAGS,
    LOCAL_MODELS_REQUIRED_FILES,
    MERGE_TAGS,
)
from ..data_models import BenchmarkConfig, DatasetConfig, HFModelInfo, ModelConfig, Task
from ..enums import (
    BatchingPreference,
    GenerativeType,
    InferenceBackend,
    ModelType,
    TaskGroup,
)
from ..exceptions import (
    HuggingFaceHubDown,
    InvalidBenchmark,
    InvalidModel,
    NeedsAdditionalArgument,
    NeedsEnvironmentVariable,
    NeedsExtraInstalled,
    NoInternetConnection,
)
from ..languages import get_all_languages
from ..task_utils import (
    multiple_choice_classification,
    question_answering,
    token_classification,
)
from ..types import ExtractLabelsFunction
from ..utils import (
    block_terminal_output,
    create_model_cache_dir,
    get_bos_token,
    get_class_by_name,
    get_eos_token,
    internet_connection_available,
)
from .base import BenchmarkModule

logger = logging.getLogger("euroeval")


class HuggingFaceEncoderModel(BenchmarkModule):
    """An encoder model from the Hugging Face Hub."""

    fresh_model = False
    batching_preference = BatchingPreference.NO_PREFERENCE
    high_priority = True

    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.
        """
        model, tokenizer = load_model_and_tokenizer(
            model_config=model_config,
            dataset_config=dataset_config,
            benchmark_config=benchmark_config,
        )
        self._model: PreTrainedModel = model
        self._tokenizer: PreTrainedTokenizer = tokenizer

        self._model, self._tokenizer = align_model_and_tokenizer(
            model=self._model,
            tokenizer=self._tokenizer,
            model_max_length=self.model_max_length,
            raise_errors=benchmark_config.raise_errors,
        )

        super().__init__(
            model_config=model_config,
            dataset_config=dataset_config,
            benchmark_config=benchmark_config,
        )

    @cached_property
    def num_params(self) -> int:
        """The number of parameters in the model.

        Returns:
            The number of parameters in the model.
        """
        token = (
            self.benchmark_config.api_key or os.getenv("HUGGINGFACE_API_KEY") or True
        )
        hf_api = HfApi(token=token)
        try:
            repo_info = hf_api.model_info(
                repo_id=self.model_config.adapter_base_model_id
                or self.model_config.model_id,
                revision=self.model_config.revision,
            )
        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 (
            hasattr(self._model.config, "num_params")
            and self._model.config.num_params is not None
        ):
            num_params = self._model.config.num_params
        elif hasattr(self._model, "parameters"):
            num_params = sum(p.numel() for p in self._model.parameters())
        else:
            logger.warning(
                "The number of parameters could not be determined for the model, since "
                "the model is not stored in the safetensors format. If this is your "
                "own model, then you can use this Hugging Face Space to convert your "
                "model to the safetensors format: "
                "https://huggingface.co/spaces/safetensors/convert."
            )
            num_params = -1
        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",
            "max_sequence_length",
            "model_max_length",
            "sliding_window",
            "sliding_window_size",
            "n_positions",
        ]
        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

    @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.task_group:
            case (
                TaskGroup.SEQUENCE_CLASSIFICATION
                | TaskGroup.TEXT_TO_TEXT
                | TaskGroup.QUESTION_ANSWERING
                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
            ):
                return DataCollatorWithPadding(self._tokenizer, padding="longest")
            case TaskGroup.TOKEN_CLASSIFICATION:
                return DataCollatorForTokenClassification(
                    tokenizer=self._tokenizer, label_pad_token_id=-100
                )
            case _:
                raise NotImplementedError(
                    f"Unsupported task group: {self.dataset_config.task.task_group}."
                )

    @property
    def generative_type(self) -> GenerativeType | None:
        """Get the generative type of the model.

        Returns:
            The generative type of the model, or None if it has not been set yet.
        """
        return None

    @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."
        )

    @property
    def trainer_class(self) -> t.Type["Trainer"]:
        """The Trainer class to use for finetuning.

        Returns:
            The Trainer class.
        """
        match self.dataset_config.task.task_group:
            case (
                TaskGroup.SEQUENCE_CLASSIFICATION
                | TaskGroup.TEXT_TO_TEXT
                | TaskGroup.TOKEN_CLASSIFICATION
            ):
                return Trainer
            case TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION:
                return (
                    multiple_choice_classification.MultipleChoiceClassificationTrainer
                )
            case TaskGroup.QUESTION_ANSWERING:
                return question_answering.QuestionAnsweringTrainer
            case _:
                raise NotImplementedError(
                    f"Unsupported task group: {self.dataset_config.task.task_group}."
                )

    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) -> 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) -> BatchEncoding:
            return self._tokenizer(text=examples["text"], truncation=True, padding=True)

        match task.task_group:
            case TaskGroup.SEQUENCE_CLASSIFICATION:
                dataset = dataset.map(
                    numericalise_labels, batched=True, load_from_cache_file=False
                ).map(tokenise, batched=True, load_from_cache_file=False)

            case TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION:
                dataset = DatasetDict(
                    train=dataset["train"].map(
                        partial(
                            multiple_choice_classification.prepare_examples,
                            tokenizer=self._tokenizer,
                        ),
                        batched=True,
                        batch_size=1,
                        remove_columns=dataset["train"].column_names,
                        load_from_cache_file=False,
                        keep_in_memory=True,
                    ),
                    val=dataset["val"].map(
                        partial(
                            multiple_choice_classification.prepare_examples,
                            tokenizer=self._tokenizer,
                        ),
                        batched=True,
                        batch_size=1,
                        remove_columns=dataset["val"].column_names,
                        load_from_cache_file=False,
                        keep_in_memory=True,
                    ),
                    test=dataset["test"].map(
                        partial(
                            multiple_choice_classification.prepare_examples,
                            tokenizer=self._tokenizer,
                        ),
                        batched=True,
                        batch_size=1,
                        remove_columns=dataset["test"].column_names,
                        load_from_cache_file=False,
                        keep_in_memory=True,
                    ),
                )

            case TaskGroup.TEXT_TO_TEXT:
                dataset = dataset.map(
                    tokenise,
                    batched=True,
                    load_from_cache_file=False,
                    keep_in_memory=True,
                )

            case TaskGroup.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 TaskGroup.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 group: {task.task_group}.")

        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 not in GENERATIVE_PIPELINE_TAGS
        )

    @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.")

        language_mapping = get_all_languages()
        language_codes = list(language_mapping.keys())

        model_config = ModelConfig(
            model_id=model_id,
            revision=revision,
            task=model_info.pipeline_tag,
            languages=[
                language_mapping[tag]
                for tag in model_info.tags
                if tag in language_codes
            ],
            merge=any(tag in model_info.tags for tag in MERGE_TAGS),
            inference_backend=InferenceBackend.TRANSFORMERS,
            model_type=ModelType.ENCODER,
            fresh=False,
            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(
docs
    model_config: ModelConfig,
    dataset_config: DatasetConfig,
    benchmark_config: BenchmarkConfig,
) -> tuple[PreTrainedModel, PreTrainedTokenizer]:
    """Load the model and tokenizer.

    Args:
        model_config:
            The model configuration.
        dataset_config:
            The dataset configuration.
        benchmark_config:
            The benchmark configuration

    Returns:
        The loaded model and tokenizer.
    """
    config: "PretrainedConfig"
    block_terminal_output()

    model_id = model_config.model_id
    task_group = dataset_config.task.task_group
    ignore_mismatched_sizes = False

    # Special case where there is a mismatch between the labels during training and
    # testing
    if dataset_config.task.task_group == TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION:
        id2label = {0: "0", 1: "1"}
    else:
        id2label = dataset_config.id2label

    config = load_hf_model_config(
        model_id=model_id,
        num_labels=len(id2label),
        id2label=id2label,
        label2id={label: idx for idx, label in id2label.items()},
        revision=model_config.revision,
        model_cache_dir=model_config.model_cache_dir,
        api_key=benchmark_config.api_key,
        trust_remote_code=benchmark_config.trust_remote_code,
        run_with_cli=benchmark_config.run_with_cli,
    )

    model_kwargs = dict(
        config=config,
        ignore_mismatched_sizes=ignore_mismatched_sizes,
        revision=model_config.revision,
        token=benchmark_config.api_key or os.getenv("HUGGINGFACE_API_KEY") or True,
        cache_dir=model_config.model_cache_dir,
        trust_remote_code=benchmark_config.trust_remote_code,
        torch_dtype=get_torch_dtype(
            device=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:
        # Get the model class associated with the task group
        model_cls_or_none: t.Type["PreTrainedModel"] | None = get_class_by_name(
            class_name=task_group_to_class_name(task_group=task_group),
            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 task group {task_group.value!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
            )
            break
        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 (OSError, ValueError) as e:
            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}."
            )

    if isinstance(model_or_tuple, tuple):
        model = model_or_tuple[0]
    else:
        model = model_or_tuple

    assert model is not None, "The model should not be None."

    model.eval()
    model.to(benchmark_config.device)

    if (
        isinstance(model, PreTrainedModel)
        and task_group == TaskGroup.QUESTION_ANSWERING
    ):
        model = setup_model_for_question_answering(model=model)

    tokenizer = load_tokenizer(
        model=model,
        model_id=model_id,
        trust_remote_code=benchmark_config.trust_remote_code,
    )

    return model, tokenizer


def get_model_repo_info(
docs
    model_id: str, revision: str, benchmark_config: BenchmarkConfig
) -> HFModelInfo | None:
    """Get the information about the model from the HF Hub or a local directory.

    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.
    """
    token = benchmark_config.api_key or os.getenv("HUGGINGFACE_API_KEY") or True
    hf_api = HfApi(token=token)
    model_id, revision = model_id.split("@") if "@" in model_id else (model_id, "main")

    # Get information on the model.
    # The first case is when the model is a local model, in which case we create a dummy
    # model info object.
    model_info: HfApiModelInfo | None = None
    if Path(model_id).is_dir():
        logger.debug(f"Checking for local model in {model_id}.")
        if all(
            (Path(model_id) / required_file).exists()
            for required_file in LOCAL_MODELS_REQUIRED_FILES
        ):
            model_info = HfApiModelInfo(id=model_id, tags=None, pipeline_tag=None)

    # If the model does not exist locally, then we get the model info from the Hugging
    # Face Hub
    if model_info is None:
        try:
            model_info = hf_api.model_info(
                repo_id=model_id, revision=revision, token=token
            )
        except (GatedRepoError, LocalTokenNotFoundError) as e:
            try:
                hf_whoami(token=token)
                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,
                )
        except (RepositoryNotFoundError, HFValidationError):
            return None
        except (OSError, RequestException):
            if internet_connection_available():
                raise HuggingFaceHubDown()
            else:
                raise NoInternetConnection()

    # Get all the Hugging Face repository tags for the model. If the model is an adapter
    # model, then we also get the tags for the base model
    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:
        has_adapter_config = model_info.siblings is not None and any(
            sibling.rfilename == "adapter_config.json"
            for sibling in model_info.siblings
        )
        if has_adapter_config:
            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 os.getenv("HUGGINGFACE_API_KEY")
                or True,
            )
            tags += base_model_info.tags or list()
            tags = list(set(tags))

    # Get the pipeline tag for the model. If it is not specified, then we determine it
    # by checking the model's architecture as written in the model's Hugging Face config
    pipeline_tag = model_info.pipeline_tag
    if pipeline_tag is None:
        hf_config = load_hf_model_config(
            model_id=model_id,
            num_labels=0,
            id2label=dict(),
            label2id=dict(),
            revision=revision,
            model_cache_dir=create_model_cache_dir(
                cache_dir=benchmark_config.cache_dir, model_id=model_id
            ),
            api_key=benchmark_config.api_key,
            trust_remote_code=benchmark_config.trust_remote_code,
            run_with_cli=benchmark_config.run_with_cli,
        )
        class_names = hf_config.architectures
        generative_class_names = [
            class_name
            for tag in GENERATIVE_PIPELINE_TAGS
            for class_name in TASK_MAPPING[tag].values()
        ]
        if class_names is not None and any(
            class_name in generative_class_names for class_name in class_names
        ):
            pipeline_tag = "text-generation"
        else:
            pipeline_tag = "fill-mask"

    if benchmark_config.only_allow_safetensors:
        # Check if any file ends with .safetensors
        repo_files = hf_api.list_repo_files(repo_id=model_id, revision=revision)
        has_safetensors = any(f.endswith(".safetensors") for f in repo_files)
        if not has_safetensors:
            msg = f"Model {model_id} does not have safetensors weights available. "
            if benchmark_config.run_with_cli:
                msg += "Skipping since the `--only-allow-safetensors` flag is set."
            else:
                msg += (
                    "Skipping since the `only_allow_safetensors` argument is set "
                    "to `True`."
                )
            raise InvalidModel(msg)

    return HFModelInfo(
        pipeline_tag=pipeline_tag, tags=tags, adapter_base_model_id=base_model_id
    )


def load_tokenizer(
docs
    model: "PreTrainedModel | None", 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. Can be None.
        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.
    if model is not None:
        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

    num_retries = 5
    for _ in range(num_retries):
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_id, **loading_kwargs)
            break
        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
    else:
        raise InvalidModel(
            f"Could not load tokenizer for model {model_id!r} after {num_retries} "
            "attempts."
        )

    # Ensure that BOS, EOS and PAD tokens are set
    tokenizer.bos_token, tokenizer.bos_token_id = get_bos_token(tokenizer=tokenizer)
    tokenizer.eos_token, tokenizer.eos_token_id = get_eos_token(tokenizer=tokenizer)

    return tokenizer


def get_torch_dtype(
docs
    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(
docs
    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 os.getenv("HUGGINGFACE_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, GatedRepoError) as e:
            # TEMP: When the model is gated then we cannot set cache dir, for some
            # reason (since transformers v4.38.2, still a problem in v4.48.0). 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:
            if "awaiting a review from the repo authors" in str(e):
                raise InvalidModel(
                    f"The model {model_id!r} is awaiting a review from the repository "
                    "authors. Please try again later."
                )
            if "trust_remote_code" in str(e):
                raise NeedsAdditionalArgument(
                    cli_argument="--trust-remote-code",
                    script_argument="trust_remote_code=True",
                    run_with_cli=run_with_cli,
                )
            raise InvalidModel(
                f"The config for the model {model_id!r} could not be loaded. The "
                f"error was {e!r}."
            )

docs
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(
docs
    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(
docs
    model: "PreTrainedModel",
    tokenizer: "PreTrainedTokenizer",
    model_max_length: int,
    raise_errors: bool = False,
) -> tuple["PreTrainedModel", "PreTrainedTokenizer"]:
    """Aligns the model and the tokenizer.

    Args:
        model:
            The model to fix.
        tokenizer:
            The tokenizer to fix.
        model_max_length:
            The maximum length of the model.
        raise_errors:
            Whether to raise errors instead of trying to fix them silently.

    Returns:
        The fixed model and tokenizer.
    """
    # Ensure that the model max length is at most 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

    # Move the model to the CPU, since otherwise we can't catch the IndexErrors when
    # finding the maximum sequence length of the model
    model_device = model.device
    model.to(torch.device("cpu"))

    # 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=model.device,
        )
        with torch.inference_mode():
            try:
                model(dummy_inputs, attention_mask=torch.ones_like(dummy_inputs))
                break

            # This happens if `max_length` is too large
            except IndexError:
                continue

    # Move the model back to the original device
    model.to(model_device)

    # 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"):
        if model.config.vocab_size < len(tokenizer):
            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

    return model, tokenizer


def task_group_to_class_name(task_group: TaskGroup) -> str:
docs
    """Convert a task group to a class name.

    Args:
        task_group:
            The task group.

    Returns:
        The class name.
    """
    pascal_case = task_group.title().replace("_", "")
    special_case_mapping = dict(
        MultipleChoiceClassification="SequenceClassification",
        Speed="SequenceClassification",
    )
    pascal_case = special_case_mapping.get(pascal_case, pascal_case)
    return f"AutoModelFor{pascal_case}"