Factories

def index_class(name: str | IndexName, /) -> type[Index]:
    """
    Get the index class for the given name.

    Parameters:
        name: The name of the index.
    """

    name = IndexName.lookup(name)

    match name:
        case IndexName.FAISS:
            return FaissIndex
        case IndexName.HNSWLIB:
            return HnswlibIndex
        case IndexName.POLAR:
            return PolarIndex
        case IndexName.WHITENING:
            return WhiteningIndex
        case _:
            raise ValueError(f"Unknown index name: {name}")

@common.correct_kwargs
def storage(
    storage: str | StorageName, /, *, path: str = "", name: str = "", split: str = ""
) -> Storage:
    """
    Create a single storage.

    Parameters:
        storage: The name of the storage.
        path: The path to the storage.
        name: The name of the storage.
        split: The split to use.

    Returns:
        The storage instance.

    Raises:
        ValueError: If the storage is unknown.
    """

    storage = StorageName.lookup(storage)
    match storage:
        case StorageName.PANDAS:
            return common.correct_kwargs(PandasStorage.from_jsonl_file)(path)
        case StorageName.DATASETS:
            return common.correct_kwargs(DatasetsStorage)(
                path=path, name=name, split=split
            )
        case _:
            raise ValueError(f"Unknown storage name {storage}")

def embedder(
    name: str | EmbedderName,
    /,
    *,
    model_name: str | list[str],
    device: str = "auto",
    batch_size: int,
) -> Embedder:
    """
    Create an embedder.

    Parameters:
        name: The name of the embedder.
        model_name: The model name to use.
        device: The device to use.
        batch_size: The batch size to use.

    Returns:
        The embedder instance.

    Raises:
        ValueError: If the name is unknown.
        TypeError: If the model name is not a string for SBERT or Huggingface,
            or not a list of strings for HuggingfaceEnsemble.
    """

    match EmbedderName.lookup(name):
        case EmbedderName.SBERT:
            if not isinstance(model_name, str):
                raise TypeError(
                    "SbertEmbedder requires a single model name. "
                    f"Got {model_name} instead."
                )

            return common.correct_kwargs(SbertEmbedder)(
                model_name=model_name,
                device=common.auto_device(device),
                batch_size=batch_size,
            )
        case EmbedderName.HUGGINGFACE:
            if not isinstance(model_name, str):
                raise TypeError(
                    "HuggingfaceEmbedder requires a single model name. "
                    f"Got {model_name} instead."
                )
            return common.correct_kwargs(HuggingfaceEmbedder)(
                path=model_name,
                device=common.auto_device(device),
                batch_size=batch_size,
            )
        case EmbedderName.HUGGINGFACE_ENSEMBLE:
            if not isinstance(model_name, list):
                raise TypeError(
                    "HuggingfaceEnsembleEmbedder requires a list of model names. "
                    f"Got {model_name} instead."
                )

            device_list = common.auto_device_list(device, len(model_name))
            return common.correct_kwargs(EnsembleEmbedder)(
                [
                    HuggingfaceEmbedder(path=model, device=dev, batch_size=batch_size)
                    for model, dev in zip(model_name, device_list)
                ]
            )
        case _:
            raise ValueError(f"Unknown embedder name: {name}")

def corpus(
    name: str | CorpusName = CorpusName.COMPOSED,
    /,
    *,
    storage: Storage,
    embedder: Embedder,
    keys: Sequence[str],
    index_name: str | IndexName,
    **index_kwargs: Any,
) -> Corpus:
    """
    Create a corpus.

    Parameters:
        name: The name of the corpus.
        storage: The storage to use.
        embedder: The embedder to use.
        keys: The key to use for the index.
        index_name: The name of the index backend to use.
        **index_kwargs: The keyword arguments to pass to the index backend.

    Returns:
        The corpus instance.

    Raises:
        ValueError: If the name is unknown.
    """

    if CorpusName.lookup(name) is not CorpusName.COMPOSED:
        raise ValueError(f"Unknown corpus name: {name}")

    return common.correct_kwargs(ComposedCorpus.index_storage)(
        storage=storage,
        embedder=embedder,
        keys=keys,
        index_backend=indices.index_class(index_name),
        **indices.index_set_backends(index_kwargs),
    )

