[Add] cache loading data from hf

src/dataloading.py

@@ -3,6 +3,7 @@ import numpy as np
 from huggingface_hub import HfApi
 
 from functools import lru_cache
+from datasets.exceptions import DatasetNotFoundError
 
 
 def get_leaderboard_models():
@@ -17,9 +18,7 @@ def get_leaderboard_models():
             dataset_id = dataset.id
             try:
                 # Check if the dataset can be loaded
-                print(dataset_id)
                 check_gated = datasets.get_dataset_config_names(dataset_id)
-                print(check_gated)
                 # Format: "open-llm-leaderboard/<provider>__<model_name>-details"
                 model_part = dataset_id.split("/")[-1].replace("-details", "")
                 if "__" in model_part:
@@ -27,7 +26,7 @@ def get_leaderboard_models():
                     models.append(f"{provider}/{model}")
                 else:
                     models.append(model_part)
-            except Exception as e:
+            except DatasetNotFoundError as e:
                 pass
     
     return sorted(models)
@@ -77,6 +76,7 @@ def filter_labels(doc):
     return labels
 
 
+
 def load_run_data(model_name, dataset_name):
     try:
         model_name = model_name.replace("/", "__")
@@ -104,6 +104,7 @@ def load_run_data(model_name, dataset_name):
     return log_probs, labels
 
 
-
-
+@lru_cache(maxsize=8)
+def load_run_data_cached(model_name, dataset_name):
+    return load_run_data(model_name, dataset_name)
 

src/similarity.py

@@ -2,7 +2,7 @@ import numpy as np
 
 from lmsim.metrics import Metrics, CAPA, EC
 
-from src.dataloading import load_run_data
+from src.dataloading import load_run_data_cached
 from src.utils import softmax, one_hot
 
 def load_data_and_compute_similarities(models: list[str], dataset: str, metric_name: str) -> np.array:
@@ -10,7 +10,7 @@ def load_data_and_compute_similarities(models: list[str], dataset: str, metric_n
     probs = []
     gts = []
     for model in models:
-        model_probs, model_gt = load_run_data(model, dataset)
+        model_probs, model_gt = load_run_data_cached(model, dataset)
         probs.append(model_probs)
         gts.append(model_gt)