[Fix] similarity heatmap creation

app.py

@@ -1,21 +1,16 @@
 import gradio as gr
 import plotly.graph_objects as go
 import numpy as np
-
-
 from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
 
-
 def create_heatmap(selected_models, selected_dataset):
-    if not selected_models:
+    if not selected_models or not selected_dataset:
         return gr.Plot(visible=False)
     
-    # Generate random similarity matrix (replace with actual computation)
+    # Generate random similarity matrix
     size = len(selected_models)
     similarities = np.random.rand(size, size)
-    
-    # Create symmetric matrix 
-    similarities = (similarities + similarities.T) / 2
+    similarities = (similarities + similarities.T) / 2  # Make symmetric
     
     # Create plot
     fig = go.Figure(data=go.Heatmap(
@@ -31,9 +26,8 @@ def create_heatmap(selected_models, selected_dataset):
         height=800
     )
     
-    with gr.Loading():
-        return gr.Plot(value=fig, visible=True)
-
+    # Return both the figure and visibility update
+    return gr.Plot.update(value=fig, visible=True)
 
 def validate_inputs(selected_models, selected_dataset):
     if not selected_models:
@@ -41,11 +35,9 @@ def validate_inputs(selected_models, selected_dataset):
     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")
     
-    # Model selection section
     with gr.Row():
         dataset_dropdown = gr.Dropdown(
             choices=get_leaderboard_datasets(),
@@ -61,20 +53,13 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
             multiselect=True,
             filterable=True,
             allow_custom_value=False,
-            info="Search and select multiple models (click selected models to remove)"
+            info="Search and select multiple models"
         )
     
-    # Add generate button
     generate_btn = gr.Button("Generate Heatmap", variant="primary")
+    heatmap = gr.Plot(label="Similarity Heatmap", visible=False)
     
-    # Heatmap display
-    heatmap = gr.Plot(
-        label="Similarity Heatmap",
-        visible=False,
-        container=False
-    )
-
-    # Button click handler
+    # Event handling
     generate_btn.click(
         fn=validate_inputs,
         inputs=[model_dropdown, dataset_dropdown],
@@ -84,14 +69,13 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
         inputs=[model_dropdown, dataset_dropdown],
         outputs=heatmap
     )
-
+    
+    # Clear button should reset to empty lists
     clear_btn = gr.Button("Clear Selection")
     clear_btn.click(
-        lambda: [None, None, gr.Plot(visible=False)],
+        lambda: [[], [], gr.Plot.update(visible=False)],
         outputs=[model_dropdown, dataset_dropdown, heatmap]
     )
 
-
-
 if __name__ == "__main__":
     demo.launch()