Upload model

modeling_ESGBertReddit.py

@@ -11,7 +11,7 @@ class ClassificationModel(PreTrainedModel):
         self.W = nn.Linear(self.bert.config.hidden_size, config.num_classes)
         self.num_classes = config.num_classes
         
-    def forward(self, input_ids, attention_mask, token_type_ids):
+    def forward(self,input_ids,attention_mask,token_type_ids,**kw):
         h, _, attn = self.bert(input_ids=input_ids, 
                                attention_mask=attention_mask,
                                token_type_ids=token_type_ids).values()
@@ -26,9 +26,9 @@ class BertModelForESGClassification(PreTrainedModel):
         super().__init__(config)
         self.model = ClassificationModel(config)
 
-    def forward(self, inputs, labels=None):
+    def forward(self,**inputs):
         logits,_ = self.model(**inputs)
-        if labels is not None:
+        if "labels" in inputs:
             loss = torch.nn.cross_entropy(logits, labels)
             return {"loss": loss, "logits": logits}
         return {"logits": logits}