base_IIXIV / fla /models /mamba /__init__.py
mainline777's picture
Duplicate from silx-ai/Quasar-Preview
41865df
Raw
History Blame Contribute Delete
474 Bytes
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from fla.models.mamba.configuration_mamba import MambaConfig
from fla.models.mamba.modeling_mamba import 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']