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
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:699ed6f59fb948186f449c5031e0dc659d504c90d7e018302aa1e190cdb40220
-size 516567560
+oid sha256:1bd3ca62665de4bfabff9d443f87a11090a10e505c0ccb56e6f9ca495b6e05bd
+size 671027808

modeling_mamba.py

@@ -313,12 +313,29 @@ class MambaModel(MambaPreTrainedModel):
 class MambaModelForCausalLM(MambaPreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config):
-        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
+    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,
+        )
+
+        self.backbone = MambaModel(
+            config=self.config,
+        )
+
+        self.lm_head = nn.Linear(
+            in_features=self.config.d_model,
+            out_features=self.config.vocab_size,
+            bias=False,
+        )
+
         self.post_init()
 
     # def get_input_embeddings(self):
@@ -339,47 +356,91 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
     # def get_decoder(self):
     #     return self.model
 
+    # def forward(
+    #     self,
+    #     input_ids: torch.LongTensor = None,
+    #     labels: Optional[torch.LongTensor] = None,
+    #     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,
+    #         return_dict=return_dict,
+    #     )
+    #     hidden_states = outputs[0]
+    #     logits = self.lm_head(hidden_states)
+    #     logits = logits.float()
+    #     loss = None
+
+    #     if labels is not None:
+    #         shift_logits = logits[..., :-1, :].contiguous()
+    #         shift_labels = labels[..., 1:].contiguous()
+    #         loss_fct = CrossEntropyLoss()
+    #         shift_logits = shift_logits.view(-1, self.config.vocab_size)
+    #         shift_labels = shift_labels.view(-1)
+
+    #         shift_labels = shift_labels.to(shift_logits.device)
+    #         loss = loss_fct(shift_logits, shift_labels)
+
+    #     if not return_dict:
+    #         output = (logits,) + outputs[1:]
+    #         return (loss,) + output if loss is not None else output
+
+    #     return CausalLMOutputWithPast(
+    #         loss=loss,
+    #         logits=logits,
+    #         hidden_states=outputs.hidden_states,
+    #     )
+
     def forward(
         self,
-        input_ids: torch.LongTensor = None,
+        input_ids,
         labels: Optional[torch.LongTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
+        output_hidden_states=False,
         **kwargs,
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
+    ) -> CausalLMOutputWithPast:
+        batch_size = input_ids.shape[0]
+        sequence_length = input_ids.shape[1]
+        vocab_size = self.config.vocab_size
+        output_hidden_states = output_hidden_states or self.config.output_hidden_states
+
         outputs = self.backbone(
             input_ids=input_ids,
-            return_dict=return_dict,
+            output_hidden_states=output_hidden_states,
         )
-        hidden_states = outputs[0]
-        logits = self.lm_head(hidden_states)
-        logits = logits.float()
-        loss = None
 
-        if labels is not None:
+        last_hidden_state = outputs.last_hidden_state
+
+        logits: torch.FloatTensor[batch_size, sequence_length, vocab_size] = (
+            self.lm_head(
+                last_hidden_state,
+            )
+        )
+
+        if labels:
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
             loss_fct = CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_logits = shift_logits.view(-1, vocab_size)
             shift_labels = shift_labels.view(-1)
 
             shift_labels = shift_labels.to(shift_logits.device)
             loss = loss_fct(shift_logits, shift_labels)
 
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
+        else:
+            loss = None
 
         return CausalLMOutputWithPast(
-            loss=loss,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
             logits=logits,
-            hidden_states=outputs.hidden_states,
+            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):