[Add, Ref] integrate similarity computation, fix one-hot for EC, add login option

app.py

@@ -1,17 +1,21 @@
+import os
 import gradio as gr
 import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
 from io import BytesIO
 from PIL import Image
-from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
 from huggingface_hub import login
 
+from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
+from src.similarity import load_data_and_compute_similarities
+
 # Set matplotlib backend for non-GUI environments
 plt.switch_backend('Agg')
 
 # Login to Hugging Face Hub
-login(token="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")
+token = os.getenv("HF_TOKEN")
+login(token=token)
 
 
 def create_heatmap(selected_models, selected_dataset, selected_metric):
@@ -20,13 +24,8 @@ def create_heatmap(selected_models, selected_dataset, selected_metric):
     
     # Sort models and get short names
     selected_models = sorted(selected_models)
-    selected_models_short = [model.split("/")[-1] for model in selected_models]
     
-    # Generate random similarity matrix
-    size = len(selected_models)
-    similarities = np.random.rand(size, size)
-    similarities = (similarities + similarities.T) / 2
-    similarities = np.round(similarities, 2)
+    similarities = load_data_and_compute_similarities(selected_models, selected_dataset, selected_metric)
 
     # Create figure and heatmap using seaborn
     plt.figure(figsize=(8, 6))
@@ -37,8 +36,8 @@ def create_heatmap(selected_models, selected_dataset, selected_metric):
         cmap="viridis",
         vmin=0,
         vmax=1,
-        xticklabels=selected_models_short,
-        yticklabels=selected_models_short
+        xticklabels=selected_models,
+        yticklabels=selected_models
     )
     
     # Customize plot

src/dataloading.py

@@ -57,7 +57,7 @@ def filter_labels(doc):
     labels = []
     if "answer_index" in doc[0].keys():
         for d in doc:
-            labels.append(int(d["answer_index"]))
+            labels.append(d["answer_index"])
     else:
         for d in doc:
             if d["answer"] == "False":
@@ -66,6 +66,7 @@ def filter_labels(doc):
                 labels.append(1)
             else:
                 raise ValueError("Invalid label")
+    return labels
 
 def load_run_data(model_name, dataset_name):
     try:
@@ -94,10 +95,6 @@ def load_run_data(model_name, dataset_name):
     return log_probs, labels
 
 
-if __name__ == "__main__":
-    model_ids = [
-        'Qwen/Qwen2.5-7B-Instruct'
-    ]
-    
-    datasets = get_leaderboard_datasets(model_ids)
-    print(datasets)   
+
+
+

src/similarity.py

@@ -1,11 +1,10 @@
 import numpy as np
 
-from src.dataloading import load_run_data
-from lmsim.metrics import Metric, Kappa_p, EC
+from dataloading import load_run_data, get_leaderboard_datasets
+from lmsim.metrics import Metrics, Kappa_p, EC
 
 
-
-def load_data_and_compute_similarities(models, dataset, metric_name):
+def load_data_and_compute_similarities(models: list[str], dataset: str, metric_name: str) -> np.array:
     # Load data
     probs = []
     gts = []
@@ -15,11 +14,11 @@ def load_data_and_compute_similarities(models, dataset, metric_name):
         gts.append(model_gt)
     
     # Compute pairwise similarities
-    similarities = compute_pairwise_similarities(probs, gts, metric_name)
+    similarities = compute_pairwise_similarities(metric_name, probs, gts)
     return similarities
 
 
-def compute_similarity(metric: Metric, probs_a: list[np.array], gt_a: list[int], probs_b: list[np.array], gt_b: list[int]) -> float:
+def compute_similarity(metric: Metrics, probs_a: list[np.array], gt_a: list[int], probs_b: list[np.array], gt_b: 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)}"
 
@@ -44,16 +43,22 @@ def compute_pairwise_similarities(metric_name: str, probs: list[list[np.array]],
     if metric_name == "Kappa_p (prob.)":
         metric = Kappa_p()
     elif metric_name == "Kappa_p (det.)":
-        metric = Kappa_p()
+        metric = Kappa_p(prob=False)
+        # Convert probabilities to one-hot
+        probs = [[one_hot(p) for p in model_probs] for model_probs in probs]
     elif metric_name == "Error Consistency":
         metric = EC()
     else:
-        raise ValueError(f"Invalid metric: {metric_name}")
-    
+        raise ValueError(f"Invalid metric: {metric_name}") 
 
     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])
             similarities[j, i] = similarities[i, j]
-    return similarities
+    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