Upload DebertaArgClassifier

config.json

@@ -30,7 +30,6 @@
   },
   "label2id": null,
   "model_type": "deberta_arg_classifier",
-  "number_labels": 20,
   "torch_dtype": "float32",
-  "transformers_version": "4.26.1"
+  "transformers_version": "4.27.1"
 }

configuration_deberta_arg_classifier.py

@@ -3,6 +3,7 @@ from transformers import PretrainedConfig
 class DebertaConfig(PretrainedConfig):
     model_type = "deberta_arg_classifier"
 
-    def __init__(self, num_labels: int=20, **kwargs):
-        self.number_labels = num_labels
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
+
+#%%

modeling_deberta_arg_classifier.py

@@ -11,7 +11,7 @@ class DebertaArgClassifier(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.bert = AutoModel.from_pretrained("microsoft/deberta-large")
-        self.classifier = nn.Linear(self.bert.config.hidden_size, config.number_labels)
+        self.classifier = nn.Linear(self.bert.config.hidden_size, config.num_labels)
         self.criterion = nn.BCEWithLogitsLoss()
 
 
@@ -20,10 +20,11 @@ class DebertaArgClassifier(PreTrainedModel):
         output = self._cls_embeddings(output)
         output_cls = self.classifier(output)
         output = torch.sigmoid(output_cls)
+        loss = None
         if labels is not None:
             loss = self.cirterion(output_cls, labels)
-            return {"loss": loss, "logits": output}
-        return {"logits": output}
+            return {"loss": loss, "output": output}
+        return {"loss": loss, "output": output}
 
 
     def _cls_embeddings(self, output):