[Fix] heatmap not generated and deselection causes error

app.py

@@ -4,6 +4,7 @@ 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}")
     if not selected_models or not selected_dataset:
         return gr.Plot(visible=False)
     
@@ -30,6 +31,7 @@ def create_heatmap(selected_models, selected_dataset):
     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:
@@ -39,43 +41,18 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     gr.Markdown("## Model Similarity Comparison Tool")
     
     with gr.Row():
-        dataset_dropdown = gr.Dropdown(
-            choices=get_leaderboard_datasets(),
-            label="Select Dataset",
-            filterable=True,
-            interactive=True,
-            info="Leaderboard benchmark datasets"
-        )
+        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)
 
-        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 should reset to empty lists
-    clear_btn = gr.Button("Clear Selection")
-    clear_btn.click(
-        lambda: [[], [], gr.Plot.update(visible=False)],
-        outputs=[model_dropdown, dataset_dropdown, heatmap]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()