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

@@ -311,18 +311,9 @@ class MambaModel(MambaPreTrainedModel):
         )
 
 class MambaModelForCausalLM(MambaPreTrainedModel):
-    _tied_weights_keys = [
-        "lm_head.weight", # will remove this since it's a duplicate of backbone.embedding.weight
-    ]
+    _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config, **kwargs):
-        # super().__init__(config)
-        # self.backbone = MambaModel(config)
-        # self.vocab_size = config.vocab_size
-        # self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
-        # self.lm_head.weight = self.backbone.embedding.weight
-        # self.post_init()
-
         super().__init__(
             config,
             **kwargs,
@@ -338,7 +329,6 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
             bias=False,
         )
 
-        # self.lm_head.weight = self.backbone.embedding.weight
         self.post_init()
 
     def _tie_weights(self):
@@ -444,9 +434,9 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
             loss=loss,
         )
 
-    # def prepare_inputs_for_generation(self, input_ids, **kwargs):
-    #     model_inputs = {"input_ids": input_ids}
-    #     return model_inputs
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        model_inputs = {"input_ids": input_ids}
+        return model_inputs
 
 
 class MambaModelForSequenceClassification(MambaPreTrainedModel):