api.py

import abc
from collections import OrderedDict
from enum import Enum
from typing import Any, List, Mapping, Optional, Union

import datasets
import typing_extensions
from datasets import Dataset


class DatasetSplit(Enum):
    """Enum for train, validation, and test sets."""

    TRAIN = "train"
    VALIDATION = "validation"
    TEST = "test"


class TaskType(Enum):
    FREE_FORM = "free_form"
    MULTIPLE_CHOICE = "multiple_choice"
    SEQUENCE_LABELING = "sequence_labeling"


class PreparationStrategy(Enum):
    PROMPT_BASED_TESTING = "prompt_based_testing"


EvaluationResult = OrderedDict[str, float]


class TaskInterface(abc.ABC):
    """The base class for GenBench tasks.

    This class defines the interface for GenBench tasks. All tasks should
    inherit from this class and implement the methods defined here.
    """

    def __init__(self):
        pass

    @abc.abstractmethod
    def format_example(self, example: Mapping[str, Any]) -> Mapping[str, Any]:
        """Perform preprocessing/formatting on an example-level.

        By default, this method does nothing more than mapping original data source
        fields to the expected fields.

        `example` directly comes from the data source (e.g. downloaded HF dataset),
        and it may contain fields such as `question` or `answer`. This method should
        prepare the example used in the task. i.e. should create fields `input`,
        `target`, `target_scores`, or `target_labels` depending on the task type.

        Args:
            example: A dictionary containing key-value pairs for an example from the source dataset.


        Returns:
            A dictionary containing key-value pairs for the preprocessed/formatted example.
            The dictionary should contain keys `input`, `target`, `target_scores`, or `target_label`
            depending on the task type.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def evaluate_predictions(
        self,
        *,
        predictions: List[Mapping[str, Any]] = None,
        gold: datasets.Dataset = None,
    ) -> EvaluationResult:
        """Evaluate the predictions of the model against the gold data.

        Args:
            predictions: A list of dictionaries, where each dictionary contains the predicted
                         values for an example. The keys are strings and the values can be any type.
            gold: A HuggingFace `datasets.Dataset` object containing the ground truth data for the task.

        Returns:
            A dictionary containing key-value pairs for the evaluation metric(s) computed on the predicted
            values. The keys are strings representing the name of the evaluation metric and the values are
            floating-point numbers.

        Raises:
            ValueError: If a metric returns None.
        """
        ...

    @abc.abstractmethod
    def get_datasets_raw(self) -> Mapping[str, datasets.Dataset]:
        """Get the raw dataset.

        By default, this method loads the dataset specified in the task's
        config.jsonnet, and re-split it based on split.json if it exists.
        If the task creator wishes to mix and match data points, they can
        override this method.

        Returns:
            A dictionary containing key-value pairs for the raw datasets.
            The keys are strings representing the name of the dataset split
            (e.g., "train", "validation", "test") and the values are
            HuggingFace `datasets.Dataset` objects containing the raw data for the corresponding split.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_prepared_datasets(
        self,
        preparation_strategy: PreparationStrategy,
        shot_list: Optional[List[int]] = None,
        random_seed: int = 42,
    ) -> Union[Mapping[DatasetSplit, Dataset], Mapping[int, Dataset]]:
        """
        Get the prepared datasets based on the specified preparation strategy.

        This method typically relies on `get_dataset_raw` to load the raw dataset and then
        applies the given preparation strategy to prepare the datasets.

        Args:
            preparation_strategy (PreparationStrategy):
                The strategy to be used for dataset preparation.
            shot_list (Optional[List[int]]):
                A list of integers representing the number of shots to be used
                in few-shot learning tasks. If None, all data is used. Defaults to None.
            random_seed (int, optional):
                Seed for random number generation, used for reproducibility. Defaults to 42.

        Returns:
            Union[Mapping[DatasetSplit, Dataset], Mapping[int, Dataset]]:
                A dictionary containing key-value pairs for the prepared datasets.
                If shot_list is provided, the keys are integers representing the number of shots,
                otherwise, they are of type DatasetSplit representing the dataset split (e.g., "train", "validation",
                "test"). The values are HuggingFace `datasets.Dataset` objects containing the prepared data for the
                corresponding split or number of shots.
        """

        raise NotImplementedError()


class Formatter(typing_extensions.Protocol):
    def __call__(self, example: Mapping[str, Any], **kwargs) -> Mapping[str, Any]:
        """A callable that formats each example."""