Update hf_mamba_classification.py

hf_mamba_classification.py

@@ -1,14 +1,13 @@
 import torch
 from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn import CrossEntropyLoss
 from transformers.models.mamba.modeling_mamba import (
-    MambaPreTrainedModel, 
+    MambaPreTrainedModel,
     MambaModel,
-    MambaCache, 
+    MambaCache,
     MAMBA_INPUTS_DOCSTRING,
     MAMBA_START_DOCSTRING,
 )
-from transformers.modeling_outputs import SequenceClassifierOutputWithPast
 from typing import List, Optional, Tuple, Union
 from transformers.utils import (
     ModelOutput,
@@ -45,33 +44,21 @@ class MambaSequenceClassifierOutput(ModelOutput):
 
     loss: Optional[torch.FloatTensor] = None
     logits: torch.FloatTensor = None
-    # cache_params: Optional[MambaCache] = None,
     cache_params: Optional[List[torch.FloatTensor]] = None
-    # cache_params: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
     hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    
-    
+
+
 class MambaClassificationHead(nn.Module):
     """Head for sentence-level classification tasks."""
 
     def __init__(self, config):
         super().__init__()
-        # self.activation = ACT2FN[config.hidden_act]
-        # self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        # self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.out_proj = nn.Linear(config.hidden_size, config.num_labels, bias=False)
-
-        # module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
         self.out_proj.weight.data.normal_(mean=0.0, std=config.initializer_range)
 
         self.config = config
 
     def forward(self, features, **kwargs):
-        # x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
-        # x = self.dropout(x)
-        # x = self.dense(x)
-        # x = self.activation(x)
-        # x = self.dropout(x)
         x = features
         x = self.out_proj(x)
         return x
@@ -86,19 +73,15 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        # self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
         self.backbone = MambaModel(config)
-        # self.classifier = MambaClassificationHead(config)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels, bias=False)
-        # self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
-        
-        for param in self.base_model.parameters():
-            param.requires_grad = False
+        self.classifier = MambaClassificationHead(config)
 
         # Initialize weights and apply final processing
         self.post_init()
 
-    @add_start_docstrings_to_model_forward(MAMBA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_start_docstrings_to_model_forward(
+        MAMBA_INPUTS_DOCSTRING.format("batch_size, sequence_length")
+    )
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC,
         output_type=MambaSequenceClassifierOutput,
@@ -122,19 +105,9 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
             If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
             If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
         """
-        # use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # if inputs_embeds is None:
-        #     inputs_embeds = self.backbone.embeddings(input_ids)
-
-        # if self.backbone.gradient_checkpointing and self.training and use_cache:
-        #     use_cache = False
-
-        # if cache_params is None and use_cache:
-        #     cache_params = MambaCache(
-        #         self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype
-        #     )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         mamba_outputs = self.backbone(
             input_ids,
@@ -154,13 +127,15 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
         assert (
             self.config.pad_token_id is not None or batch_size == 1
         ), "Cannot handle batch sizes > 1 if no padding token is defined."
-        
+
         if self.config.pad_token_id is None:
             sequence_lengths = -1
         else:
             if input_ids is not None:
                 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
-                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                sequence_lengths = (
+                    torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                )
                 sequence_lengths = sequence_lengths % input_ids.shape[-1]
                 sequence_lengths = sequence_lengths.to(logits.device)
             else:
@@ -170,34 +145,13 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
                     "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
                 )
 
-        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
-        loss = None
-        if labels is not None:
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = "regression"
-                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
-                    self.config.problem_type = "single_label_classification"
-                else:
-                    self.config.problem_type = "multi_label_classification"
-
-            if self.config.problem_type == "regression":
-                loss_fct = MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(pooled_logits, labels)
-            elif self.config.problem_type == "single_label_classification":
-                loss_fct = CrossEntropyLoss()
-                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
-            elif self.config.problem_type == "multi_label_classification":
-                loss_fct = BCEWithLogitsLoss()
-                loss = loss_fct(pooled_logits, labels)
-        
-        # if use_cache:
-        #     cache_params.seqlen_offset += inputs_embeds.shape[1]
-                
+        pooled_logits = logits[
+            torch.arange(batch_size, device=logits.device), sequence_lengths
+        ]
+
+        loss_fct = CrossEntropyLoss()
+        loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+
         if not return_dict:
             output = (pooled_logits,) + mamba_outputs[1:]
             return ((loss,) + output) if loss is not None else output
@@ -207,4 +161,4 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
             logits=pooled_logits,
             cache_params=mamba_outputs.cache_params,
             hidden_states=mamba_outputs.hidden_states,
-        )
+        )