Merge pull request #204 from Machine-Learning-for-Medical-Language/ho…

…xfix-logit-to-prob

[FIX] Updated `structure_labels` function to output softmax probability instead of logit for the best category

src/cnlpt/cnlp_predict.py

@@ -15,6 +15,11 @@
 logger = logging.getLogger(__name__)
 
 
+def simple_softmax(x: list):
+    """Softmax values for 1-D score array"""
+    return np.exp(x) / np.sum(np.exp(x), axis=0)
+
+
 def restructure_prediction(
     task_names: List[str],
     raw_prediction: EvalPrediction,
@@ -76,7 +81,9 @@ def structure_labels(
     else:
         preds = np.argmax(p.predictions[task_ind], axis=1)
         if output_prob:
-            prob_values = np.max(p.predictions[task_ind], axis=1)
+            prob_values = np.max(
+                [simple_softmax(logits) for logits in p.predictions[task_ind]], axis=1
+            )
 
     # for inference
     if not hasattr(p, "label_ids") or p.label_ids is None: