subclass from transformers

markrogersjr · Mar 7, 2024 · 240a343 · 240a343
1 parent 34076d6
commit 240a343
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/mamba_ssm/models/config_mamba.py b/mamba_ssm/models/config_mamba.py
@@ -1,8 +1,10 @@
 from dataclasses import dataclass, field
 
+from transformers import PretrainedConfig
+
 
 @dataclass
-class MambaConfig:
+class MambaConfig(PretrainedConfig):
 
     d_model: int = 2560
     n_layer: int = 64

diff --git a/mamba_ssm/models/mixer_seq_simple.py b/mamba_ssm/models/mixer_seq_simple.py
@@ -9,6 +9,8 @@
 
 import torch
 import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
 
 from mamba_ssm.models.config_mamba import MambaConfig
 from mamba_ssm.modules.mamba_simple import Mamba, Block
@@ -173,7 +175,7 @@ def forward(self, input_ids, inference_params=None):
         return hidden_states
 
 
-class MambaLMHeadModel(nn.Module, GenerationMixin):
+class MambaLMHeadModel(PreTrainedModel, GenerationMixin):
 
     def __init__(
         self,
@@ -193,7 +195,8 @@ def __init__(
         pad_vocab_size_multiple = config.pad_vocab_size_multiple
         factory_kwargs = {"device": device, "dtype": dtype}
 
-        super().__init__()
+        PreTrainedModel.__init__(self, config)
+        GenerationMixin.__init__(self)
         if vocab_size % pad_vocab_size_multiple != 0:
             vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)
         self.backbone = MixerModel(
@@ -235,7 +238,6 @@ def forward(self, input_ids, position_ids=None, inference_params=None, num_last_
         if num_last_tokens > 0:
             hidden_states = hidden_states[:, -num_last_tokens:]
         lm_logits = self.lm_head(hidden_states)
-        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
         return CausalLMOutput(logits=lm_logits)
 
     @classmethod