fix model

model.py

@@ -1,19 +1,21 @@
-import torch
-import torch.nn as nn
-from transformers import AutoModel
-
-class EmotionClassifier(nn.Module):
-    def __init__(self, model_name, num_labels=5):
-        super().__init__()
-        self.transformer = AutoModel.from_pretrained(model_name)
-        hidden_size = self.transformer.config.hidden_size
-        self.classifier = nn.Linear(hidden_size, num_labels)
-
-    def forward(self, input_ids, attention_mask):
-        outputs = self.transformer(
-            input_ids=input_ids,
-            attention_mask=attention_mask
-        )
-        cls_embeddings = outputs.last_hidden_state[:, 0, :]
-        logits = self.classifier(cls_embeddings)
-        return logits
+import torch
+import torch.nn as nn
+from transformers import AutoModel
+
+class EmotionClassifier(nn.Module):
+    def __init__(self, model_name="microsoft/deberta-v3-base"):
+        super().__init__()
+        self.backbone = AutoModel.from_pretrained(model_name)
+        hidden_size = self.backbone.config.hidden_size
+        # Must be named `out` because your trained weights use out.weight & out.bias
+        self.out = nn.Linear(hidden_size, 5)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+        last_hidden_state = outputs.last_hidden_state
+        cls_token = last_hidden_state[:, 0, :]
+        logits = self.out(cls_token)
+        return logits