Update app.py

app.py

@@ -8,24 +8,41 @@ import gradio as gr
 def calculate_weight_diff(base_weight, chat_weight):
     return torch.abs(base_weight - chat_weight).mean().item()
 
-def calculate_layer_diffs(base_model, chat_model):
+def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False):
     layer_diffs = []
-    for base_layer, chat_layer in tqdm(zip(base_model.model.layers, chat_model.model.layers), total=len(base_model.model.layers)):
-        layer_diff = {
-            'input_layernorm': calculate_weight_diff(base_layer.input_layernorm.weight, chat_layer.input_layernorm.weight),
-            'mlp_down_proj': calculate_weight_diff(base_layer.mlp.down_proj.weight, chat_layer.mlp.down_proj.weight),
-            'mlp_gate_proj': calculate_weight_diff(base_layer.mlp.gate_proj.weight, chat_layer.mlp.gate_proj.weight),
-            'mlp_up_proj': calculate_weight_diff(base_layer.mlp.up_proj.weight, chat_layer.mlp.up_proj.weight),
-            'post_attention_layernorm': calculate_weight_diff(base_layer.post_attention_layernorm.weight, chat_layer.post_attention_layernorm.weight),
-            'self_attn_q_proj': calculate_weight_diff(base_layer.self_attn.q_proj.weight, chat_layer.self_attn.q_proj.weight),
-            'self_attn_k_proj': calculate_weight_diff(base_layer.self_attn.k_proj.weight, chat_layer.self_attn.k_proj.weight),
-            'self_attn_v_proj': calculate_weight_diff(base_layer.self_attn.v_proj.weight, chat_layer.self_attn.v_proj.weight),
-            'self_attn_o_proj': calculate_weight_diff(base_layer.self_attn.o_proj.weight, chat_layer.self_attn.o_proj.weight)
-        }
-        layer_diffs.append(layer_diff)
+    layers = zip(base_model.model.layers, chat_model.model.layers)
+    
+    if load_one_at_a_time:
+        for base_layer, chat_layer in tqdm(layers, total=len(base_model.model.layers)):
+            layer_diff = {
+                'input_layernorm': calculate_weight_diff(base_layer.input_layernorm.weight, chat_layer.input_layernorm.weight),
+                'mlp_down_proj': calculate_weight_diff(base_layer.mlp.down_proj.weight, chat_layer.mlp.down_proj.weight),
+                'mlp_gate_proj': calculate_weight_diff(base_layer.mlp.gate_proj.weight, chat_layer.mlp.gate_proj.weight),
+                'mlp_up_proj': calculate_weight_diff(base_layer.mlp.up_proj.weight, chat_layer.mlp.up_proj.weight),
+                'post_attention_layernorm': calculate_weight_diff(base_layer.post_attention_layernorm.weight, chat_layer.post_attention_layernorm.weight),
+                'self_attn_q_proj': calculate_weight_diff(base_layer.self_attn.q_proj.weight, chat_layer.self_attn.q_proj.weight),
+                'self_attn_k_proj': calculate_weight_diff(base_layer.self_attn.k_proj.weight, chat_layer.self_attn.k_proj.weight),
+                'self_attn_v_proj': calculate_weight_diff(base_layer.self_attn.v_proj.weight, chat_layer.self_attn.v_proj.weight),
+                'self_attn_o_proj': calculate_weight_diff(base_layer.self_attn.o_proj.weight, chat_layer.self_attn.o_proj.weight)
+            }
+            layer_diffs.append(layer_diff)
 
-        base_layer, chat_layer = None, None
-        del base_layer, chat_layer
+            base_layer, chat_layer = None, None
+            del base_layer, chat_layer
+    else:
+        for base_layer, chat_layer in tqdm(layers, total=len(base_model.model.layers)):
+            layer_diff = {
+                'input_layernorm': calculate_weight_diff(base_layer.input_layernorm.weight, chat_layer.input_layernorm.weight),
+                'mlp_down_proj': calculate_weight_diff(base_layer.mlp.down_proj.weight, chat_layer.mlp.down_proj.weight),
+                'mlp_gate_proj': calculate_weight_diff(base_layer.mlp.gate_proj.weight, chat_layer.mlp.gate_proj.weight),
+                'mlp_up_proj': calculate_weight_diff(base_layer.mlp.up_proj.weight, chat_layer.mlp.up_proj.weight),
+                'post_attention_layernorm': calculate_weight_diff(base_layer.post_attention_layernorm.weight, chat_layer.post_attention_layernorm.weight),
+                'self_attn_q_proj': calculate_weight_diff(base_layer.self_attn.q_proj.weight, chat_layer.self_attn.q_proj.weight),
+                'self_attn_k_proj': calculate_weight_diff(base_layer.self_attn.k_proj.weight, chat_layer.self_attn.k_proj.weight),
+                'self_attn_v_proj': calculate_weight_diff(base_layer.self_attn.v_proj.weight, chat_layer.self_attn.v_proj.weight),
+                'self_attn_o_proj': calculate_weight_diff(base_layer.self_attn.o_proj.weight, chat_layer.self_attn.o_proj.weight)
+            }
+            layer_diffs.append(layer_diff)
 
     return layer_diffs
 
@@ -50,11 +67,11 @@ def visualize_layer_diffs(layer_diffs):
     plt.tight_layout()
     return fig
 
-def gradio_interface(base_model_name, chat_model_name):
+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)
 
-    layer_diffs = calculate_layer_diffs(base_model, chat_model)
+    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)
     
     return fig
@@ -63,7 +80,8 @@ iface = gr.Interface(
     fn=gradio_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 chat model name"),
+        gr.Checkbox(label="Load one layer at a time")
     ],
     outputs="image",
     title="Model Weight Difference Visualizer"