Adaptors

@abc.abstractmethod
def evaluate(self, data: Mapping[str, Sequence[Any]]) -> Sequence[float] | NDArray:
    """
    Evaluate a particular set of entries with a language model.
    Returns a list of scores, one for each entry, in the same order.

    Parameters:
        data: A mapping from column names to the data in that column.

    Returns:
        The scores for each entry. Scores must be floating point numbers.
    """

    ...

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)

@staticmethod
def task_choices(
    name: Literal["sst2", "mrpc", "mnli", "qqp", "rte", "qnli"],
    split: Literal["train", "validation", "test"],
) -> Sequence[str]:
    """
    Get the valid choices for a particular task and split.

    Parameters:
        name: The name of the task.
        split: The split of the task.

    Returns:
        The valid choices for the task and split.
    """

    LOGGER.debug("Getting choices for task", task=name)

    # Perform checks for supported kinds of datasets.
    match name:
        case "sst2" | "mrpc" | "mnli" | "qqp" | "rte" | "qnli":
            pass
        case _:
            raise ValueError(f"Unknown task name {name}")

    # Perform checks for supported kinds of splits.
    match split:
        case "train" | "validation" | "test":
            pass
        case _:
            raise ValueError(f"Unknown split {split}")

    # The actual mux.
    match name, split:
        case "sst2", _:
            return ["negative", "positive"]
        case "mrpc", _:
            return ["not equivalent", "equivalent"]
        # All following cases all use "unlabeled" for "test".
        case _, "test":
            return ["unlabeled"]
        case "mnli", _:
            return ["entailment", "neutral", "contradiction"]
        case "qqp", _:
            return ["not duplicate", "duplicate"]
        case "rte", _:
            return ["entailment", "not entailment"]
        case "qnli", _:
            return ["entailment", "not entailment"]

    raise RuntimeError("Unreachable")

@abc.abstractmethod
def evaluate(self, data: Mapping[str, Sequence[Any]]) -> Sequence[float] | NDArray:
    """
    Evaluate a particular set of entries with a language model.
    Returns a list of scores, one for each entry, in the same order.

    Parameters:
        data: A mapping from column names to the data in that column.

    Returns:
        The scores for each entry. Scores must be floating point numbers.
    """

    ...

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)

@staticmethod
def numeric_choices(question: str, choices: Sequence[str]) -> str:
    """
    Convert a multiple choice question into a numeric choice question.
    Returns a tuple of generated prompt and list of valid choices.
    """

    return (
        f"{question}\nSelect from one of the following (answer in number):\n"
        + "\n".join(f"{i}) {choice}" for i, choice in enumerate(choices, 1))
    )

def on_storage(self, storage: Storage, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a storage.
    Given indices and a storage,
    this method will extract the corresponding entries from the storage,
    and evaluate them with `Adaptor.evaluate`.

    Parameters:
        storage: The storage to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    indices = np.array(indices).astype("i")

    # Reshape the indices into 1D to evaluate.
    indices_shape = indices.shape
    indices = indices.ravel()

    items = storage[indices.tolist()]
    result = np.array(self.evaluate(data=items))

    # Reshape back.
    return result.reshape(indices_shape)

def on_corpus(self, corpus: Corpus, indices: ArrayLike) -> NDArray:
    """
    Evaluate a particular set of indices on a corpus.
    A convenience wrapper around `Adaptor.on_storage`.

    Parameters:
        corpus: The corpus to evaluate.
        indices: The indices to evaluate.

    Returns:
        The scores for each entry. The shape must be the same as the indices.
    """

    return self.on_storage(storage=corpus.storage, indices=indices)