Upload NLIScorer

pipeline.py

@@ -363,12 +363,13 @@ class NLIScorer(Pipeline):
     def _sanitize_parameters(self, **kwargs):
         preprocess_kwargs = {}
         postprocess_kwargs = {}
-        # if "threshold" in kwargs:
-        #     postprocess_kwargs["threshold"] = kwargs["threshold"]
+        if "task_name" in kwargs:
+            postprocess_kwargs["task_name"] = kwargs["task_name"]
+        if "threshold" in kwargs:
+            postprocess_kwargs["threshold"] = kwargs["threshold"]
         return preprocess_kwargs, {}, postprocess_kwargs
 
-    def preprocess(self, inputs):
-        task_name = inputs.pop("task_type")
+    def preprocess(self, inputs, task_name):
         TaskClass = TASK_CLASSES[task_name]
         task_class = TaskClass(tokenizer=self.tokenizer, **inputs)
         return task_class.as_model_inputs
@@ -377,17 +378,11 @@ class NLIScorer(Pipeline):
         outputs = self.model(**model_inputs)
         return outputs
 
-    def postprocess(self, model_outputs):
+    def postprocess(self, model_outputs, threshold=0.5):
         pos_scores = model_outputs["logits"].softmax(-1)[0][1]
-        return {"score": pos_scores}
-
-    def __call__(self, inputs, **kwargs):
-        task_name = inputs.get("task_type")
-        task_threshold = TASK_THRESHOLDS[task_name]
-        outputs = super().__call__(inputs, **kwargs)
-        best_class = int(outputs["score"] > task_threshold)
+        best_class = int(pos_scores > threshold)
         if best_class == 1:
-            score = outputs["score"]
+            score = pos_scores
         else:
-            score = 1 - outputs["score"]
+            score = 1 - pos_scores
         return {"score": score.item(), "label": best_class}