Update app.py

app.py

@@ -67,9 +67,9 @@ def visualize_layer_diffs(layer_diffs):
     plt.tight_layout()
     return fig
 
-def gradio_interface(base_model_name, chat_model_name, load_one_at_a_time=False):
-    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.bfloat16)
-    chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16)
+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)
 
     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)
@@ -81,6 +81,7 @@ iface = gr.Interface(
     inputs=[
         gr.Textbox(lines=2, placeholder="Enter base model name"),
         gr.Textbox(lines=2, placeholder="Enter chat model name"),
+        gr.Textbox(lines=2, placeholder="Enter Hugging Face token"),
         gr.Checkbox(label="Load one layer at a time")
     ],
     outputs="image",