EuroEval 15.12.0__py3-none-any.whl → 16.7.1__py3-none-any.whl

euroeval/tokenisation_utils.py

@@ -0,0 +1,603 @@
+"""Utility functions related to tokenisation."""
+
+import collections.abc as c
+import logging
+import re
+import typing as t
+
+import torch
+from transformers.tokenization_mistral_common import MistralCommonTokenizer
+
+from .constants import BOS_TOKENS, EOS_TOKENS, PAD_TOKENS
+from .enums import GenerativeType
+from .exceptions import InvalidModel
+from .logging_utils import log, log_once
+
+if t.TYPE_CHECKING:
+    from transformers.tokenization_utils import PreTrainedTokenizer
+    from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+
+    from .data_models import DatasetConfig, ModelConfig
+
+
+def get_special_token_metadata(tokeniser: "PreTrainedTokenizerBase") -> dict:
+    """Get the special token metadata for a tokeniser.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+
+    Returns:
+        The special token metadata.
+    """
+    # Create some test input IDs, to check if the tokeniser is adding special tokens
+    test_input_ids = tokeniser("Test").input_ids
+
+    # Extract the CLS token IDs from the tokeniser, if it's using them
+    has_cls_token = True
+    if tokeniser.cls_token_id in test_input_ids:
+        cls_token_id = tokeniser.cls_token_id
+        cls_token = tokeniser.cls_token
+    elif tokeniser.bos_token_id in test_input_ids:
+        cls_token_id = tokeniser.bos_token_id
+        cls_token = tokeniser.bos_token
+    elif tokeniser.cls_token is not None:
+        cls_token_id = tokeniser.cls_token_id
+        cls_token = tokeniser.cls_token
+        has_cls_token = False
+    else:
+        cls_token_id = tokeniser.bos_token_id
+        cls_token = tokeniser.bos_token
+        has_cls_token = False
+
+    # Extract the SEP token IDs from the tokeniser, if it's using them
+    has_sep_token = True
+    if tokeniser.sep_token_id in test_input_ids:
+        sep_token = tokeniser.sep_token
+    elif tokeniser.eos_token_id in test_input_ids:
+        sep_token = tokeniser.eos_token
+    elif tokeniser.sep_token is not None:
+        sep_token = tokeniser.sep_token
+        has_sep_token = False
+    else:
+        sep_token = tokeniser.eos_token
+        has_sep_token = False
+
+    return dict(
+        cls_token_id=cls_token_id,
+        cls_token=cls_token,
+        sep_token=sep_token,
+        has_cls_token=has_cls_token,
+        has_sep_token=has_sep_token,
+    )
+
+
+def should_prompts_be_stripped(
+    labels_to_be_generated: c.Sequence[str], tokeniser: "PreTrainedTokenizer"
+) -> bool:
+    """Determine if we should strip the prompts for few-shot evaluation.
+
+    This is the case if the tokeniser needs to include the space as part of the label
+    token. The strategy is thus to tokenise a label with a preceeding colon (as in the
+    prompts), i.e., ": positive", and check if the tokenisation starts with the tokens
+    of ": ". If this is the case, then we should not strip the prompts, since the
+    tokeniser produces the whitespace token separately.
+
+    Args:
+        labels_to_be_generated:
+            The labels that are to be generated.
+        tokeniser:
+            The tokeniser used to tokenise the labels.
+
+    Returns:
+        Whether we should strip the prompts.
+    """
+    strip_prompts = True
+    for label in labels_to_be_generated:
+        colon_tokens = tokeniser(": ", add_special_tokens=False).input_ids
+        label_tokens = tokeniser(": " + label, add_special_tokens=False).input_ids
+
+        if isinstance(colon_tokens, torch.Tensor):
+            colon_tokens = list(colon_tokens.squeeze(0))
+        if isinstance(label_tokens, torch.Tensor):
+            label_tokens = list(label_tokens.squeeze(0))
+
+        label_tokens_start_with_colon_tokens = (
+            label_tokens[: len(colon_tokens)] == colon_tokens
+        )
+        if label_tokens_start_with_colon_tokens:
+            strip_prompts = False
+
+    return strip_prompts
+
+
+def should_prefix_space_be_added_to_labels(
+    labels_to_be_generated: c.Sequence[str], tokeniser: "PreTrainedTokenizer"
+) -> bool:
+    """Determine if we should add a prefix space to the labels.
+
+    This is the case if the prompts are stripped and the tokeniser doesn't
+    automatically add prefix whitespaces to the labels.
+
+    Args:
+        labels_to_be_generated:
+            The labels that are to be generated.
+        tokeniser:
+            The tokeniser used to tokenise the labels.
+
+    Returns:
+        Whether we should add a prefix space to the labels.
+    """
+    if not should_prompts_be_stripped(
+        labels_to_be_generated=labels_to_be_generated, tokeniser=tokeniser
+    ):
+        return False
+
+    whitespace_token = tokeniser.convert_ids_to_tokens(
+        ids=tokeniser(" ", add_special_tokens=False).input_ids[0]
+    )[0]
+
+    add_prefix_space = True
+    for label in labels_to_be_generated:
+        label_tokens = tokeniser(label, add_special_tokens=False).input_ids
+        if isinstance(label_tokens, torch.Tensor):
+            label_tokens = list(label_tokens.squeeze(0))
+        first_label_token: int = int(label_tokens[0])
+        first_character_of_label = tokeniser.convert_ids_to_tokens(first_label_token)[0]
+        has_prefix_space = first_character_of_label == whitespace_token
+        if has_prefix_space:
+            add_prefix_space = False
+            break
+
+    return add_prefix_space
+
+
+def get_bos_token(
+    tokeniser: "PreTrainedTokenizer",
+) -> tuple[str, int] | tuple[None, None]:
+    """Get the beginning-of-sequence token from a tokeniser.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+
+    Returns:
+        A pair (token, token_id) representing the beginning-of-sequence token and its
+        token ID, or (None, None) if no BOS token is found.
+    """
+    if isinstance(tokeniser.bos_token, str) and isinstance(tokeniser.bos_token_id, int):
+        return tokeniser.bos_token, tokeniser.bos_token_id
+
+    vocab: dict[str, int] = tokeniser.get_vocab()
+
+    for candidate_bos_token in BOS_TOKENS:
+        if candidate_bos_token in vocab:
+            bos_token = candidate_bos_token
+            bos_token_id = vocab[bos_token]
+            break
+    else:
+        log_once(
+            "The model does not have a beginning-of-sequence token. Please ensure that "
+            "this has been set in the tokeniser's configuration. Using no BOS token."
+            " This may lead to unexpected behavior in the model.",
+            level=logging.WARNING,
+        )
+        return None, None
+
+    log_once(
+        f"Beginning-of-sequence token was not set, but detected it as {bos_token!r} "
+        f"with ID {bos_token_id}.",
+        level=logging.DEBUG,
+    )
+    return bos_token, bos_token_id
+
+
+def get_eos_token(
+    tokeniser: "PreTrainedTokenizer",
+) -> tuple[str, int] | tuple[None, None]:
+    """Get the end-of-sequence token from a tokeniser.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+
+    Returns:
+        A pair (token, token_id) representing the end-of-sequence token and its token
+        ID, or (None, None) if no EOS token is found.
+    """
+    if isinstance(tokeniser.eos_token, str) and isinstance(tokeniser.eos_token_id, int):
+        return tokeniser.eos_token, tokeniser.eos_token_id
+
+    vocab: dict[str, int] = tokeniser.get_vocab()
+
+    for candidate_eos_token in EOS_TOKENS:
+        if candidate_eos_token in vocab:
+            eos_token = candidate_eos_token
+            eos_token_id = vocab[eos_token]
+            break
+    else:
+        log_once(
+            "The model does not have an end-of-sequence token. Please ensure that this "
+            "has been set in the tokeniser's configuration. Using no EOS token. This "
+            "may lead to unexpected behavior in the model.",
+            level=logging.WARNING,
+        )
+        return None, None
+
+    log_once(
+        f"End-of-sequence token was not set, but detected it as {eos_token!r} with "
+        f"ID {eos_token_id}.",
+        level=logging.WARNING,
+    )
+    return eos_token, eos_token_id
+
+
+def get_pad_token(
+    tokeniser: "PreTrainedTokenizer",
+) -> tuple[str, int] | tuple[None, None]:
+    """Get the padding token from a tokeniser.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+
+    Returns:
+        A pair (token, token_id) representing the padding token and its token ID, or
+        (None, None) if no padding token is found.
+    """
+    # If the tokeniser already has a padding token, return it
+    if tokeniser.pad_token is not None and tokeniser.pad_token_id is not None:
+        assert isinstance(tokeniser.pad_token, str), (
+            "Expected tokeniser.pad_token to be a string, but got "
+            f"{type(tokeniser.pad_token)}."
+        )
+        assert isinstance(tokeniser.pad_token_id, int), (
+            "Expected tokeniser.pad_token_id to be an integer, but got "
+            f"{type(tokeniser.pad_token_id)}."
+        )
+        return (tokeniser.pad_token, tokeniser.pad_token_id)
+
+    # If the tokeniser has a BOS token, use it as the padding token
+    if tokeniser.bos_token is not None and tokeniser.bos_token_id is not None:
+        assert isinstance(tokeniser.bos_token, str), (
+            "Expected tokeniser.bos_token to be a string, but got "
+            f"{type(tokeniser.bos_token)}."
+        )
+        assert isinstance(tokeniser.bos_token_id, int), (
+            "Expected tokeniser.bos_token_id to be an integer, but got "
+            f"{type(tokeniser.bos_token_id)}."
+        )
+        pad_token = tokeniser.bos_token
+        pad_token_id = tokeniser.bos_token_id
+
+    # If the tokeniser has an EOS token, use it as the padding token
+    elif tokeniser.eos_token is not None and tokeniser.eos_token_id is not None:
+        assert isinstance(tokeniser.eos_token, str), (
+            "Expected tokeniser.eos_token to be a string, but got "
+            f"{type(tokeniser.eos_token)}."
+        )
+        assert isinstance(tokeniser.eos_token_id, int), (
+            "Expected tokeniser.eos_token_id to be an integer, but got "
+            f"{type(tokeniser.eos_token_id)}."
+        )
+        pad_token = tokeniser.eos_token
+        pad_token_id = tokeniser.eos_token_id
+
+    # Otherwise, try to find a candidate padding token in the vocabulary
+    else:
+        for candidate in PAD_TOKENS:
+            if candidate in tokeniser.get_vocab():
+                pad_token = candidate
+                pad_token_id = tokeniser.get_vocab()[candidate]
+                break
+        else:
+            log_once(
+                "Could not identify a padding token for the model. Please ensure that "
+                "this has been set in the tokeniser's configuration. Using no padding "
+                "token. This may lead to unexpected behavior in the model.",
+                level=logging.WARNING,
+            )
+            return None, None
+
+    log_once(
+        f"Padding token was not set, but detected it as {pad_token!r} with ID "
+        f"{pad_token_id}.",
+        level=logging.DEBUG,
+    )
+    return pad_token, pad_token_id
+
+
+def get_end_of_chat_token_ids(
+    tokeniser: "PreTrainedTokenizer", generative_type: GenerativeType | None
+) -> c.Sequence[int] | None:
+    """Get the end token ID for chat models.
+
+    This is only relevant for tokenisers with a chat template.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+        generative_type:
+            The generative type, or None if not available.
+
+    Returns:
+        The token IDs used to end chats, or None if the tokeniser does not have a chat
+        template or if no end-of-chat token could be found.
+    """
+    if generative_type == GenerativeType.BASE:
+        return None
+
+    user_message: dict[str, str] = dict(role="user", content="X")
+    try:
+        token_ids = apply_chat_template(
+            conversation=[user_message],
+            tokeniser=tokeniser,
+            tokenise=True,
+            add_generation_prompt=False,
+            enable_thinking=generative_type == GenerativeType.REASONING,
+        )
+    except InvalidModel as e:
+        if "does not have a chat template" in str(e):
+            return None
+        raise e
+    assert isinstance(token_ids, list)
+
+    for idx, token in enumerate(tokeniser.convert_ids_to_tokens(token_ids)):
+        if "X" in token:
+            x_token_index = idx
+            break
+    else:
+        log(
+            "Could not locate the end-of-chat token for the model.", level=logging.DEBUG
+        )
+        return None
+
+    end_of_chat_tokens = token_ids[x_token_index + 1 :]
+    if len(end_of_chat_tokens) == 0:
+        log(
+            "Could not locate the end-of-chat token for the model.", level=logging.DEBUG
+        )
+        return None
+
+    log_once(
+        f"Detected end-of-chat token IDs as {end_of_chat_tokens}, corresponding to "
+        f"tokens {tokeniser.convert_ids_to_tokens(end_of_chat_tokens)}.",
+        level=logging.DEBUG,
+    )
+    return end_of_chat_tokens
+
+
+def get_first_label_token_mapping(
+    dataset_config: "DatasetConfig",
+    model_config: "ModelConfig",
+    tokeniser: "PreTrainedTokenizer | None",
+    generative_type: "GenerativeType | None",
+    log_metadata: bool,
+) -> dict[str, str] | bool:
+    """Check if the model should output scores.
+
+    Args:
+        dataset_config:
+            The dataset configuration.
+        model_config:
+            The model configuration.
+        tokeniser:
+            The tokeniser, or None if not available.
+        generative_type:
+            The generative type, or None if not available.
+        log_metadata:
+            Whether to log metadata.
+
+    Returns:
+        A mapping from labels to the first token in each label, or alternatively a
+        Boolean value indicating whether the model should output scores (if the mapping
+        is outputted then the model will always output scores).
+    """
+    if not (dataset_config.task.uses_logprobs and dataset_config.labels):
+        if log_metadata:
+            log_once(
+                "We will not use logprobs with the model, since the dataset does not "
+                "have labels.",
+                level=logging.DEBUG,
+            )
+        return False
+    elif generative_type == GenerativeType.REASONING:
+        if log_metadata:
+            log_once(
+                f"The model {model_config.model_id!r} is a reasoning model and "
+                "thus does not support logprobs, so we do not enable it.",
+                level=logging.DEBUG,
+            )
+        return False
+    elif tokeniser is None:
+        if log_metadata:
+            log_once(
+                f"We will use logprobs with the model {model_config.model_id!r} "
+                "since the dataset supports it and no tokeniser is available.",
+                level=logging.DEBUG,
+            )
+        return True
+
+    local_labels = [
+        dataset_config.prompt_label_mapping[label].strip()
+        for label in dataset_config.labels
+    ]
+
+    # Tokenise some text containing each label, which we will use to extract the
+    # first token of each label
+    all_tokens: c.Sequence[c.Sequence[str]]
+    if not has_chat_template(tokeniser=tokeniser):
+        add_prefix_space = should_prefix_space_be_added_to_labels(
+            labels_to_be_generated=local_labels, tokeniser=tokeniser
+        )
+        all_tokens = [
+            [
+                tokeniser.decode(token_id)
+                for token_id in tokeniser.encode(
+                    text=f" {label}" if add_prefix_space else label,
+                    add_special_tokens=False,
+                )
+            ]
+            for label in local_labels
+        ]
+    else:
+        all_tokens = [
+            tokeniser.convert_ids_to_tokens(
+                ids=apply_chat_template(
+                    conversation=[
+                        dict(role="user", content=""),
+                        dict(role="assistant", content=label),
+                        # Adding extra user message as Mistral tokenisers require
+                        # conversations to end with a user message
+                        dict(role="user", content=""),
+                    ],
+                    tokeniser=tokeniser,
+                    tokenise=True,
+                    add_generation_prompt=True,
+                    enable_thinking=generative_type == GenerativeType.REASONING,
+                )
+            )
+            for label in local_labels
+        ]
+
+    # Remove any non-alphabetic characters from the tokens
+    all_tokens = [
+        [
+            re.sub(
+                pattern=r"^[^a-zæøåüöä0-9 ]+|[^a-zæøåüöä0-9 ]+$",
+                repl="",
+                string=token.lower(),
+            )
+            for token in token_list
+        ]
+        for token_list in all_tokens
+    ]
+
+    # Extract the first token of each label
+    first_tokens: list[str] = list()
+    for token_list, label in zip(all_tokens, local_labels):
+        matching_tokens = [
+            tok for tok in token_list if tok and label.startswith(tok.strip())
+        ]
+        if not matching_tokens:
+            if log_metadata:
+                log_once(
+                    f"No matching token found in token_list for label {label!r}, so "
+                    "we will not use logprobs with the model.",
+                    level=logging.DEBUG,
+                )
+            return False
+        first_tokens.append(matching_tokens[0])
+
+    # Build a mapping from labels to the first token in each label if the first
+    # tokens are distinct
+    if len(first_tokens) == len(set(first_tokens)):
+        mapping = {
+            label: first_token for label, first_token in zip(local_labels, first_tokens)
+        }
+        if log_metadata:
+            log_once(
+                "Using logprobs as evaluation strategy for the model, with the "
+                f"following mapping from labels to their first token: {mapping}.",
+                level=logging.DEBUG,
+            )
+        return mapping
+    else:
+        if log_metadata:
+            log_once(
+                "We will not use logprobs with the model since the first tokens of the "
+                "labels are not distinct. The first tokens for the labels "
+                f"{local_labels} are {first_tokens}",
+                level=logging.DEBUG,
+            )
+        return False
+
+
+def has_chat_template(tokeniser: "PreTrainedTokenizer") -> bool:
+    """Check if a tokeniser has a chat template.
+
+    Args:
+        tokeniser:
+            The tokeniser.
+
+    Returns:
+        Whether the tokeniser has a chat template.
+    """
+    if isinstance(tokeniser, MistralCommonTokenizer):
+        log_once(
+            "The tokeniser is a Mistral tokeniser, so assuming that the model is "
+            "instruction tuned.",
+            level=logging.DEBUG,
+        )
+        return True
+    elif hasattr(tokeniser, "chat_template"):
+        has_template = tokeniser.chat_template is not None
+        if has_template:
+            log_once(
+                "The tokeniser has a chat template, so assuming that the model is "
+                "instruction tuned.",
+                level=logging.DEBUG,
+            )
+        return has_template
+    else:
+        log_once(
+            "We cannot find a chat template for the tokeniser, so assuming that the "
+            "model isn't instruction tuned.",
+            level=logging.DEBUG,
+        )
+        return False
+
+
+def apply_chat_template(
+    conversation: c.Sequence[dict[str, str]],
+    tokeniser: "PreTrainedTokenizer",
+    tokenise: bool,
+    add_generation_prompt: bool,
+    **extra_kwargs,
+) -> str | c.Sequence[int]:
+    """Apply the chat template to a prompt.
+
+    Args:
+        conversation:
+            The conversation to apply the chat template to.
+        tokeniser:
+            The tokeniser.
+        tokenise:
+            Whether to tokenise the resulting prompt, returning a list of token IDs
+            instead of a string.
+        add_generation_prompt:
+            Whether to add a generation prompt at the end of the conversation. This is
+            only relevant for regular Hugging Face tokenisers, as Mistral tokenisers
+            always add a generation prompt.
+        **extra_kwargs:
+            Extra keyword arguments to pass to the tokeniser's `apply_chat_template`
+            method. Only relevant for regular Hugging Face tokenisers.
+
+    Returns:
+        The prompt with the chat template applied, either as a string or a list of
+        token IDs, depending on the value of `tokenise`.
+
+    Raises:
+        InvalidModel:
+            If the tokeniser does not have a chat template.
+    """
+    # Ensure that the first user message is not empty, as this can cause issues with
+    # Jinja2
+    conversation[0]["content"] = conversation[0]["content"] or " "
+
+    if not has_chat_template(tokeniser=tokeniser):
+        raise InvalidModel(
+            "The tokeniser does not have a chat template, so cannot apply it."
+        )
+    elif isinstance(tokeniser, MistralCommonTokenizer):
+        templated_prompt = tokeniser.apply_chat_template(
+            conversation=conversation, tokenize=tokenise
+        )
+    else:
+        templated_prompt = tokeniser.apply_chat_template(
+            conversation=conversation,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=tokenise,
+            **extra_kwargs,
+        )
+    return templated_prompt

