Language Models

@abc.abstractmethod
def generate(self, prompts: Sequence[str], /) -> Sequence[str]:
    """
    TODO:
        Add logits.

    Generate a sequence of responses given prompts.
    The length of the response is the same as the prompt.
    The response would be a continuation of the prompt,
    and the prompts would be the prefix of the response.

    Parameters:
        prompts: The prompts to generate.

    Returns:
        The generated responses. The length must be the same as the prompts.
    """

    ...

def classify(self, prompts: Sequence[str], /) -> NDArray:
    """
    Classify the given prompts.

    Parameters:
        prompts: The prompts to classify.

    Returns:
        The logits for each prompt and choice.

    Raises:
        ValueError: If the shape of the logits is not [len(prompts), len(choices)].
    """

    classified = self._classify(prompts)

    if list(classified.shape) != [len(prompts), len(self.choices)]:
        raise ValueError(
            f"Expected logits to have shape {[len(prompts), len(self.choices)]}, "
            f"but got {classified.shape}"
        )

    return classified

@abc.abstractmethod
def _classify(self, prompts: Sequence[str], /) -> NDArray:
    """
    Generate logits given prompts.

    Parameters:
        prompts: The prompts to classify.

    Returns:
        The logits for each prompt and choice.
    """

    ...

def classify(self, prompts: Sequence[str], /) -> NDArray:
    """
    Classify the given prompts.

    Parameters:
        prompts: The prompts to classify.

    Returns:
        The logits for each prompt and choice.

    Raises:
        ValueError: If the shape of the logits is not [len(prompts), len(choices)].
    """

    classified = self._classify(prompts)

    if list(classified.shape) != [len(prompts), len(self.choices)]:
        raise ValueError(
            f"Expected logits to have shape {[len(prompts), len(self.choices)]}, "
            f"but got {classified.shape}"
        )

    return classified

def classify(self, prompts: Sequence[str], /) -> NDArray:
    """
    Classify the given prompts.

    Parameters:
        prompts: The prompts to classify.

    Returns:
        The logits for each prompt and choice.

    Raises:
        ValueError: If the shape of the logits is not [len(prompts), len(choices)].
    """

    classified = self._classify(prompts)

    if list(classified.shape) != [len(prompts), len(self.choices)]:
        raise ValueError(
            f"Expected logits to have shape {[len(prompts), len(self.choices)]}, "
            f"but got {classified.shape}"
        )

    return classified

def to(self, device: str, /) -> Self:
    """
    Move the tokenizer to the given device.

    Parameters:
        device: The device to move to.
    """
    self._device = device
    return self

def tokenize(self, prompts: Sequence[str], /, max_length: int | None = None):
    """
    Tokenize, pad, truncate, cast to device, and yield the encoded results.
    Returning `BatchEncoding` but not marked in the type hint
    due to optional dependency.

    Parameters:
        prompts: The prompts to tokenize.

    Returns:
        (BatchEncoding): The tokenized prompts.
    """
    if not isinstance(prompts, list):
        prompts = list(prompts)

    inputs = self._tokenizer(
        prompts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    )
    return inputs.to(self.device)

def encode(
    self,
    prompts: Sequence[str],
    /,
    return_tensors: str | None = None,
    add_special_tokens: bool = True,
):
    """
    Encode the given prompts.

    Parameters:
        prompts: The prompts to encode.
        return_tensors: Whether to return tensors.
        add_special_tokens: Whether to add special tokens.

    Returns:
        (Any): The encoded prompts.
    """

    return self._tokenizer.encode(
        prompts,
        return_tensors=return_tensors,
        add_special_tokens=add_special_tokens,
    )

def decode(self, outputs: Any, /, skip_special_tokens: bool = True) -> str:
    """
    Decode the given outputs.

    Parameters:
        outputs: The outputs to decode.
        skip_special_tokens: Whether to skip special tokens.

    Returns:
        The decoded outputs.
    """

    return self._tokenizer.decode(outputs, skip_special_tokens=skip_special_tokens)

def batch_decode(
    self, outputs: Any, /, skip_special_tokens: bool = True
) -> list[str]:
    """
    Batch decode the given outputs.

    Parameters:
        outputs: The outputs to decode.
        skip_special_tokens: Whether to skip special tokens.

    Returns:
        The batch decoded outputs.
    """

    return self._tokenizer.batch_decode(
        outputs, skip_special_tokens=skip_special_tokens
    )