File size: 501 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 |
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from fla.models.mamba.configuration_mamba import MambaConfig
from fla.models.mamba.modeling_mamba import MambaBlock, MambaForCausalLM, MambaModel
AutoConfig.register(MambaConfig.model_type, MambaConfig, exist_ok=True)
AutoModel.register(MambaConfig, MambaModel, exist_ok=True)
AutoModelForCausalLM.register(MambaConfig, MambaForCausalLM, exist_ok=True)
__all__ = ['MambaConfig', 'MambaForCausalLM', 'MambaModel', 'MambaBlock']
|