[Fix] convert logits to softmax for kappa_p

src/similarity.py

@@ -1,8 +1,9 @@
 import numpy as np
 
-from src.dataloading import load_run_data
 from lmsim.metrics import Metrics, Kappa_p, EC
 
+from src.dataloading import load_run_data
+from src.utils import softmax, one_hot
 
 def load_data_and_compute_similarities(models: list[str], dataset: str, metric_name: str) -> np.array:
     # Load data
@@ -18,22 +19,12 @@ def load_data_and_compute_similarities(models: list[str], dataset: str, metric_n
     return similarities
 
 
-def compute_similarity(metric: Metrics, probs_a: list[np.array], gt_a: list[int], probs_b: list[np.array], gt_b: list[int]) -> float:
+def compute_similarity(metric: Metrics, outputs_a: list[np.array], outputs_b: list[np.array], gt: list[int],) -> float:
     # Check that the models have the same number of responses
-    assert len(probs_a) == len(probs_b), f"Models must have the same number of responses: {len(probs_a)} != {len(probs_b)}"
-
-    # Only keep responses where the ground truth is the same
-    output_a = []
-    output_b = []
-    gt = []
-    for i in range(len(probs_a)):
-        if gt_a[i] == gt_b[i]:
-            output_a.append(probs_a[i])
-            output_b.append(probs_b[i])
-            gt.append(gt_a[i])
-
-    # Placeholder similarity value
-    similarity = metric.compute_k(output_a, output_b, gt)
+    assert len(outputs_a) == len(outputs_b) == len(gt), f"Models must have the same number of responses: {len(outputs_a)} != {len(outputs_b)} != {len(gt)}" 
+    
+    # Compute similarity values
+    similarity = metric.compute_k(outputs_a, outputs_b, gt)
     
     return similarity
 
@@ -54,11 +45,28 @@ def compute_pairwise_similarities(metric_name: str, probs: list[list[np.array]],
     similarities = np.zeros((len(probs), len(probs)))
     for i in range(len(probs)):
         for j in range(i, len(probs)):
-            similarities[i, j] = compute_similarity(metric, probs[i], gts[i], probs[j], gts[j])
+            outputs_a = probs[i]
+            outputs_b = probs[j]
+            gt_a = gts[i]
+            gt_b = gts[j]
+
+            # Format softmax outputs
+            if metric_name == "Kappa_p (prob.)":
+                outputs_a = [softmax(logits) for logits in outputs_a]
+                outputs_b = [softmax(logits) for logits in outputs_b]
+
+            # Assert that the ground truth index is the same
+            indices_to_remove = []
+            if gt_a != gt_b:
+                for idx, (a, b) in enumerate(zip(gt_a, gt_b)):
+                    if a != b:
+                        indices_to_remove.append(idx)
+                for idx in sorted(indices_to_remove, reverse=True):
+                    del outputs_a[idx]
+                    del outputs_b[idx]
+                    del gt_a[idx]
+                    del gt_b[idx] 
+
+            similarities[i, j] = compute_similarity(metric, outputs_a, outputs_b, gt_a)
             similarities[j, i] = similarities[i, j]
     return similarities
-
-def one_hot(probs: np.array) -> np.array:
-    one_hot = np.zeros_like(probs)
-    one_hot[np.argmax(probs)] = 1
-    return one_hot

src/utils.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+
+def softmax(logits: np.ndarray) -> np.ndarray:
+        exp_logits = np.exp(logits - np.max(logits))
+        return exp_logits / exp_logits.sum(axis=0)
+
+def one_hot(probs: np.array) -> np.array:
+    one_hot = np.zeros_like(probs)
+    one_hot[np.argmax(probs)] = 1
+    return one_hot