Update hf_mamba_classification.py

hf_mamba_classification.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from torch.nn import CrossEntropyLoss
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.models.mamba.modeling_mamba import (
     MambaPreTrainedModel,
     MambaModel,
@@ -44,7 +44,9 @@ 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
 
 
@@ -149,8 +151,32 @@ class MambaForSequenceClassification(MambaPreTrainedModel):
             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))
+        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 not return_dict:
             output = (pooled_logits,) + mamba_outputs[1:]