Update app.py

app.py

@@ -1,129 +1,213 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import io
 import torch
+import pandas as pd
+import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
-from tqdm import tqdm
+import plotly.graph_objects as go
 import gradio as gr
-import io
 import PIL.Image
+from transformers import AutoModelForCausalLM
+from tqdm import tqdm
+
+# Set style for matplotlib
+sns.set_theme(style="whitegrid")
 
 def calculate_weight_diff(base_weight, chat_weight):
+    """Calculates the mean absolute difference between two tensors."""
     return torch.abs(base_weight - chat_weight).mean().item()
 
-def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False):
+def calculate_layer_diffs(base_model, chat_model):
+    """Iterates through layers and calculates differences for specific projections."""
     layer_diffs = []
+    
+    # We zip the layers to iterate through them simultaneously
     layers = zip(base_model.model.layers, chat_model.model.layers)
+    total_layers = len(base_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
-    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)
+    # List of components we want to track
+    components_to_track = [
+        ('input_layernorm', lambda l: l.input_layernorm.weight),
+        ('self_attn_q_proj', lambda l: l.self_attn.q_proj.weight),
+        ('self_attn_k_proj', lambda l: l.self_attn.k_proj.weight),
+        ('self_attn_v_proj', lambda l: l.self_attn.v_proj.weight),
+        ('self_attn_o_proj', lambda l: l.self_attn.o_proj.weight),
+        ('post_attention_layernorm', lambda l: l.post_attention_layernorm.weight),
+        ('mlp_gate_proj', lambda l: l.mlp.gate_proj.weight),
+        ('mlp_up_proj', lambda l: l.mlp.up_proj.weight),
+        ('mlp_down_proj', lambda l: l.mlp.down_proj.weight),
+    ]
+
+    print("Calculating differences...")
+    for base_layer, chat_layer in tqdm(layers, total=total_layers):
+        layer_data = {}
+        for name, getter in components_to_track:
+            try:
+                val = calculate_weight_diff(getter(base_layer), getter(chat_layer))
+                layer_data[name] = val
+            except AttributeError:
+                # Handle cases where architecture might differ slightly (e.g., bias terms)
+                layer_data[name] = 0.0
+        
+        layer_diffs.append(layer_data)
 
     return layer_diffs
 
-def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name):
+def visualize_2d_heatmap(layer_diffs, base_model_name, chat_model_name):
+    """Generates the static 2D Heatmap image."""
+    if not layer_diffs:
+        return None
+        
     num_layers = len(layer_diffs)
-    num_components = len(layer_diffs[0])
+    components = list(layer_diffs[0].keys())
+    num_components = len(components)
     
-    # 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
+    # Dynamically adjust figure size
+    height = max(8, num_layers / 6)
+    width = max(20, num_components * 2.5)
     
-    # Create figure with subplots arranged in 2 rows if there are many components
+    # Logic for subplot arrangement
     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.suptitle(f"{base_model_name} <> {chat_model_name}", fontsize=16)
+    fig, axs = plt.subplots(nrows, ncols, figsize=(width, height * (1.2 if nrows > 1 else 1)))
+    axs = axs.flatten() if num_components > 1 else [axs]
 
-    # 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))
+    fig.suptitle(f"Weight Differences: {base_model_name} vs {chat_model_name}", fontsize=16, y=0.98)
 
-    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, 
+    # Font sizing logic
+    tick_font_size = max(6, min(10, 300 / num_layers))
+    
+    for i, component in enumerate(components):
+        # Extract data for this specific component across all layers
+        data = [[row[component]] for row in layer_diffs]
+        
+        sns.heatmap(data, 
                     annot=True, 
-                    fmt=".9f", 
-                    cmap="YlGnBu", 
+                    fmt=".6f", 
+                    cmap="viridis", 
                     ax=axs[i], 
                     cbar=False,
-                    annot_kws={'size': annot_font_size})
+                    annot_kws={'size': tick_font_size * 0.8})
         
-        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_title(component, fontsize=12, fontweight='bold')
         axs[i].set_yticks(range(num_layers))
         axs[i].set_yticklabels(range(num_layers), fontsize=tick_font_size)
-        axs[i].invert_yaxis()
+        axs[i].set_xticks([]) # Hide x-axis ticks for the single column heatmap
+        axs[i].invert_yaxis() # Layer 0 at bottom or top? Usually 0 is bottom in diagrams, but top in matrices.
+                              # Let's keep 0 at top (standard matrix view) or remove invert for 0 at bottom.
+                              # Standard heatmap has index 0 at top.
 
