
"""Organize the models."""

from .gen import GEN, GENPYG, GENcatpos, GENCross, GENDist, GENnopos
from .gen_nograph import GEN_nograph
from .gka import GKA
from .msa import MSA, MSAEncoderOnly
from .nps import CNP, CNPCross, CNPDist
from .perceiver import Perceiver
from .transformer import (
    TFS,
    EncoderOnly,
    TFSCross,
    TFSDist,
    TFSfull,
    TFSfullpn,
    TFSOne,
    TFSpos,
)

default_models = {
    "encoder_only": EncoderOnly,
    "perceiver": Perceiver,
    "tfs": TFS,
    "tfspos": TFSpos,
    "tfsfull": TFSfull,
    "tfsfullpn": TFSfullpn,
    "tfscross": TFSCross,
    "tfsdist": TFSDist,
    "tfsone": TFSOne,
    "tfsnormalone": TFS,
    "tfsnormaltwo": TFS,
    "tfsnormalfour": TFS,
    "np": CNP,
    "gka": GKA,
    "npdist": CNPDist,
    "npcross": CNPCross,
    "gen_nograph": GEN_nograph,
    "msa": MSA,
    "msa_encoder_only": MSAEncoderOnly,
}

gen_models = {
    "gen": GEN,
    "gennormalone": GEN,
    "gennormaltwo": GEN,
    "gennormalfour": GEN,
    "genpyg": GENPYG,
    "gencatpos": GENcatpos,
    "gennopos": GENnopos,
    "gencross": GENCross,
    "gendist": GENDist,
}
