Upload sentiment classifier trained on Amazon Reviews

sentiment_classifier.py

@@ -1,10 +1,11 @@
 """Sentiment classifier for text classification."""
 
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 import torch
 import torch.nn as nn
 from transformers import AutoModel, PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
 
 # Handle imports for both local usage and HuggingFace Hub
 try:
@@ -104,8 +105,9 @@ class SentimentClassifier(PreTrainedModel):
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
         labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
         **kwargs,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> Union[SequenceClassifierOutput, Dict[str, torch.Tensor]]:
         """
         Forward pass for classification.
 
@@ -113,11 +115,14 @@ class SentimentClassifier(PreTrainedModel):
             input_ids: Input token IDs [batch_size, seq_len].
             attention_mask: Attention mask [batch_size, seq_len].
             labels: Ground truth sentiment labels [batch_size].
+            return_dict: Whether to return a SequenceClassifierOutput or dict.
             **kwargs: Additional arguments.
 
         Returns:
-            Dictionary containing loss and logits.
+            SequenceClassifierOutput or dictionary containing loss and logits.
         """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
         # Encode with transformer
         outputs = self.encoder(
             input_ids=input_ids,
@@ -140,10 +145,16 @@ class SentimentClassifier(PreTrainedModel):
             loss_fct = nn.CrossEntropyLoss(weight=self.class_weights)
             loss = loss_fct(logits, labels)
 
-        return {
-            "loss": loss,
-            "logits": logits,
-        }
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=None,
+        )
 
     def predict(
         self,
@@ -164,7 +175,7 @@ class SentimentClassifier(PreTrainedModel):
 
         with torch.no_grad():
             outputs = self.forward(input_ids, attention_mask)
-            logits = outputs["logits"]
+            logits = outputs.logits
             label_predictions = torch.argmax(logits, dim=-1)
 
         return label_predictions
@@ -188,7 +199,7 @@ class SentimentClassifier(PreTrainedModel):
 
         with torch.no_grad():
             outputs = self.forward(input_ids, attention_mask)
-            logits = outputs["logits"]
+            logits = outputs.logits
             probabilities = torch.softmax(logits, dim=-1)
 
         return probabilities