from .embedders import (
    CaduceusEmbedder,
    DNABertEmbedder,
    NucleotideTransformerEmbedder,
    GPNEmbedder,
    SegmentNTEmbedder,
    ESMEmbedder,
    ESMDBPEmbedder,
    ProGenEmbedder,
)


def get_embedder(name, device, for_dna=True):
    name = name.lower()
    if for_dna:
        if name == "caduceus":
            return CaduceusEmbedder(device)
        if name == "dnabert":
            return DNABertEmbedder(device)
        if name == "nucleotide":
            return NucleotideTransformerEmbedder(device)
        if name == "gpn":
            return GPNEmbedder(device)
        if name == "segmentnt":
            return SegmentNTEmbedder(device)
    else:
        if name in ("esm",):
            return ESMEmbedder(device)
        if name in ("esm-dbp", "esm_dbp"):
            return ESMDBPEmbedder(device)
        if name == "progen":
            return ProGenEmbedder(device)
    raise ValueError(f"Unknown model {name} (for_dna={for_dna})")