Update tunBertClassificationPipeline.py

tunBertClassificationPipeline.py

@@ -3,18 +3,18 @@ import torch
 
 class TBCP(Pipeline):
     def _sanitize_parameters(self, **kwargs):
-        preprocess_kwargs = {}
+        postprocess_kwargs = {}
         if "text_pair" in kwargs:
-            preprocess_kwargs["text_pair"] = kwargs["text_pair"]
-        return preprocess_kwargs, {}, {}
+            postprocess_kwargs["top_k"] = kwargs["top_k"]
+        return {}, {}, postprocess_kwargs
 
-    def preprocess(self, text, text_pair=None):
-        return self.tokenizer(text, text_pair=text_pair, return_tensors="pt")
+    def preprocess(self, text):
+        return self.tokenizer(text, return_tensors="pt")
 
     def _forward(self, model_inputs):
         return self.model(**model_inputs)
 
-    def postprocess(self, model_outputs):
+    def postprocess(self, model_outputs,top_k = None):
         logits = model_outputs.logits
         probabilities = torch.nn.functional.softmax(logits, dim=-1)