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:1bd3ca62665de4bfabff9d443f87a11090a10e505c0ccb56e6f9ca495b6e05bd
-size 671027808
+oid sha256:699ed6f59fb948186f449c5031e0dc659d504c90d7e018302aa1e190cdb40220
+size 516567560

modeling_mamba.py

@@ -1,43 +1,38 @@
-import json
-import math
-import os
-from collections import namedtuple
-from dataclasses import dataclass
-from functools import partial
-from typing import Dict, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import transformers
 from einops import einsum, rearrange, repeat
-from torch import FloatTensor, Tensor, nn
-from transformers import GenerationMixin, PreTrainedModel
+from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import (
-    BaseModelOutput,
     BaseModelOutputWithPast,
-    CausalLMOutput,
-    ImageClassifierOutput,
+    CausalLMOutputWithPast,
     QuestionAnsweringModelOutput,
     SequenceClassifierOutput,
 )
-from trl import PreTrainedModelWrapper
+from transformers.modeling_utils import PreTrainedModel
 
 from .configuration_mamba import MambaConfig
 
 
-# Inspired by:
-# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L31
-# - https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L177
-# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/modules/mamba_simple.py#L31
-class MambaBlock(nn.Module):
-    def __init__(self, config: MambaConfig):
-        """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1].
-        Furthermore, in section E.2.2 of the paper, the authors describe the Mamba block as:
-            "[T]he Mamba block is simply the standard SwiGLU block with an extra conv → SSM path added."
-        """
+class MambaRMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5):
         super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        output = (
+            x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+        )
+        return output
 
+
+class Mamba(nn.Module):
+    def __init__(self, config: MambaConfig):
+        """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1]."""
+        super().__init__()
         self.config = config
 
         self.in_proj = nn.Linear(config.d_model, config.d_inner * 2, bias=config.bias)
@@ -62,9 +57,8 @@ class MambaBlock(nn.Module):
         A = repeat(torch.arange(1, config.d_state + 1), "n -> d n", d=config.d_inner)
         self.A_log = nn.Parameter(torch.log(A))
         self.D = nn.Parameter(torch.ones(config.d_inner))
-
         self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
-        # self.norm = RMSNorm(config.d_model)
+        # self.norm = MambaRMSNorm(config.d_model)
 
     def forward(self, x):
         """Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].
@@ -80,9 +74,10 @@ class MambaBlock(nn.Module):
             mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
 
         """
+
         (b, l, d) = x.shape
-        # x_copy = x # There was a separate class for residual, I deleted that part and added it here.
-        # x = self.norm(x)
+        x_copy = x  # There was a separate class for residual, I deleted that part and added it here.
+        x = self.norm(x)
         x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)
         (x, res) = x_and_res.split(
             split_size=[self.config.d_inner, self.config.d_inner], dim=-1
@@ -96,13 +91,9 @@ class MambaBlock(nn.Module):
 
         y = self.ssm(x)
 
-        y = y * F.silu(
-            res
-        )  # SwiGLU: Swish_β(xW + b) ⊗ (xV + c) => torch.kron(F.silu(xW + b), xV + c) => torch.kron(F.silu(res), y)
+        y = y * F.silu(res)
 
-        output = self.out_proj(y)  # output = self.out_proj(y) + x_copy
-
-        # "the Mamba block is simply the standard SwiGLU block with an extra 𝖼𝗈𝗇𝗏 → 𝖲𝖲𝖬 path added"
+        output = self.out_proj(y) + x_copy
 
         return output
 
@@ -177,21 +168,17 @@ class MambaBlock(nn.Module):
         # Discretize continuous parameters (A, B)
         # - A is discretized using zero-order hold (ZOH) discretization (see Section 2 Equation 4 in the Mamba paper [1])
         # - B is discretized using a simplified Euler discretization instead of ZOH. From a discussion with authors:
-        #   "A is the more important term and the performance doesn't change much with the simplification on B"
-        deltaA = torch.exp(einsum(delta, A, "b l d_in, d_in n -> b l d_in n"))
-        deltaB_u = einsum(delta, B, u, "b l d_in, b l n, b l d_in -> b l d_in n")
+        #   "A is the more important term and the performance doesn't change much with the simplication on B"
+        deltaA = torch.exp(einsum(delta, A, "b l d_in, d_in n -> b d_in l n"))
+        deltaB_u = einsum(delta, B, u, "b l d_in, b l n, b l d_in -> b d_in l n")
 
         # Perform selective scan (see scan_SSM() in The Annotated S4 [2])
-        # Note that the below is sequential, while the official implementation does a much faster parallel scan that
-        # is additionally hardware-aware (like FlashAttention).
         x = torch.zeros((b, d_in, n), device=deltaA.device)
         ys = []
-
         for i in range(l):
-            x = deltaA[:, i] * x + deltaB_u[:, i]
+            x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
             y = einsum(x, C[:, i, :], "b d_in n, b n -> b d_in")
             ys.append(y)
-
         y = torch.stack(ys, dim=1)  # shape (b, l, d_in)
 
         y = y + u * D
@@ -199,395 +186,163 @@ class MambaBlock(nn.Module):
         return y
 
 
-# Inspired by:
-# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L19
-# - https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L328
-# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/ops/triton/layernorm.py#L481
-class RMSNorm(nn.Module):
-    def __init__(self, d_model: int, eps: float = 1e-5):
-        super().__init__()
-
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(d_model))
-
-    def forward(self, x):
-        output = (
-            x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
-        )
-
-        return output
-
-
-class ResidualBlock(
-    nn.Module
-):  # Copied and modified from https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L143
+class Block(nn.Module):
     def __init__(self, config: MambaConfig):
