app fixes_v5

app.py

@@ -83,42 +83,23 @@ class DebertPaperClassifierV5(torch.nn.Module):
         else:
             self.loss_fct = torch.nn.BCEWithLogitsLoss()
 
-        def forward(self, input_ids, attention_mask, labels=None):
-            outputs = self.deberta(
-                input_ids=input_ids,
-                attention_mask=attention_mask
-            )
-            logits = self.classifier(outputs.last_hidden_state[:, 0, :])
-            loss = None
-            if labels is not None:
-                loss = self.loss_fct(logits, labels)
-            return (loss, logits) if loss is not None else logits
-
-    def _init_weights(self):
-        for module in self.classifier.modules():
-            if isinstance(module, torch.nn.Linear):
-                module.weight.data.normal_(mean=0.0, std=0.02)
-                if module.bias is not None:
-                    module.bias.data.zero_()
-
-def forward(self, 
-                input_ids, 
-                attention_mask,
-                labels=None,
-               ):
+    def forward(self, input_ids, attention_mask, labels=None):
         outputs = self.deberta(
             input_ids=input_ids,
             attention_mask=attention_mask
         )
-
-        cls_output = outputs.last_hidden_state[:, 0, :]
-        logits = self.classifier(cls_output)
-
+        logits = self.classifier(outputs.last_hidden_state[:, 0, :])
         loss = None
         if labels is not None:
             loss = self.loss_fct(logits, labels)
-
         return (loss, logits) if loss is not None else logits
+
+    def _init_weights(self):
+        for module in self.classifier.modules():
+            if isinstance(module, torch.nn.Linear):
+                module.weight.data.normal_(mean=0.0, std=0.02)
+                if module.bias is not None:
+                    module.bias.data.zero_()
     
 @st.cache_resource
 def load_model(test=False):