Upload model

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "MambaModel"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig"
+    "AutoConfig": "configuration_mamba.MambaConfig",
+    "AutoModel": "modeling_mamba.MambaModel"
   },
   "d_model": 768,
   "fused_add_norm": true,
@@ -10,6 +14,7 @@
   "residual_in_fp32": true,
   "rms_norm": true,
   "ssm_cfg": {},
+  "torch_dtype": "float16",
   "transformers_version": "4.37.2",
   "vocab_size": 50277
 }

modeling_mamba.py

@@ -1,6 +1,7 @@
-from typing import Optional
+from typing import Optional, Tuple
 
 from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+import torch
 from transformers import GenerationMixin, PreTrainedModel
 from transformers.generation import TextStreamer
 
@@ -48,18 +49,18 @@ class MambaModelForCausalLM(MambaModel, GenerationMixin):
     def generate(
         self,
         input_ids,
-        max_length,
-        top_k=1,
-        top_p=0.0,
-        temperature=1.0,
-        return_dict_in_generate=False,
-        output_scores=False,
-        repetition_penalty=1.0,
-        eos_token_id=None,
-        teacher_outputs=None,
-        vocab_size=None,
-        cg=False,
-        enable_timing=False,
+        max_length: int = 2048,
+        top_k: int = 1,
+        top_p: float = 0.0,
+        temperature: float = 1.0,
+        return_dict_in_generate: bool = False,
+        output_scores: bool = False,
+        repetition_penalty: float = 1.0,
+        eos_token_id: Optional[int] = None,
+        teacher_outputs: Optional[torch.Tensor] = None,
+        vocab_size: Optional[int] = None,
+        cg: bool = False,
+        enable_timing: bool = False,
         streamer: Optional[TextStreamer] = None,
         **kwargs,
     ):
@@ -77,5 +78,5 @@ class MambaModelForCausalLM(MambaModel, GenerationMixin):
             vocab_size=vocab_size,
             cg=cg,
             enable_timing=enable_timing,
-            streamer = streamer,
+            streamer=streamer,
         )