-        """Simple block wrapping Mamba block with normalization and residual connection."""
+        """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1]."""
         super().__init__()
+        self.config = config
 
-        # self.args = args
-        self.mixer = MambaBlock(config)
-        self.norm = RMSNorm(config.d_model)
-        # self.norm = partial(
-        #     nn.LayerNorm if not config.rms_norm else RMSNorm, eps=config.norm_epsilon,
-        # )
-
-        # if config.rms_norm:
-        #     self.norm = RMSNorm(config.d_model, eps=config.norm_epsilon)
-
-        # else:
-        #     self.norm = nn.LayerNorm(config.d_model, eps=config.norm_epsilon)
-
-    def forward(self, x):
-        """
-        Args:
-            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)
-
-        Returns:
-            output: shape (b, l, d)
-
-        Official Implementation:
-            Block.forward(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L297
-
-            Note: the official repo chains residual blocks that look like
-                [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> ...
-            where the first Add is a no-op. This is purely for performance reasons as this
-            allows them to fuse the Add->Norm.
-
-            We instead implement our blocks as the more familiar, simpler, and numerically equivalent
-                [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> ....
-
-        """
-        output = self.mixer(self.norm(x)) + x
-
-        return output
-
-
-# Inspired by:
-# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L181
-# class MambaPretrainedModel(PreTrainedModel, nn.Module):
-class MambaPretrainedModel(PreTrainedModel):
-    r"""
-    Base class for all models.
-
-    [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
-    downloading and saving models as well as a few methods common to all models to:
-
-        - resize the input embeddings,
-        - prune heads in the self-attention heads.
-
-    Class attributes (overridden by derived classes):
-
-        - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
-          for this model architecture.
-        - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model,
-          taking as arguments:
+        self.mixer = Mamba(config)
+        self.norm = MambaRMSNorm(config.d_model)
 
-            - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint.
-            - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model.
-            - **path** (`str`) -- A path to the TensorFlow checkpoint.
 
-        - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived
-          classes of the same architecture adding modules on top of the base model.
-        - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization.
-        - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
-          models, `pixel_values` for vision models and `input_values` for speech models).
-    """
+class MambaBlock(Block):
+    pass
 
-    config_class = MambaConfig  # TODO: Build on top of MambaConfig?
-    # base_model_prefix = "backbone"
-    base_model_prefix = "mamba"
-    main_input_name = "input_ids"
-    model_tags = None
 
-    _auto_class = None
-    _no_split_modules = ["MambaBlock"]
-    _skip_keys_device_placement = None
-    _keep_in_fp32_modules = None
-
-    # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing
-    # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings.
-    _keys_to_ignore_on_load_missing = None
-    # a list of `re` patterns of `state_dict` keys that should be removed from the list of
-    # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary
-    # warnings.
-    _keys_to_ignore_on_load_unexpected = None
-    # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't
-    # trained, but which are either deterministic or tied variables)
-    _keys_to_ignore_on_save = None
-    # a list of `state_dict` keys that are potentially tied to another key in the state_dict.
-    _tied_weights_keys = None
-
-    is_parallelizable = False
+class MambaPreTrainedModel(PreTrainedModel):
+    config_class = MambaConfig
+    base_model_prefix = "backbone"
     supports_gradient_checkpointing = True
+    _no_split_modules = ["MambaBlock"]
 
