Additional Types¶

class EncodeKwargs(TypedDict):
    """Keyword arguments for encoding methods.

    Attributes:
        batch_size: The batch size to use for encoding.
        show_progress_bar: Whether to show a progress bar during encoding.
    """

    batch_size: NotRequired[int]
    show_progress_bar: NotRequired[bool]

class PromptType(str, Enum):
    """The type of prompt used in the input for retrieval models. Used to differentiate between queries and documents.

    Attributes:
        query: A prompt that is a query.
        document: A prompt that is a document.
    """

    query = "query"
    document = "document"

class ConversationTurn(TypedDict):
    """A conversation, consisting of a list of messages.

    Attributes:
        role: The role of the message sender.
        content: The content of the message.
    """

    role: str
    content: str

class TextInput(TypedDict):
    """The input to the encoder for text.

    Attributes:
        text: The text to encode. Can be a list of texts or a list of lists of texts.
    """

    text: list[str]

class CorpusInput(TextInput):
    """The input to the encoder for retrieval corpus.

    Attributes:
        title: The title of the text to encode. Can be a list of titles or a
            list of lists of titles.
        body: The body of the text to encode. Can be a list of bodies or a
            list of lists of bodies.
    """

    title: list[str]
    body: list[str]

class QueryInput(TextInput):
    """The input to the encoder for queries.

    Attributes:
        query: The query to encode. Can be a list of queries or a list of lists of queries.
        conversation: Optional. A list of conversations, each conversation is a list of messages.
        instruction: Optional. A list of instructions to encode.
    """

    query: list[str]
    conversation: NotRequired[list[Conversation]]
    instruction: NotRequired[list[str]]

class ImageInput(TypedDict):
    """The input to the encoder for images.

    Attributes:
        image: The image to encode. Can be a list of images or a list of lists of images.
    """

    image: list[Image.Image]

class AudioInputItem(TypedDict):
    """An audio item for the AudioInput.

    Dataset based on `datasets.Audio` will be converted to this format during encoding.

    Attributes:
        array: The audio array as bytes.
        sampling_rate: The sampling rate of the audio.
    """

    array: npt.NDArray[np.floating]
    sampling_rate: int

class AudioInput(TypedDict):
    """The input to the encoder for audio.

    Attributes:
        audio: The audio to encode. Can be a list of audio files or a list of lists of audio files.
    """

    audio: list[AudioInputItem]

class VideoInputItem(TypedDict):
    """A video item for the VideoInput.

    Dataset based on `datasets.Video` will be converted to this format during encoding.

    Attributes:
        frames: The video frames as Tensor.
        audio: The audio array as AudioInputItem.
    """

    frames: torch.Tensor
    audio: AudioInputItem

class VideoInput(TypedDict):
    """The input to the encoder for videos.

    Attributes:
        video: The video to encode. VideoDecoder object.
    """

    video: VideoInputItem

class MultimodalInput(  # type: ignore[misc]
    TextInput, CorpusInput, QueryInput, ImageInput, AudioInput, VideoInput
):
    """The input to the encoder for multimodal data."""

    pass

class OutputDType(str, Enum):
    """Enum for valid compression levels.

    Used by the CompressionWrapper class and specified by models to indicate the dtypes of output embeddings they
    support internally.
    """

    FLOAT16 = "float16"
    BF16 = "bfloat16"
    INT8 = "int8"
    INT4 = "int4"
    UINT8 = "uint8"
    UINT4 = "uint4"
    BINARY = "binary"
    FLOAT8_E4M3FN = "float8_e4m3fn"
    FLOAT8_E5M2 = "float8_e5m2"
    FLOAT8_E8M0FNU = "float8_e8m0fnu"
    FLOAT8_E4M3FNUZ = "float8_e4m3fnuz"
    FLOAT8_E5M2FNUZ = "float8_e5m2fnuz"

    def get_dtype(self) -> torch.dtype:
        """Returns the PyTorch dtype that matches the enum.

        Output types that are not natively supported by PyTorch like 4-bit integers require specific mapping to the
        desired dtype.
        """
        if self == OutputDType.UINT4:
            return torch.uint8
        elif self == OutputDType.INT4:
            return torch.int8
        elif self == OutputDType.BINARY:
            return torch.bool
        return getattr(torch, self.value)

