"""Module implementing the model architectures and utily functions."""

from mt_dao.models.deptoe import (
    DEPToEBlock,
    DEPToEBlockConfig,
    DEPToECausalLM,
    DEPToEForCausalLM,
    DEPToEModel,
    DEPToEPreTrainedModel,
)
from mt_dao.models.deptoe.configuration_deptoe import DEPToEConfig
from mt_dao.models.ffns import (
    SigmaMoE,
    SigmaMoEConfig,
    entropy_l,
    entropy_reg,
    log_mean,
)
from mt_dao.models.mpt import (
    ComposerMPTMuPCausalLM,
    MPTForQuestionAnswering,
    MPTForSequenceClassification,
    MPTForTokenClassification,
    MPTMuPForCausalLM,
)
from mt_dao.models.rope import RotaryPositionalEncoding
from mt_dao.models.switch_head import (
    AttentionMask,
    KVCache,
    SwitchHead,
    SwitchHeadConfig,
)
from mt_dao.models.utils import (
    CROSS_ENTROPY_IGNORE_INDEX,
    compute_loss_from_logits,
    get_targets,
)

__all__ = [
    "CROSS_ENTROPY_IGNORE_INDEX",
    "AttentionMask",
    "ComposerMPTMuPCausalLM",
    "DEPToEBlock",
    "DEPToEBlockConfig",
    "DEPToECausalLM",
    "DEPToEConfig",
    "DEPToEForCausalLM",
    "DEPToEModel",
    "DEPToEPreTrainedModel",
    "KVCache",
    "MPTForQuestionAnswering",
    "MPTForSequenceClassification",
    "MPTForTokenClassification",
    "MPTMuPForCausalLM",
    "RotaryPositionalEncoding",
    "SigmaMoE",
    "SigmaMoEConfig",
    "SwitchHead",
    "SwitchHeadConfig",
    "compute_loss_from_logits",
    "entropy_l",
    "entropy_reg",
    "get_targets",
    "log_mean",
]
