[Fix] heatmap not generated and deselection causes error

app.py

@@ -4,9 +4,11 @@ import numpy as np
 from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
 
 def create_heatmap(selected_models, selected_dataset):
-    print(f"Creating heatmap with models: {selected_models} and dataset: {selected_dataset}")
+    print(f"Selected Models: {selected_models}")  # Debugging
+    print(f"Selected Dataset: {selected_dataset}")  # Debugging
     if not selected_models or not selected_dataset:
-        return gr.Plot(visible=False)
+        print("No models or dataset selected")  # Debugging
+        return gr.Plot.update(visible=False)
     
     # Generate random similarity matrix
     size = len(selected_models)
@@ -27,11 +29,10 @@ def create_heatmap(selected_models, selected_dataset):
         height=800
     )
     
-    # Return both the figure and visibility update
+    print("Heatmap generated successfully")  # Debugging
     return gr.Plot.update(value=fig, visible=True)
 
 def validate_inputs(selected_models, selected_dataset):
-    print(f"Validating inputs: models={selected_models}, dataset={selected_dataset}")
     if not selected_models:
         raise gr.Error("Please select at least one model!")
     if not selected_dataset:
@@ -41,18 +42,43 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     gr.Markdown("## Model Similarity Comparison Tool")
     
     with gr.Row():
-        model_selector = gr.Dropdown(label="Select Models", choices=get_leaderboard_models_cached(), multiselect=True)
-        dataset_selector = gr.Dropdown(label="Select Dataset", choices=get_leaderboard_datasets())
-        heatmap_output = gr.Plot(visible=False)
-    
-    def on_submit(selected_models, selected_dataset):
-        try:
-            validate_inputs(selected_models, selected_dataset)
-            return create_heatmap(selected_models, selected_dataset)
-        except gr.Error as e:
-            return gr.Markdown(str(e))
-    
-    submit_button = gr.Button("Generate Heatmap")
-    submit_button.click(on_submit, inputs=[model_selector, dataset_selector], outputs=heatmap_output)
+        dataset_dropdown = gr.Dropdown(
+            choices=get_leaderboard_datasets(),
+            label="Select Dataset",
+            filterable=True,
+            interactive=True,
+            info="Leaderboard benchmark datasets"
+        )
 
-demo.launch()
+        model_dropdown = gr.Dropdown(
+            choices=get_leaderboard_models_cached(),
+            label="Select Models",
+            multiselect=True,
+            filterable=True,
+            allow_custom_value=False,
+            info="Search and select multiple models"
+        )
+    
+    generate_btn = gr.Button("Generate Heatmap", variant="primary")
+    heatmap = gr.Plot(label="Similarity Heatmap", visible=False)
+    
+    # Event handling
+    generate_btn.click(
+        fn=validate_inputs,
+        inputs=[model_dropdown, dataset_dropdown],
+        queue=False
+    ).then(
+        fn=create_heatmap,
+        inputs=[model_dropdown, dataset_dropdown],
+        outputs=heatmap
+    )
+    
+    # Clear button
+    clear_btn = gr.Button("Clear Selection")
+    clear_btn.click(
+        lambda: [None, None, gr.Plot.update(visible=False)],
+        outputs=[model_dropdown, dataset_dropdown, heatmap]
+    )
+
+if __name__ == "__main__":
+    demo.launch()