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

@@ -230,48 +230,26 @@ class MambaModel(MambaPreTrainedModel):
         self.config = config
 
         self.embedding = nn.Embedding(config.vocab_size, config.d_model)
-        self.layers = nn.ModuleList([MambaBlock(config, layer_idx) for layer_idx in range(config.n_layer)])
+        self.layers = nn.ModuleList(
+            [MambaBlock(config, layer_idx) for layer_idx in range(config.n_layer)]
+        )
         self.norm_f = MambaRMSNorm(config.d_model)
 
         self.gradient_checkpointing = False
         self.post_init()
 
-    # def get_input_embeddings(self):
-    #     return self.embedding
-
-    # def set_input_embeddings(self, value):
-    #     self.embedding = value
-
-    # def forward(
-    #     self,
-    #     input_ids: torch.LongTensor = None,
-    #     **kwargs,
-    # ) -> Union[Tuple, BaseModelOutputWithPast]:
-    #     x = self.embedding(input_ids)
-    #     all_hidden_states = list()
-    #     for layer in self.layers:
-    #         x = layer(x)
-    #         all_hidden_states.append(x)
-
-    #     hidden_states = self.norm_f(x)
-
-    #     return BaseModelOutputWithPast(
-    #         last_hidden_state=hidden_states,
-    #         hidden_states=all_hidden_states,
-    #     )
-
     def forward(
         self,
         input_ids: torch.LongTensor = None,
         output_hidden_states=False,
         return_dict: Optional[bool] = None,
         **kwargs,
-        # ) -> BaseModelOutput:
-    # ) -> Union[Tuple, BaseModelOutputWithPast]:
     ) -> BaseModelOutputWithPast:
         batch_size = input_ids.shape[0]
         hidden_size = self.config.d_model
-        hidden_states: Tuple[torch.Tensor[(batch_size, sequence_length, hidden_size)]] = ()
+        hidden_states: Tuple[
+            torch.Tensor[(batch_size, sequence_length, hidden_size)]
+        ] = ()
         sequence_length = input_ids.shape[1]
         output_hidden_states = output_hidden_states or self.config.output_hidden_states
 
@@ -304,12 +282,12 @@ class MambaModel(MambaPreTrainedModel):
             len(hidden_states) == self.config.n_layer + 2
         ), f"{len(hidden_states)} != {self.config.n_layer + 2}"
 
-        # return BaseModelOutput(
         return BaseModelOutputWithPast(
             hidden_states=hidden_states if output_hidden_states else None,
             last_hidden_state=last_hidden_state,
         )
 
+
 class MambaModelForCausalLM(MambaPreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -334,62 +312,6 @@ class MambaModelForCausalLM(MambaPreTrainedModel):
     def _tie_weights(self):
         self.lm_head.weight = self.backbone.embedding.weight
 
-    # def get_input_embeddings(self):
-    #     return self.model.embedding
-
-    # def set_input_embeddings(self, value):
-    #     self.model.embedding = value
-
-    # def get_output_embeddings(self):
-    #     return self.lm_head
-
-    # def set_output_embeddings(self, new_embeddings):
-    #     self.lm_head = new_embeddings
-
-    # def set_decoder(self, decoder):
-    #     self.model = decoder
-
-    # 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,