Update app.py

app.py

@@ -46,7 +46,7 @@ def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False):
 
     return layer_diffs
 
-def visualize_layer_diffs(layer_diffs):
+def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):  # Added model names as parameters
     num_layers = len(layer_diffs)
     num_components = len(layer_diffs[0])
 
@@ -72,7 +72,7 @@ def gradio_interface(base_model_name, chat_model_name, hf_token, load_one_at_a_t
     chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
 
     layer_diffs = calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=load_one_at_a_time)
-    fig = visualize_layer_diffs(layer_diffs)
+    fig = visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name)  # Pass model names to visualization
     
     return fig
 
@@ -88,4 +88,4 @@ iface = gr.Interface(
     title="Model Weight Difference Visualizer"
 )
 
-iface.launch()
+iface.launch()