Update app.py

app.py

@@ -37,10 +37,11 @@ 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):
+def generate_heatmap_image(model_entry, max_size=(400, 400)):
     """
     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 image as bytes.
+    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.
@@ -69,14 +70,16 @@ def generate_heatmap_image(model_entry):
     plt.savefig(buf, format="png")
     plt.close()
     buf.seek(0)
-    # Convert the buffer into a PIL Image.
+    
+    # Convert the buffer into a PIL Image and resize it
     image = Image.open(buf).convert("RGB")
+    image.thumbnail(max_size)  # e.g., (400, 400)
     return image
 
 def generate_heatmaps(selected_model_names):
     """
     Filter the global DATA for entries matching the selected model names,
-    generate a heatmap for each, and return a list of image bytes.
+    generate a heatmap for each, and return a list of PIL images.
     """
     filtered_entries = [entry for entry in DATA if entry["Meta"]["Model Name"] in selected_model_names]
     images = []
@@ -93,10 +96,15 @@ with gr.Blocks() as demo:
     gr.Markdown("Select the models you want to compare and generate their heatmaps below.")
     
     with gr.Row():
-        model_dropdown = gr.Dropdown(choices=MODEL_NAMES, label="Select Model(s)", multiselect=True, value=MODEL_NAMES[:3])
+        model_dropdown = gr.Dropdown(
+            choices=MODEL_NAMES,
+            label="Select Model(s)",
+            multiselect=True,
+            value=MODEL_NAMES[:3]
+        )
     
     generate_btn = gr.Button("Generate Heatmaps")
-    # Use the 'columns' parameter to set a grid layout in the gallery.
+    # The 'columns' parameter will display images in a grid with 2 columns.
     gallery = gr.Gallery(label="Heatmaps", columns=2)
     
     generate_btn.click(fn=generate_heatmaps, inputs=model_dropdown, outputs=gallery)