change mabel score to probability

Files changed (3) hide show

__pycache__/bert_for_sequence_classification.cpython-37.pyc CHANGED Viewed

Binary files a/__pycache__/bert_for_sequence_classification.cpython-37.pyc and b/__pycache__/bert_for_sequence_classification.cpython-37.pyc differ

__pycache__/bibert_multitask_classification.cpython-37.pyc CHANGED Viewed

Binary files a/__pycache__/bibert_multitask_classification.cpython-37.pyc and b/__pycache__/bibert_multitask_classification.cpython-37.pyc differ

bibert_multitask_classification.py CHANGED Viewed

@@ -41,13 +41,13 @@ class BiBert_MultiTaskPipeline(Pipeline):
         scores = softmax(outputs)
         if top_k == 1 and _legacy:
-            return {"label": self.model.config.id2label[scores.argmax().item()], "score": scores.max().item()}
         dict_scores = [
-            {"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(scores)
         ]
         if not _legacy:
-            dict_scores.sort(key=lambda x: x["score"], reverse=True)
             if top_k is not None:
                 dict_scores = dict_scores[:top_k]
         return dict_scores

         scores = softmax(outputs)
         if top_k == 1 and _legacy:
+            return {"label": self.model.config.id2label[scores.argmax().item()], "probability": scores.max().item()}
         dict_scores = [
+            {"label": self.model.config.id2label[i], "probability": score.item()} for i, score in enumerate(scores)
         ]
         if not _legacy:
+            dict_scores.sort(key=lambda x: x["probability"], reverse=True)
             if top_k is not None:
                 dict_scores = dict_scores[:top_k]
         return dict_scores