update for pipeline

config.json

@@ -8,5 +8,6 @@
   },
   "model_type": "HHEMv2Config",
   "torch_dtype": "float32",
-  "transformers_version": "4.39.3"
+  "transformers_version": "4.39.3",
+  "id2label": {"0": "hallucinated", "1": "consistent"}
 }

modeling_hhem_v2.py

@@ -45,8 +45,15 @@ class HHEMv2ForSequenceClassification(PreTrainedModel):
     #     combined_model = PeftModel.from_pretrained(base_model, checkpoint, is_trainable=False)
     #     self.t5 = combined_model
 
-    def forward(self, **kwargs):
-        return self.t5(**kwargs)
+    def forward(self, **kwargs): # To cope with `text-classiication` pipeline
+        self.t5.eval()
+        with torch.no_grad():
+            outputs = self.t5(**kwargs)
+        logits = outputs.logits    
+        logits = logits[:, 0, :]
+        outputs.logits = logits
+        return outputs
+        # return self.t5(**kwargs)
 
     def predict(self, text_pairs):
         tokenizer = self.tokenzier