-    # Flash Attention 2 support
-    _supports_flash_attn_2 = False
-
-    # SDPA support
-    _supports_sdpa = False
-
-    # Has support for a `Cache` instance as `past_key_values`
-    _supports_cache_class = False
-
-    def __init__(self, *inputs, **kwargs):
-        super().__init__(*inputs, **kwargs)
-
-    # https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L54
-    def _init_weights(
-        self,
-        module,
-        initializer_range=0.02,  # Now only used for embedding layer.
-        rescale_prenorm_residual=True,
-        n_residuals_per_layer=1,  # Change to 2 if we have MLP
-    ):
-        if isinstance(module, nn.Linear):
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
-                if not getattr(module.bias, "_no_reinit", False):
-                    nn.init.zeros_(module.bias)
-
+                module.bias.data.zero_()
         elif isinstance(module, nn.Embedding):
-            nn.init.normal_(module.weight, std=initializer_range)
-
-        if rescale_prenorm_residual:
-            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
-            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
-            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
-            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
-            #
-            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
-            for name, p in module.named_parameters():
-                if name in [
-                    "out_proj.weight",
-                    "fc2.weight",
-                ]:
-                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
-                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
-                    # We need to reinit p since this code could be called multiple times
-                    # Having just p *= scale would repeatedly scale it down
-                    nn.init.kaiming_uniform_(p, a=math.sqrt(5))
-                    with torch.no_grad():
-                        p /= math.sqrt(n_residuals_per_layer * self.config.n_layer)
-
-    # def _set_gradient_checkpointing(self, module, value=False):
-    #     if isinstance(module, GPT2Model):
-    #         module.gradient_checkpointing = value
-
-# Inspired by:
-# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L86
-class MambaModel(MambaPretrainedModel):
-    def __init__(self, config: MambaConfig = MambaConfig(), **kwargs) -> None:
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class MambaModel(MambaPreTrainedModel):
+    def __init__(self, config: MambaConfig):
         """Full Mamba model.
         Mamba model decoder consisting of *config.n_layer* layers. Each layer is a [`MambaBlock`]
+
         Args:
             config: MambaConfig
         """
-        super().__init__(
-            config,
-            **kwargs,
-        )
-
-        self.embedding = nn.Embedding(
-            num_embeddings=self.config.vocab_size,
-            embedding_dim=self.config.d_model,
-        )
-
-        self.layers = nn.ModuleList(
-            [ResidualBlock(self.config) for _ in range(self.config.n_layer)]
-        )
-        self.norm_f = RMSNorm(d_model=self.config.d_model)
-        # self.norm_f = (nn.LayerNorm if not self.config.rms_norm else RMSNorm)(
-        #     # self.config.d_model, eps=self.config.norm_epsilon, **factory_kwargs
-        #     self.config.d_model, eps=self.config.norm_epsilon,
-        # )
+        super().__init__(config)
+        self.config = config
 
-        # self.gradient_checkpointing = False
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.layers = nn.ModuleList([MambaBlock(config) for _ in range(config.n_layer)])
+        self.norm_f = MambaRMSNorm(config.d_model)
 
-        # Initialize weights and apply final processing
+        self.gradient_checkpointing = False
         self.post_init()
 
-    # def _init_weights(self, module):
-    #     std = 0.02
+    def get_input_embeddings(self):
+        return self.embedding
 
-    #     if isinstance(module, (nn.Linear, nn.Conv1d)):
-    #         module.weight.data.normal_(mean=0.0, std=std)
-
-    #         if module.bias is not None:
-    #             module.bias.data.zero_()
-
-    #     elif isinstance(module, nn.Embedding):
-    #         module.weight.data.normal_(mean=0.0, std=std)
-
-    #         if module.padding_idx is not None:
-    #             module.weight.data[module.padding_idx].zero_()
-
-    # def get_input_embeddings(self):
-    #     return self.embed_out
-
-    # def set_input_embeddings(self, value):
-    #     self.embed_out = value
+    def set_input_embeddings(self, value):
+        self.embedding = value
 
     def forward(
         self,
         input_ids: torch.LongTensor = None,
-        output_hidden_states=False,
         return_dict: Optional[bool] = None,
-        **kwargs,
-        # ) -> BaseModelOutput:
     ) -> Union[Tuple, BaseModelOutputWithPast]:
-        batch_size = input_ids.shape[0]
-        hidden_size = self.config.d_model
-        hidden_states: Tuple[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
-
-        last_hidden_state = self.embedding(input_ids)
-        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,)
-
+        x = self.embedding(input_ids)
+        all_hidden_states = list()
         for layer in self.layers:
