| import mamba_ssm | |
| from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig, PreTrainedModel | |
| mamba_config_defaults = mamba_ssm.models.config_mamba.MambaConfig() | |
| class MambaConfig(PretrainedConfig): | |
| model_type = "mamba" | |
| def __init__( | |
| self, | |
| d_model: int = mamba_config_defaults.d_model, | |
| fused_add_norm: bool = mamba_config_defaults.fused_add_norm, | |
| n_layer: int = mamba_config_defaults.n_layer, | |
| pad_vocab_size_multiple: int = mamba_config_defaults.pad_vocab_size_multiple, | |
| residual_in_fp32: bool = mamba_config_defaults.residual_in_fp32, | |
| rms_norm: bool = mamba_config_defaults.rms_norm, | |
| ssm_cfg: dict = mamba_config_defaults.ssm_cfg, | |
| vocab_size: int = mamba_config_defaults.vocab_size, | |
| **kwargs, | |
| ): | |
| self.d_model = d_model | |
| self.fused_add_norm = fused_add_norm | |
| self.n_layer = n_layer | |
| self.pad_vocab_size_multiple = pad_vocab_size_multiple | |
| self.residual_in_fp32 = residual_in_fp32 | |
| self.rms_norm = rms_norm | |
| self.ssm_cfg = ssm_cfg | |
| self.vocab_size = vocab_size | |
| super().__init__(**kwargs) | |