[Add] clear button, load the right data, create plot on click

app.py

@@ -6,35 +6,41 @@ import numpy as np
 from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
 
 
-def create_heatmap(selected_models, benchmark):
+def create_heatmap(selected_models, selected_dataset):
     if not selected_models:
-        return gr.update(visible=False)
+        return gr.Plot(visible=False)
     
-    # Generate random similarity matrix
+    # Generate random similarity matrix (replace with actual computation)
     size = len(selected_models)
     similarities = np.random.rand(size, size)
     
-    # Create symmetric matrix (for demo purposes)
+    # Create symmetric matrix 
     similarities = (similarities + similarities.T) / 2
     
-    # Create heatmap with Plotly
+    # Create plot
     fig = go.Figure(data=go.Heatmap(
         z=similarities,
         x=selected_models,
         y=selected_models,
-        colorscale='Viridis',
-        hoverongaps=False
+        colorscale='Viridis'
     ))
     
     fig.update_layout(
-        title=f"Model Similarity for {benchmark}",
-        xaxis_title="Models",
-        yaxis_title="Models",
-        height=600,
-        width=800
+        title=f"Similarity Matrix for {selected_dataset}",
+        width=800,
+        height=800
     )
     
-    return fig
+    with gr.Loading():
+        return gr.Plot(value=fig, visible=True)
+
+
+def validate_inputs(selected_models, selected_dataset):
+    if not selected_models:
+        raise gr.Error("Please select at least one model!")
+    if not selected_dataset:
+        raise gr.Error("Please select a dataset!")
+
 
 with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     gr.Markdown("## Model Similarity Comparison Tool")
@@ -42,12 +48,12 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     # Model selection section
     with gr.Row():
         dataset_dropdown = gr.Dropdown(
-                choices=get_leaderboard_datasets(),
-                label="Select Dataset",
-                filterable=True,
-                interactive=True,
-                info="Leaderboard benchmark datasets"
-            )
+            choices=get_leaderboard_datasets(),
+            label="Select Dataset",
+            filterable=True,
+            interactive=True,
+            info="Leaderboard benchmark datasets"
+        )
 
         model_dropdown = gr.Dropdown(
             choices=get_leaderboard_models_cached(),
@@ -58,20 +64,34 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
             info="Search and select multiple models (click selected models to remove)"
         )
     
+    # Add generate button
+    generate_btn = gr.Button("Generate Heatmap", variant="primary")
+    
     # Heatmap display
     heatmap = gr.Plot(
         label="Similarity Heatmap",
         visible=False,
         container=False
     )
-    
-    # Interactive updates
-    model_dropdown.input(
+
+    # Button click handler
+    generate_btn.click(
+        fn=validate_inputs,
+        inputs=[model_dropdown, dataset_dropdown],
+        queue=False
+    ).then(
         fn=create_heatmap,
-        inputs=(model_dropdown, dataset_dropdown),
+        inputs=[model_dropdown, dataset_dropdown],
         outputs=heatmap
     )
 
+    clear_btn = gr.Button("Clear Selection")
+    clear_btn.click(
+        lambda: [None, None, gr.Plot(visible=False)],
+        outputs=[model_dropdown, dataset_dropdown, heatmap]
+    )
+
+
 
 if __name__ == "__main__":
     demo.launch()

src/dataloading.py

@@ -5,22 +5,18 @@ from functools import lru_cache
 def get_leaderboard_models():
     api = HfApi()
     
-    # List all files in the repository
-    files = api.list_repo_files(
-        repo_id="open-llm-leaderboard/open_llm_leaderboard",
-        repo_type="space"
-    )
+    # List all datasets in the open-llm-leaderboard organization
+    datasets = api.list_datasets(author="open-llm-leaderboard")
     
     models = []
-    for file in files:
-        # Filter files in the "open_llm_leaderboard" directory
-        if file.startswith("open_llm_leaderboard/") and "-details" in file and "__" in file:
-            # Extract provider and model name from filename
-            filename = file.split("/")[-1].replace("-details", "")
-            provider, model = filename.split("__", 1)
+    for dataset in datasets:
+        if dataset.id.endswith("-details"):
+            # Format: "open-llm-leaderboard/<provider>__<model_name>-details"
+            model_part = dataset.id.split("/")[-1].replace("-details", "")
+            provider, model = model_part.split("__", 1)
             models.append(f"{provider}/{model}")
     
-    return sorted(list(set(models)))  # Remove duplicates
+    return sorted(models)