Update app.py

app.py

@@ -51,22 +51,52 @@ def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False):
 def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):
     num_layers = len(layer_diffs)
     num_components = len(layer_diffs[0])
+    
+    # Dynamically adjust figure size based on number of layers
+    height = max(8, num_layers / 8)  # Minimum height of 8, scales up for more layers
+    width = max(24, num_components * 3)  # Minimum width of 24, scales with components
+    
+    # Create figure with subplots arranged in 2 rows if there are many components
+    if num_components > 6:
+        nrows = 2
+        ncols = (num_components + 1) // 2
+        fig, axs = plt.subplots(nrows, ncols, figsize=(width, height * 1.5))
+        axs = axs.flatten()
+    else:
+        nrows = 1
+        ncols = num_components
+        fig, axs = plt.subplots(1, num_components, figsize=(width, height))
 
-    fig, axs = plt.subplots(1, num_components, figsize=(24, 8))
     fig.suptitle(f"{base_model_name} <> {chat_model_name}", fontsize=16)
 
+    # Adjust font sizes based on number of layers
+    tick_font_size = max(6, min(10, 300 / num_layers))
+    annot_font_size = max(6, min(10, 200 / num_layers))
+
     for i, component in tqdm(enumerate(layer_diffs[0].keys()), total=len(layer_diffs[0].keys())):
         component_diffs = [[layer_diff[component]] for layer_diff in layer_diffs]
-        sns.heatmap(component_diffs, annot=True, fmt=".9f", cmap="YlGnBu", ax=axs[i], cbar=False)
-        axs[i].set_title(component)
-        axs[i].set_xlabel("Difference")
-        axs[i].set_ylabel("Layer")
+        sns.heatmap(component_diffs, 
+                    annot=True, 
+                    fmt=".9f", 
+                    cmap="YlGnBu", 
+                    ax=axs[i], 
+                    cbar=False,
+                    annot_kws={'size': annot_font_size})
+        
+        axs[i].set_title(component, fontsize=max(10, tick_font_size * 1.2))
+        axs[i].set_xlabel("Difference", fontsize=tick_font_size)
+        axs[i].set_ylabel("Layer", fontsize=tick_font_size)
         axs[i].set_xticks([])
         axs[i].set_yticks(range(num_layers))
-        axs[i].set_yticklabels(range(num_layers))
+        axs[i].set_yticklabels(range(num_layers), fontsize=tick_font_size)
         axs[i].invert_yaxis()
 
-    plt.tight_layout()
+    # Remove any empty subplots if using 2 rows
+    if num_components > 6:
+        for j in range(i + 1, len(axs)):
+            fig.delaxes(axs[j])
+
+    plt.tight_layout(rect=[0, 0.03, 1, 0.95])  # Adjust layout to prevent overlap
     
     # Convert plot to image
     buf = io.BytesIO()
@@ -76,7 +106,6 @@ def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):
     return PIL.Image.open(buf)
 
 def gradio_interface(base_model_name, chat_model_name, hf_token, load_one_at_a_time=False):
-    # Update to use 'token' instead of 'use_auth_token' to handle deprecation warning
     base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.bfloat16, token=hf_token)
     chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16, token=hf_token)