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,
@@ -15,6 +19,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:287cad4048030ae246aeda26c0e703b838c50422fe89f19099298c034b25e7b5
-size 516565384
+oid sha256:699ed6f59fb948186f449c5031e0dc659d504c90d7e018302aa1e190cdb40220
+size 516567560

modeling_mamba.py

@@ -8,8 +8,7 @@ from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
     CausalLMOutputWithPast,
-    QuestionAnsweringModelOutput,
-    SequenceClassifierOutput,
+    SequenceClassifierOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
 
@@ -320,9 +319,9 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
         **kwargs,
     ) -> CausalLMOutputWithPast:
         batch_size = input_ids.shape[0]
+        output_hidden_states = output_hidden_states or self.config.output_hidden_states
         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,
@@ -337,7 +336,7 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
             )
         )
 
-        if labels:
+        if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
             loss_fct = CrossEntropyLoss()
@@ -364,17 +363,75 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
         }
 
 
-class MambaModelForSequenceClassification(MambaPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = MambaModel(config)
-        # self.classifier = nn.Linear(config.d_model, config.num_labels)
-        # self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> SequenceClassifierOutput:
-        pass
+# class MambaModelForSequenceClassification(MambaModelForCausalLM):
+#     def __init__(
+#             self,
+#             config,
+#             id2label={0: "NEGATIVE", 1: "POSITIVE"},
+#             label2id={"NEGATIVE": 0, "POSITIVE": 1},
+#             num_labels=2,
+#             **kwargs,
+#         ):
+#         super().__init__(
+#             config,
+#             **kwargs,
+#         )
+
+#         self.id2label = id2label
+#         self.label2id = label2id
+#         self.num_labels = num_labels # TODO: config.num_labels
+
+#         self.score = nn.Linear(
+#             in_features=self.config.vocab_size,
+#             out_features=self.num_labels,
+#             bias=False,
+#         )
+
+#     def forward(
+#         self,
+#         input_ids: Optional[torch.Tensor] = None,
+#         labels: Optional[torch.Tensor] = None,
+#         output_hidden_states=False,
+#         **kwargs,
+#     ) -> SequenceClassifierOutputWithPast:
+#         batch_size = input_ids.shape[0]
+#         hidden_size = self.config.vocab_size
+#         hidden_states: Tuple[
+#             torch.Tensor[(batch_size, sequence_length, hidden_size)]
+#         ] = ()
+#         num_labels = self.num_labels # TODO: config.num_labels
+#         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 = super().forward(
+#             input_ids=input_ids,
+#             labels=None,
+#             output_hidden_states=output_hidden_states,
+#             **kwargs,
+#         )
+
+#         last_hidden_state = outputs.logits
+#         assert last_hidden_state.shape == (
+#             batch_size,
+#             sequence_length,
+#             hidden_size,
+#         ), f"{last_hidden_state.shape} != {(batch_size, sequence_length, hidden_size)}"
+#         hidden_states += (last_hidden_state,)
+
+#         logits: torch.FloatTensor[batch_size, num_labels] = self.score(
+#             last_hidden_state[:, -1, :] # TODO: Check if this makes sense
+#         )
+
+#         if labels is not None:
+#             loss_fct = CrossEntropyLoss()
+#             loss = loss_fct(logits, labels)
+        
+#         else:
+#             loss = None
+
+#         return SequenceClassifierOutputWithPast(
+#             loss=loss,
+#             logits=logits,
+#             hidden_states=hidden_states,
+#         )