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
 }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6504b24e9ba95e4a6bad94a346c849040623647d1a99a47f4f5e1cd32cbd9572
+size 259551392

modeling_mamba.py

@@ -0,0 +1,80 @@
+from typing import Optional
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.generation import TextStreamer
+
+from mamba_model.configuration_mamba import MambaConfig
+
+class MambaModel(PreTrainedModel):
+    config_class = MambaConfig
+
+    def __init__(
+        self,
+        config,
+        initializer_cfg=None,
+        device=None,
+        dtype=None,
+        **kwargs,
+    ):
+        super().__init__(
+            config,
+            **kwargs,
+        )
+
+        self.model = MambaLMHeadModel(
+            config,
+            initializer_cfg=initializer_cfg,
+            device=device,
+            dtype=dtype,
+        )
+
+    def forward(
+            self,
+            input_ids,
+            position_ids=None,
+            inference_params=None,
+            num_last_tokens=0,
+            **kwargs,
+        ):
+        return self.model.forward(
+            input_ids,
+            position_ids,
+            inference_params,
+            num_last_tokens
+        )
+
+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,
+        streamer: Optional[TextStreamer] = None,
+        **kwargs,
+    ):
+        return self.model.generate(
+            input_ids=input_ids,
+            max_length=max_length,
+            top_k=top_k,
+            top_p=top_p,
+            temperature=temperature,
+            return_dict_in_generate=return_dict_in_generate,
+            output_scores=output_scores,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=eos_token_id,
+            teacher_outputs=teacher_outputs,
+            vocab_size=vocab_size,
+            cg=cg,
+            enable_timing=enable_timing,
+            streamer = streamer,
+        )