Spaces:
Running
Running
| import streamlit as st | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| import os | |
| import io | |
| 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): | |
| layer_diffs = [] | |
| for base_layer, chat_layer in zip(base_model.model.layers, chat_model.model.layers): | |
| layer_diff = { | |
| 'input_layernorm': calculate_weight_diff(base_layer.input_layernorm.weight, chat_layer.input_layernorm.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]) | |
| fig, axs = plt.subplots(1, num_components, figsize=(24, 8)) | |
| fig.suptitle(f"{base_model_name} <> {chat_model_name}", fontsize=16) | |
| for i, component in enumerate(layer_diffs[0].keys()): | |
| component_diffs = [[layer_diff[component]] for layer_diff in layer_diffs] | |
| sns.heatmap(component_diffs, annot=True, fmt=".6f", cmap="YlGnBu", ax=axs[i], cbar_kws={"shrink": 0.8}) | |
| axs[i].set_title(component) | |
| axs[i].set_xlabel("Layer") | |
| axs[i].set_ylabel("Difference") | |
| axs[i].set_xticks([]) | |
| axs[i].set_yticks(range(num_layers)) | |
| axs[i].set_yticklabels(range(num_layers)) | |
| axs[i].invert_yaxis() | |
| plt.tight_layout() | |
| return fig | |
| def main(): | |
| st.set_page_config( | |
| page_title="Model Weight Comparator", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| st.title("Language Model Weight Comparator") | |
| # Config sidebar for input parameters | |
| with st.sidebar: | |
| st.header("Configuration") | |
| base_model_name = st.text_input( | |
| "Base Model Name", | |
| value="meta-llama/Meta-Llama-3-70B-Instruct", | |
| help="Enter the name of the base model from Hugging Face" | |
| ) | |
| chat_model_name = st.text_input( | |
| "Chat Model Name", | |
| value="mattshumer/Reflection-Llama-3.1-70B", | |
| help="Enter the name of the chat model from Hugging Face" | |
| ) | |
| if st.button("Compare Models"): | |
| if not base_model_name or not chat_model_name: | |
| st.error("Please enter both model names") | |
| return | |
| try: | |
| st.info("Loading models... This might take some time.") | |
| base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.bfloat16) | |
| chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=torch.bfloat16) | |
| st.info("Calculating weight differences...") | |
| layer_diffs = calculate_layer_diffs(base_model, chat_model) | |
| st.info("Generating visualization...") | |
| fig = visualize_layer_diffs(layer_diffs, base_model_name, chat_model_name) | |
| st.pyplot(fig) | |
| # visualization | |
| buf = io.BytesIO() | |
| fig.savefig(buf, format='png', dpi=300, bbox_inches='tight') | |
| buf.seek(0) | |
| st.download_button( | |
| label="Download Visualization", | |
| data=buf, | |
| file_name="model_comparison.png", | |
| mime="image/png" | |
| ) | |
| except Exception as e: | |
| st.error(f"An error occurred: {str(e)}") | |
| if __name__ == "__main__": | |
| main() |