add process method

modeling_tunbert.py

@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 from transformers import  PreTrainedModel, BertModel
 from transformers.modeling_outputs import SequenceClassifierOutput
@@ -38,6 +39,12 @@ class TunBERT(PreTrainedModel):
             loss = loss_func(logits.view(-1,self.config.type_vocab_size),labels.view(-1)) 
         return SequenceClassifierOutput(loss = loss, logits= logits, hidden_states=outputs.last_hidden_state,attentions=outputs.attentions)
 
+    def process(self,**inputs):
+        with torch.no_grad():
+            out = self.forward(**inputs)
+        out = torch.argmax(output.logits,dim=1)
+        return ["positive" if index == 0 else "negative" for index in out.tolist()]
+
 
 TunBertConfig.register_for_auto_class()
 TunBERT.register_for_auto_class("AutoModelForSequenceClassification")