| 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) |