euroeval/types.py

@@ -1,18 +1,19 @@
 """Types used throughout the project."""
 
+import collections.abc as c
 import typing as t
 
 from transformers.trainer_utils import EvalPrediction
 
 if t.TYPE_CHECKING:
+    from datasets.arrow_dataset import Dataset
     from numpy.typing import NDArray
 
-    from .data_models import GenerativeModelOutput
+    from .data_models import BenchmarkConfig, GenerativeModelOutput
 
-
-ScoreDict: t.TypeAlias = dict[str, dict[str, float] | list[dict[str, float]]]
-Predictions: t.TypeAlias = "NDArray | list[str] | list[list[str]]"
-Labels: t.TypeAlias = "NDArray | list[str] | list[list[str]]"
+ScoreDict: t.TypeAlias = dict[str, dict[str, float] | c.Sequence[dict[str, float]]]
+Predictions: t.TypeAlias = "NDArray | c.Sequence[str] | c.Sequence[c.Sequence[str]]"
+Labels: t.TypeAlias = "NDArray | c.Sequence[str] | c.Sequence[c.Sequence[str]]"
 
 
 class ComputeMetricsFunction(t.Protocol):
@@ -22,15 +23,20 @@ class ComputeMetricsFunction(t.Protocol):
         self,
         model_outputs_and_labels: EvalPrediction
         | tuple[
-            "NDArray | list[str] | list[list[str]]",
-            "NDArray | list[str] | list[list[str]]",
+            "NDArray | c.Sequence[str] | c.Sequence[c.Sequence[str]]",
+            "NDArray | c.Sequence[str] | c.Sequence[c.Sequence[str]]",
         ],
