Spaces:
Sleeping
Sleeping
File size: 2,351 Bytes
92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 92d7f1e 2ddd1e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import streamlit as st
import torch
from transformers import AutoModelForCausalLM
import difflib
@st.cache_data
def get_model_structure(model_id):
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="cpu",
)
structure = {k: str(v.shape) for k, v in model.state_dict().items()}
return structure
def compare_structures(struct1, struct2):
struct1_lines = [f"{k}: {v}" for k, v in struct1.items()]
struct2_lines = [f"{k}: {v}" for k, v in struct2.items()]
diff = difflib.ndiff(struct1_lines, struct2_lines)
return diff
def display_diff(diff):
left_lines = []
right_lines = []
for line in diff:
if line.startswith('- '):
left_lines.append(f'<span style="background-color: #ffdddd;">{line[2:]}</span>')
right_lines.append('')
elif line.startswith('+ '):
right_lines.append(f'<span style="background-color: #ddffdd;">{line[2:]}</span>')
left_lines.append('')
elif line.startswith(' '):
left_lines.append(line[2:])
right_lines.append(line[2:])
else:
pass
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([1, 1]) # Adjust the ratio to make columns wider
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)
# Apply custom CSS for wider layout
st.markdown(
"""
<style>
.reportview-container .main .block-container {
max-width: 90%;
padding-left: 5%;
padding-right: 5%;
}
.stMarkdown {
white-space: pre-wrap;
}
</style>
""",
unsafe_allow_html=True
)
|