-            last_hidden_state = layer(last_hidden_state)
-            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,)
-
-        last_hidden_state = self.norm_f(last_hidden_state)
-        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,)
-
-        assert (
-            len(hidden_states) == self.config.n_layer + 2
-        ), f"{len(hidden_states)} != {self.config.n_layer + 2}"
-
-        # return BaseModelOutput(
+            x = layer(x)
+            all_hidden_states.append(x)
+
+        hidden_states = self.norm_f(x)
+
         return BaseModelOutputWithPast(
-            hidden_states=hidden_states if output_hidden_states else None,
-            last_hidden_state=last_hidden_state,
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
         )
 
 
-# Influences:
-# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L238
-# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L176
-# class MambaModelForCausalLM(MambaModel, GenerationMixin):
-# class MambaModelForCausalLM(PreTrainedModel, GenerationMixin):
-# class MambaLMHeadModel(MambaPretrainedModel, GenerationMixin):
-class MambaLMHeadModel(MambaPretrainedModel):
-    _tied_weights_keys = [
-        "backbone.embedding.weight",
-        "lm_head.weight",
-    ]
+class MambaModelForCausalLM(MambaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(
-        self,
-        config: MambaConfig = MambaConfig(),
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            config,
-            **kwargs,
-        )
+    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
+        self.post_init()
 
-        self.backbone = MambaModel(
-            config=self.config,
-        )
+    # def get_input_embeddings(self):
+    #     return self.model.embedding
 
-        self.lm_head = nn.Linear(
-            in_features=self.config.d_model,
-            out_features=self.config.vocab_size,
-            bias=False,
-        )
+    # def set_input_embeddings(self, value):
+    #     self.model.embedding = value
 
-        #         # self.head.weight = self.backbone.embedding.weight # TODO: there's some logic in GenerationMix that does this
+    # def get_output_embeddings(self):
+    #     return self.lm_head
 
-        # Initialize weights and apply final processing
-        self.post_init()
+    # def set_output_embeddings(self, new_embeddings):
+    #     self.lm_head = new_embeddings
 
-    def forward(
-        self, input_ids, output_hidden_states=False, **kwargs
-    ) -> CausalLMOutput:
-        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
+    # 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,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
         outputs = self.backbone(
             input_ids=input_ids,
-            output_hidden_states=output_hidden_states,
+            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,
         )
 
-        last_hidden_state = outputs.last_hidden_state
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        model_inputs = {"input_ids": input_ids}
+        return model_inputs
 
-        logits: torch.FloatTensor[batch_size, sequence_length, vocab_size] = (
-            self.lm_head(
-                last_hidden_state,
-            )
-        )
 
-        return CausalLMOutput(
-            hidden_states=outputs.hidden_states if output_hidden_states else None,
-            logits=logits,
-        )
+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 prepare_inputs_for_generation(
-        self, input_ids, attention_mask=None, **model_kwargs
-    ):
-        return {
-            "input_ids": input_ids,
-        }
-
-
-# class MultimodalMambaModelForCausalLMWithValueHead(PreTrainedModelWrapper):
-#     lm_head_namings: Tuple[str, str] = ("lm_head", "embed_out")
-#     transformers_parent_class: transformers.PreTrainedModel = transformers.AutoModelForCausalLM
-
-#     # def __init__(
-#     #     self,
-#     #     config: MultimodalMambaConfig = MultimodalMambaConfig(),
-#     #     **kwargs,
-#     # ) -> None:
-#     #     super().__init__(
-#     #         config,
-#     #         **kwargs,
-#     #     )
-
-#     #     self.model = MultimodalMambaModelForCausalLM(
-#     #         config=config,
-#     #     )
-
-#     #     self.value_head = nn.Linear(
-#     #         in_features=config.embedding_dim,
-#     #         out_features=1,
-#     #         bias=False,
-#     #     )
-
-#     # def forward(
-#     #     self, input_ids, output_hidden_states=False, **kwargs
-#     # ) -> CausalLMOutput:
-#     #     outputs = self.model(
-#     #         input_ids=input_ids,
-#     #         output_hidden_states=output_hidden_states,
-#     #     )
-
-#     #     last_hidden_state = outputs.last_hidden_state
-
-#     #     value: torch.FloatTensor[batch_size, sequence_length, 1] = self.value_head(
-#     #         last_hidden_state,
-#     #     )
-
-#     #     return CausalLMOutput(
-#     #         hidden_states=outputs.hidden_states if output_hidden_states else None,
-#     #         logits=outputs.logits,
-#     #         value=value,
-#     #     )
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> SequenceClassifierOutput:
+        pass