from transformers import AutoModelForCausalLM, AutoTokenizer import torch 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() def calculate_layer_diffs(base_model, chat_model, load_one_at_a_time=False): layer_diffs = [] 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 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 def visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name): num_layers = len(layer_diffs) num_components = len(layer_diffs[0]) # 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 # Create figure with subplots arranged in 2 rows if there are many components 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) # 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)) 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, annot=True, fmt=".9f", cmap="YlGnBu", ax=axs[i], cbar=False, annot_kws={'size': annot_font_size}) 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_yticks(range(num_layers)) axs[i].set_yticklabels(range(num_layers), fontsize=tick_font_size) axs[i].invert_yaxis() # Remove any empty subplots if using 2 rows if num_components > 6: 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 # 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, 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) return visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name) 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 ) iface.launch(share=False, server_port=7860)