-    # Remove any empty subplots if using 2 rows
-    if num_components > 6:
-        for j in range(i + 1, len(axs)):
-            fig.delaxes(axs[j])
+    # Remove empty subplots
+    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
+    plt.tight_layout(rect=[0, 0, 1, 0.96])
     
-    # Convert plot to image
     buf = io.BytesIO()
-    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
+    fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
     buf.seek(0)
-    plt.close(fig)  # Close the figure to free memory
+    plt.close(fig)
     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, token=hf_token)
-    chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16, token=hf_token)
+def visualize_3d_surface(layer_diffs):
+    """Generates an interactive 3D Surface plot using Plotly."""
+    if not layer_diffs:
+        return None
 
-    layer_diffs = calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=load_one_at_a_time)
-    return visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name)
+    # Convert list of dicts to DataFrame for easier handling
+    df = pd.DataFrame(layer_diffs)
+    
+    # X axis: Components
+    x_labels = df.columns.tolist()
+    # Y axis: Layers
+    y_labels = df.index.tolist()
+    # Z axis: Values (Transposed because Plotly expects Z[y][x])
+    z_data = df.values
+    
+    fig = go.Figure(data=[go.Surface(z=z_data, x=x_labels, y=y_labels, colorscale='Viridis')])
 
-if __name__ == "__main__":
-    iface = gr.Interface(
-        fn=gradio_interface,
-        inputs=[
-            gr.Textbox(label="Base Model Name", lines=2),
-            gr.Textbox(label="Chat Model Name", lines=2),
-            gr.Textbox(label="Hugging Face Token", type="password", lines=2),
-            gr.Checkbox(label="Load one layer at a time")
-        ],
-        outputs=gr.Image(type="pil", label="Weight Differences Visualization"),
-        title="Model Weight Difference Visualizer",
-        cache_examples=False
+    fig.update_layout(
+        title='3D Landscape of Weight Differences',
+        scene=dict(
+            xaxis_title='Model Components',
+            yaxis_title='Layer Index',
+            zaxis_title='Mean Weight Diff',
+            xaxis=dict(tickangle=45),
+        ),
+        autosize=True,
+        height=800,
+        margin=dict(l=65, r=50, b=65, t=90)
     )
     
-    iface.launch(share=False, server_port=7860)
+    return fig
+
+def process_models(base_name, chat_name, hf_token):
+    try:
+        print(f"Loading {base_name}...")
+        base_model = AutoModelForCausalLM.from_pretrained(
+            base_name, 
+            torch_dtype=torch.bfloat16, 
+            token=hf_token,
+            device_map="cpu", # Force CPU to avoid GPU OOM during comparison if models are large
+            trust_remote_code=True
+        )
+        
+        print(f"Loading {chat_name}...")
+        chat_model = AutoModelForCausalLM.from_pretrained(
+            chat_name, 
+            torch_dtype=torch.bfloat16, 
+            token=hf_token,
+            device_map="cpu",
+            trust_remote_code=True
+        )
+
+        diffs = calculate_layer_diffs(base_model, chat_model)
+        
+        # Clean up memory
+        del base_model
+        del chat_model
+        torch.cuda.empty_cache()
+
+        img_2d = visualize_2d_heatmap(diffs, base_name, chat_name)
+        plot_3d = visualize_3d_surface(diffs)
+        
+        return img_2d, plot_3d
+        
+    except Exception as e:
+        raise gr.Error(f"Error processing models: {str(e)}")
+
+# --- Gradio UI Layout ---
+with gr.Blocks(title="Model Diff Visualizer") as demo:
+    gr.Markdown("# 🧠 LLM Weight Difference Visualizer")
+    gr.Markdown("Compare the weights of a Base model vs. its Instruct/Chat tuned version layer by layer.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            base_input = gr.Textbox(label="Base Model Name", placeholder="e.g., meta-llama/Llama-2-7b-hf")
+            chat_input = gr.Textbox(label="Chat/Tuned Model Name", placeholder="e.g., meta-llama/Llama-2-7b-chat-hf")
+            token_input = gr.Textbox(label="Hugging Face Token (Optional)", type="password")
+            submit_btn = gr.Button("🚀 Analyze Differences", variant="primary")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 2D Layer-wise Heatmap")
+            output_2d = gr.Image(label="2D Visualization", type="pil")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 3D Interactive Landscape")
+            output_3d = gr.Plot(label="3D Visualization")
+
+    submit_btn.click(
+        fn=process_models,
+        inputs=[base_input, chat_input, token_input],
+        outputs=[output_2d, output_3d]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=False, server_port=7860)