Update app.py

app.py

@@ -37,11 +37,10 @@ MODEL_NAMES = get_model_names(DATA)
 # Define the six metrics in the desired order.
 METRICS = ["Correctness", "Completeness", "Conciseness", "Helpfulness", "Honesty", "Harmlessness"]
 
-def generate_heatmap_image(model_entry, max_size=(20, 20)):
+def generate_heatmap_image(model_entry):
     """
     For a given model entry, extract the six metrics and compute a 6x6 similarity matrix
     using the definition: similarity = 1 - |v_i - v_j|, then return the heatmap as a PIL image.
-    The image is then resized (via thumbnail) to fit the specified max_size.
     """
     scores = model_entry["claude-3.5-sonnet Scores"]["3C3H Scores"]
     # Create a vector with the metrics in the defined order.
@@ -51,7 +50,8 @@ def generate_heatmap_image(model_entry, max_size=(20, 20)):
     # Create a mask for the upper triangle (keeping the diagonal visible).
     mask = np.triu(np.ones_like(matrix, dtype=bool), k=1)
     
-    plt.figure(figsize=(6, 5))
+    # Make the figure itself smaller
+    plt.figure(figsize=(4, 4))
     sns.heatmap(matrix,
                 mask=mask,
                 annot=True,
@@ -71,9 +71,8 @@ def generate_heatmap_image(model_entry, max_size=(20, 20)):
     plt.close()
     buf.seek(0)
     
-    # Convert the buffer into a PIL Image and resize it
+    # Convert the buffer into a PIL Image
     image = Image.open(buf).convert("RGB")
-    image.thumbnail(max_size)  # e.g., (400, 400)
     return image
 
 def generate_heatmaps(selected_model_names):
@@ -105,7 +104,13 @@ with gr.Blocks() as demo:
     
     generate_btn = gr.Button("Generate Heatmaps")
     # The 'columns' parameter will display images in a grid with 2 columns.
-    gallery = gr.Gallery(label="Heatmaps", columns=2)
+    # 'image_size=(300, 300)' ensures each image is displayed at 300x300 px.
+    gallery = gr.Gallery(
+        label="Heatmaps",
+        columns=2,
+        image_size=(200, 200),
+        object_fit="contain"
+    )
     
     generate_btn.click(fn=generate_heatmaps, inputs=model_dropdown, outputs=gallery)