Upload model.py

model.py

@@ -0,0 +1,19 @@
+from torch import nn
+
+class ClassificationModel(nn.Module):
+    def __init__(self, base_model):
+        super(ClassificationModel, self).__init__()
+        self.base_model = base_model
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, 8),
+            nn.LogSoftmax(dim=1)
+        )
+
+    def forward(self, input_ids, attention_mask):
+        hidden_states = self.base_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        cls_output = hidden_states[:, 0, :]
+        probs = self.classifier(cls_output)
+        return probs