+        dataset: "Dataset",
+        benchmark_config: "BenchmarkConfig",
     ) -> dict[str, float]:
         """Compute the metrics.
 
         Args:
             model_outputs_and_labels:
                 The model outputs and labels.
+            dataset:
+                The dataset used for evaluation. This is only used in case any
+                additional metadata is used to compute the metrics.
 
         Returns:
             The computed metrics.
@@ -43,7 +49,7 @@ class ExtractLabelsFunction(t.Protocol):
 
     def __call__(
         self, input_batch: dict[str, list], model_output: "GenerativeModelOutput"
-    ) -> list[str]:
+    ) -> c.Sequence[str]:
         """Extract the labels from the generated output.
 
         Args:
@@ -58,7 +64,7 @@ class ExtractLabelsFunction(t.Protocol):
         ...
 
 
-def is_list_of_int(x: object) -> t.TypeGuard[list[int]]:
+def is_list_of_int(x: object) -> t.TypeGuard[c.Sequence[int]]:
     """Check if an object is a list of integers.
 
     Args:
@@ -71,7 +77,7 @@ def is_list_of_int(x: object) -> t.TypeGuard[list[int]]:
     return isinstance(x, list) and all(isinstance(i, int) for i in x)
 
 
-def is_list_of_list_of_int(x: object) -> t.TypeGuard[list[list[int]]]:
+def is_list_of_list_of_int(x: object) -> t.TypeGuard[c.Sequence[c.Sequence[int]]]:
     """Check if an object is a list of list of integers.
 
     Args:
@@ -88,7 +94,7 @@ def is_list_of_list_of_int(x: object) -> t.TypeGuard[list[list[int]]]:
     )
 
 
-def is_list_of_str(x: object) -> t.TypeGuard[list[str]]:
+def is_list_of_str(x: object) -> t.TypeGuard[c.Sequence[str]]:
     """Check if an object is a list of integers.
 
     Args: