Update app.py

app.py

@@ -1,36 +1,40 @@
 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: v.shape for k, v in model.state_dict().items()}
+    structure = {k: str(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))
+    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 key, shape1, shape2 in diff:
-        left_lines.append(f"{key}: {shape1}")
-        right_lines.append(f"{key}: {shape2}")
+    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)
@@ -49,7 +53,7 @@ if model_id1 and model_id2:
     left_html, right_html = display_diff(diff)
     
     st.write("### Comparison Result")
-    col1, col2 = st.columns(2)
+    col1, col2 = st.columns([1, 1])  # Adjust the ratio to make columns wider
     
     with col1:
         st.write("### Model 1")
@@ -58,3 +62,20 @@ if model_id1 and model_id2:
     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
+)