def adaptor(name: str | AdaptorName, /, **kwargs: Any) -> Adaptor:
    """
    Create an adaptor.

    Parameters:
        name: The name of the adaptor.
        **kwargs: The keyword arguments to pass to the adaptor.
            See the documentation of the corresponding adaptor for details.

    Returns:
        The adaptor instance.

    Raises:
        ValueError: If the name is unknown.
    """

    name = AdaptorName.lookup(name)

    match name:
        case AdaptorName.BIGBENCH_MC:
            return common.correct_kwargs(BigBenchMultipleChoice)(**kwargs)
        case AdaptorName.BIGBENCH_QA:
            return common.correct_kwargs(BigBenchQuestionAnswer)(**kwargs)
        case AdaptorName.SST2:
            return common.correct_kwargs(Sst2QuestionAnswer)(**kwargs)
        case AdaptorName.GLUE:
            return common.correct_kwargs(GlueAdaptor)(**kwargs)
        case _:
            raise ValueError(f"Unknown adaptor name: {name}")

def generative(
    name: str | GeneratorName,
    /,
    *,
    model_path: str,
    batch_size: int,
    device: str = "auto",
    add_sep_token: bool = False,
) -> GenerativeModel:
    """
    Create a generative model.

    Parameters:
        name: The name of the model.
        model_path: The path to the model.
        batch_size: The batch size to use.
        device: The device to use.
        add_sep_token: Whether to add the sep token.

    Returns:
        The generative model instance.

    Raises:
        ValueError: If the name is unknown.
    """

    device = common.auto_device(device)

    match GeneratorName.lookup(name):
        case GeneratorName.HUGGINGFACE_GENERATIVE:
            return common.correct_kwargs(HuggingfaceGenerativeLM)(
                model_path=model_path,
                batch_size=batch_size,
                device=device,
                add_sep_token=add_sep_token,
            )
        case _:
            raise ValueError(f"Unknown LM name {name}")

def classifier(
    name: str | ClassifierName,
    /,
    *,
    model_path: str,
    batch_size: int,
    choices: Sequence[str],
    device: str = "auto",
    add_sep_token: bool = False,
) -> ClassifierModel:
    device = common.auto_device(device)

    match ClassifierName.lookup(name):
        case ClassifierName.HUGGINGFACE_LOGITS:
            return common.correct_kwargs(HuggingfaceLogitsLM)(
                model_path=model_path,
                batch_size=batch_size,
                device=device,
                choices=choices,
                add_sep_token=add_sep_token,
            )
        case ClassifierName.HUGGINGFACE_SEQUENCE:
            return common.correct_kwargs(HuggingfaceSequenceLM)(
                model_path=model_path,
                device=device,
                choices=choices,
                add_sep_token=add_sep_token,
            )
        case _:
            raise ValueError(f"Unknown LM name {name}")

def optimizer(
    name: str | OptimizerName, /, *, corpus: Corpus, adaptor: Adaptor, **kwargs: Any
) -> Optimizer:
    """
    Create an optimizer instance.

    Parameters:
        name: The name of the optimizer.
        corpus: The corpus to optimize.
        adaptor: The adaptor to use.
        **kwargs: Additional keyword arguments to pass to the optimizer.
            See the documentation for the specific optimizer for details.

    Returns:
        The optimizer instance.

    Raises:
        ValueError: If the name is unknown.
    """

    name = OptimizerName.lookup(name)

    klass: type[Optimizer]

    match name:
        case OptimizerName.BAYESIAN:
            klass = AxServiceOptimizer
        case OptimizerName.KMEANS:
            klass = KMeansOptimizer
        case OptimizerName.KMEDOIDS:
            klass = KMedoidsOptimizer
        case OptimizerName.BRUTE:
            klass = BruteForceOptimizer
        case OptimizerName.RANDOM:
            klass = RandomOptimizer
        case OptimizerName.UNIFORM:
            klass = UniformOptimizer
        case _:
            raise ValueError(f"Unknown optimizer name: {name}")

    corpus_evaluator = CorpusEvaluator(corpus=corpus, adaptor=adaptor)
    return klass(index_eval=corpus_evaluator, index=corpus.index, **kwargs)