File size: 361 Bytes
9a27a9d |
1 2 3 4 5 6 7 8 9 |
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) |