Upload model

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "MambaModelForCausalLM"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig"
+    "AutoConfig": "configuration_mamba.MambaConfig",
+    "AutoModelForCausalLM": "modeling_mamba.MambaModelForCausalLM"
   },
   "bias": false,
   "conv_bias": true,
@@ -14,6 +18,7 @@
   "model_type": "mamba",
   "n_layer": 24,
   "pad_vocab_size_multiple": 8,
+  "torch_dtype": "float32",
   "transformers_version": "4.37.2",
   "vocab_size": 50280
 }

modeling_mamba.py

@@ -236,16 +236,16 @@ class MambaModel(MambaPreTrainedModel):
         self.gradient_checkpointing = False
         self.post_init()
 
-    def get_input_embeddings(self):
-        return self.embedding
+    # def get_input_embeddings(self):
+    #     return self.embedding
 
-    def set_input_embeddings(self, value):
-        self.embedding = value
+    # def set_input_embeddings(self, value):
+    #     self.embedding = value
 
     def forward(
         self,
         input_ids: torch.LongTensor = None,
-        return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         x = self.embedding(input_ids)
         all_hidden_states = list()
@@ -297,6 +297,7 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         outputs = self.backbone(
             input_ids=input_ids,