custom_modules.py

from transformers import AutoConfig, AutoModel
from transformers.models.auto.configuration_auto import CONFIG_MAPPING

from .model import EmbeddingMoEConfig, EmbeddingMoE

# Register your model
CONFIG_MAPPING.register("embedding_moe", EmbeddingMoEConfig)
AutoConfig.register("embedding_moe", EmbeddingMoEConfig)
AutoModel.register("embedding_moe", EmbeddingMoE)