app.py

import streamlit as st
import torch
from transformers import AutoModelForCausalLM

def get_model_structure(model_id):
 model = AutoModelForCausalLM.from_pretrained(
 model_id,
 torch_dtype=torch.bfloat16,
 device_map="cpu",
 )
 structure = {k: v.shape for k, v in model.state_dict().items()}
 return structure

def compare_structures(struct1, struct2):
 keys1 = set(struct1.keys())
 keys2 = set(struct2.keys())
 all_keys = keys1.union(keys2)
 
 diff = []
 for key in all_keys:
 shape1 = struct1.get(key)
 shape2 = struct2.get(key)
 if shape1 != shape2:
 diff.append((key, shape1, shape2))
 return diff

def display_diff(diff):
 left_lines = []
 right_lines = []
 
 for key, shape1, shape2 in diff:
 left_lines.append(f"{key}: {shape1}")
 right_lines.append(f"{key}: {shape2}")
 
 left_html = "<br>".join(left_lines)
 right_html = "<br>".join(right_lines)
 
 return left_html, right_html

st.title("Model Structure Comparison Tool")
model_id1 = st.text_input("Enter the first HuggingFace Model ID")
model_id2 = st.text_input("Enter the second HuggingFace Model ID")

if model_id1 and model_id2:
 struct1 = get_model_structure(model_id1)
 struct2 = get_model_structure(model_id2)
 
 diff = compare_structures(struct1, struct2)
 left_html, right_html = display_diff(diff)
 
 st.write("### Comparison Result")
 col1, col2 = st.columns(2)
 
 with col1:
 st.write("### Model 1")
 st.markdown(left_html, unsafe_allow_html=True)
 
 with col2:
 st.write("### Model 2")
 st.markdown(right_html, unsafe_allow_html=True)