[Fix] plotly heatmap

app.py

@@ -1,33 +1,38 @@
 import gradio as gr
-import plotly.express as px
+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):
+    # Return nothing if no inputs are provided
     if not selected_models or not selected_dataset:
-        return None  # Hide plot if inputs are missing
+        return None
 
-    # Generate random similarity matrix
+    # Generate a random symmetric similarity matrix
     size = len(selected_models)
     similarities = np.random.rand(size, size)
-    similarities = (similarities + similarities.T) / 2  # make symmetric
+    similarities = (similarities + similarities.T) / 2
     similarities = np.round(similarities, 2)
 
-    # Use Plotly Express imshow to create a heatmap
-    fig = px.imshow(similarities,
-                    x=selected_models,
-                    y=selected_models,
-                    color_continuous_scale='Viridis',
-                    zmin=0, zmax=1,
-                    text_auto=True)
-    # Move x-axis labels to top and adjust tick angle for readability
-    fig.update_xaxes(side="top", tickangle=45)
-    # Update overall layout: title, dimensions, margins
+    # Create a heatmap trace using go.Heatmap; we set x and y to the model names.
+    fig = go.Figure(data=go.Heatmap(
+        z=similarities,
+        x=selected_models,
+        y=selected_models,
+        colorscale="Viridis",
+        zmin=0, zmax=1,
+        text=similarities,
+        hoverinfo="text"
+    ))
+
+    # Update layout: add title, axis titles, set fixed dimensions and margins
     fig.update_layout(
         title=f"Similarity Matrix for {selected_dataset}",
-        width=800 + 20 * size,
-        height=800 + 20 * size,
-        margin=dict(b=100, l=100)
+        xaxis_title="Models",
+        yaxis_title="Models",
+        width=800,
+        height=800,
+        margin=dict(l=100, r=100, t=100, b=100)
     )
     return fig
 
@@ -36,9 +41,7 @@ def validate_inputs(selected_models, selected_dataset):
         raise gr.Error("Please select at least one model!")
     if not selected_dataset:
         raise gr.Error("Please select a dataset!")
-    
 
-# Gradio interface setup
 with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     gr.Markdown("## Model Similarity Comparison Tool")
     
@@ -50,7 +53,6 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
             interactive=True,
             info="Leaderboard benchmark datasets"
         )
-
         model_dropdown = gr.Dropdown(
             choices=get_leaderboard_models_cached(),
             label="Select Models",
@@ -61,9 +63,10 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
         )
     
     generate_btn = gr.Button("Generate Heatmap", variant="primary")
+    # Initialize the Plot component without a figure (it will be updated)
     heatmap = gr.Plot(label="Similarity Heatmap", visible=True)
     
-    # Use a single output (the figure)
+    # First validate inputs, then create the heatmap; note that we use a single output.
     generate_btn.click(
         fn=validate_inputs,
         inputs=[model_dropdown, dataset_dropdown],
@@ -74,7 +77,7 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
         outputs=heatmap
     )
     
-    # Clear button: clear selections and the plot
+    # Clear button to reset selections and clear the plot
     clear_btn = gr.Button("Clear Selection")
     clear_btn.click(
         lambda: [None, None, None],

src/test.py

@@ -0,0 +1,14 @@
+import plotly.graph_objects as go
+import numpy as np
+
+models = ["model1", "model2", "model3"]
+size = len(models)
+sim = np.random.rand(size, size)
+sim = (sim + sim.T) / 2
+sim = np.round(sim, 2)
+fig = go.Figure(data=go.Heatmap(z=sim, x=models, y=models, colorscale="Viridis"))
+fig.update_layout(title="Test Heatmap", xaxis_title="Models", yaxis_title="Models", width=800, height=800)
+fig.show()
+
+# Save fig
+fig.write_html("heatmap.html")