Create tunBertClassificationPipeline.py

tunBertClassificationPipeline.py

@@ -0,0 +1,25 @@
+from transformers import Pipeline
+import torch
+
+class TBCP(Pipeline):
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        if "text_pair" in kwargs:
+            preprocess_kwargs["text_pair"] = kwargs["text_pair"]
+        return preprocess_kwargs, {}, {}
+
+    def preprocess(self, text, text_pair=None):
+        return self.tokenizer(text, text_pair=text_pair, return_tensors="pt")
+
+    def _forward(self, model_inputs):
+        return self.model(**model_inputs)
+
+    def postprocess(self, model_outputs):
+        logits = model_outputs.logits
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        
+        best_class = probabilities.argmax().item()
+        label = self.model.config.id2label[best_class]
+        score = probabilities.squeeze()[best_class].item()
+        logits = logits.squeeze().tolist()
+        return {"label": label, "score": score, "logits": logits}