Update app.py

app.py

@@ -4,6 +4,8 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 from tqdm import tqdm
 import gradio as gr
+import io
+import PIL.Image
 
 def calculate_weight_diff(base_weight, chat_weight):
     return torch.abs(base_weight - chat_weight).mean().item()
@@ -46,7 +48,7 @@ def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False):
 
     return layer_diffs
 
-def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):  # Added model names as parameters
+def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):
     num_layers = len(layer_diffs)
     num_components = len(layer_diffs[0])
 
@@ -65,26 +67,31 @@ def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):  # Add
         axs[i].invert_yaxis()
 
     plt.tight_layout()
-    return fig
+    
+    # Convert plot to image
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
+    buf.seek(0)
+    plt.close(fig)  # Close the figure to free memory
+    return PIL.Image.open(buf)
 
 def gradio_interface(base_model_name, chat_model_name, hf_token, load_one_at_a_time=False):
-    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
-    chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
+    # 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)
 
     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, base_model_name, chat_model_name)  # Pass model names to visualization
-    
-    return fig
+    return visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name)
 
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
         gr.Textbox(lines=2, placeholder="Enter base model name"),
-        gr.Textbox(lines=2, placeholder="Enter finetuned model name"),
-        gr.Textbox(lines=2, placeholder="Enter Hugging Face token"),
+        gr.Textbox(lines=2, placeholder="Enter chat model name"),
+        gr.Textbox(lines=2, placeholder="Enter Hugging Face token", type="password"),  # Hide token input
         gr.Checkbox(label="Load one layer at a time")
     ],
-    outputs="image",
+    outputs=gr.Image(type="pil"),  # Specify PIL image output
     title="Model Weight Difference Visualizer"
 )