# This file was auto-generated by Fern from our API Definition.

import enum
import typing

T_Result = typing.TypeVar("T_Result")


class DatasetType(str, enum.Enum):
    """
    The type of the dataset
    """

    EMBED_INPUT = "embed-input"
    EMBED_RESULT = "embed-result"
    CLUSTER_RESULT = "cluster-result"
    CLUSTER_OUTLIERS = "cluster-outliers"
    RERANKER_FINETUNE_INPUT = "reranker-finetune-input"
    PROMPT_COMPLETION_FINETUNE_INPUT = "prompt-completion-finetune-input"
    SINGLE_LABEL_CLASSIFICATION_FINETUNE_INPUT = "single-label-classification-finetune-input"
    CHAT_FINETUNE_INPUT = "chat-finetune-input"
    MULTI_LABEL_CLASSIFICATION_FINETUNE_INPUT = "multi-label-classification-finetune-input"

    def visit(
        self,
        embed_input: typing.Callable[[], T_Result],
        embed_result: typing.Callable[[], T_Result],
        cluster_result: typing.Callable[[], T_Result],
        cluster_outliers: typing.Callable[[], T_Result],
        reranker_finetune_input: typing.Callable[[], T_Result],
        prompt_completion_finetune_input: typing.Callable[[], T_Result],
        single_label_classification_finetune_input: typing.Callable[[], T_Result],
        chat_finetune_input: typing.Callable[[], T_Result],
        multi_label_classification_finetune_input: typing.Callable[[], T_Result],
    ) -> T_Result:
        if self is DatasetType.EMBED_INPUT:
            return embed_input()
        if self is DatasetType.EMBED_RESULT:
            return embed_result()
        if self is DatasetType.CLUSTER_RESULT:
            return cluster_result()
        if self is DatasetType.CLUSTER_OUTLIERS:
            return cluster_outliers()
        if self is DatasetType.RERANKER_FINETUNE_INPUT:
            return reranker_finetune_input()
        if self is DatasetType.PROMPT_COMPLETION_FINETUNE_INPUT:
            return prompt_completion_finetune_input()
        if self is DatasetType.SINGLE_LABEL_CLASSIFICATION_FINETUNE_INPUT:
            return single_label_classification_finetune_input()
        if self is DatasetType.CHAT_FINETUNE_INPUT:
            return chat_finetune_input()
        if self is DatasetType.MULTI_LABEL_CLASSIFICATION_FINETUNE_INPUT:
            return multi_label_classification_finetune_input()