class RetrievalEvaluationResult(NamedTuple):
    """Holds the results of retrieval evaluation metrics."""

    all_scores: dict[str, dict[str, float]]
    ndcg: dict[str, float]
    map: dict[str, float]
    recall: dict[str, float]
    precision: dict[str, float]
    naucs: dict[str, float]
    mrr: dict[str, float]
    naucs_mrr: dict[str, float]
    hit_rate: dict[str, float]

class SplitDescriptiveStatistics(TypedDict):
    """Base class for descriptive statistics for the subset."""

    pass

class DescriptiveStatistics(TypedDict, SplitDescriptiveStatistics):
    """Class for descriptive statistics for the full task.

    Attributes:
        num_samples: Total number of samples
        hf_subset_descriptive_stats: HFSubset descriptive statistics (only for multilingual datasets)
    """

    num_samples: int
    hf_subset_descriptive_stats: NotRequired[dict[HFSubset, SplitDescriptiveStatistics]]

class TextStatistics(TypedDict):
    """Class for descriptive statistics for texts.

    Attributes:
        total_text_length: Total length of all texts
        min_text_length: Minimum length of text
        average_text_length: Average length of text
        max_text_length: Maximum length of text
        unique_texts: Number of unique texts
    """

    total_text_length: int
    min_text_length: int
    average_text_length: float
    max_text_length: int
    unique_texts: int

class ImageStatistics(TypedDict):
    """Class for descriptive statistics for images.

    Attributes:
        min_image_width: Minimum width of images
        average_image_width: Average width of images
        max_image_width: Maximum width of images

        min_image_height: Minimum height of images
        average_image_height: Average height of images
        max_image_height: Maximum height of images

        unique_images: Number of unique images
    """

    min_image_width: float
    average_image_width: float
    max_image_width: float

    min_image_height: float
    average_image_height: float
    max_image_height: float

    unique_images: int

class AudioStatistics(TypedDict):
    """Class for descriptive statistics for audio.

    Attributes:
        total_duration_seconds: Total length of all audio clips in total frames
        min_duration_seconds: Minimum length of audio clip in seconds
        average_duration_seconds: Average length of audio clip in seconds
        max_duration_seconds: Maximum length of audio clip in seconds
        unique_audios: Number of unique audio clips
        average_sampling_rate: Average sampling rate
        sampling_rates: Dict of unique sampling rates and their frequencies
    """

    total_duration_seconds: float

    min_duration_seconds: float
    average_duration_seconds: float
    max_duration_seconds: float

    unique_audios: int

    average_sampling_rate: float
    sampling_rates: dict[int, int]

class LabelStatistics(TypedDict):
    """Class for descriptive statistics for texts.

    Attributes:
        min_labels_per_text: Minimum number of labels per text
        average_label_per_text: Average number of labels per text
        max_labels_per_text: Maximum number of labels per text

        unique_labels: Number of unique labels
        labels: dict of label frequencies
    """

    min_labels_per_text: int
    average_label_per_text: float
    max_labels_per_text: int

    unique_labels: int
    labels: dict[str, dict[str, int]]

class ScoreStatistics(TypedDict):
    """Class for descriptive statistics for texts.

    Attributes:
        min_score: Minimum score
        avg_score: Average score
        max_score: Maximum score
    """

    min_score: int | float
    avg_score: float
    max_score: int | float

class TopRankedStatistics(TypedDict):
    """Statistics for top ranked documents in a retrieval task.

    Attributes:
        num_top_ranked: Total number of top ranked documents across all queries.
        min_top_ranked_per_query: Minimum number of top ranked documents for any query.
        average_top_ranked_per_query: Average number of top ranked documents per query.
        max_top_ranked_per_query: Maximum number of top ranked documents for any query.
    """

    num_top_ranked: int
    min_top_ranked_per_query: int
    average_top_ranked_per_query: float
    max_top_ranked_per_query: int

class RelevantDocsStatistics(TypedDict):
    """Statistics for relevant documents in a retrieval task.

    Attributes:
        num_relevant_docs: Total number of relevant documents across all queries.
        min_relevant_docs_per_query: Minimum number of relevant documents for any query.
        average_relevant_docs_per_query: Average number of relevant documents per query.
        max_relevant_docs_per_query: Maximum number of relevant documents for any query.
        unique_relevant_docs: Number of unique relevant documents across all queries.
    """

    num_relevant_docs: int
    min_relevant_docs_per_query: int
    average_relevant_docs_per_query: float
    max_relevant_docs_per_query: float
    unique_relevant_docs: int