Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import difflib
 import requests
 import os
@@ -31,7 +31,6 @@ def get_model_structure(model_id):
     save_to_firebase(model_id, structure)
     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()]
@@ -87,23 +86,36 @@ 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, diff_found = display_diff(diff)
-    
-    st.write("### Comparison Result")
-    if not diff_found:
-        st.success("The model structures are identical.")
-    
-    col1, col2 = st.columns([1.5, 1.5])  # Adjust the ratio to make columns wider
+if st.button("Compare Models"):
+    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, diff_found = display_diff(diff)
+        
+        st.write("### Comparison Result")
+        if not diff_found:
+            st.success("The model structures are identical.")
+        
+        col1, col2 = st.columns([1.5, 1.5])  # Adjust the ratio to make columns wider
 
-    with col1:
-        st.write("### Model 1")
-        st.markdown(left_html, unsafe_allow_html=True)
+        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)
+        with col2:
+            st.write("### Model 2")
+            st.markdown(right_html, unsafe_allow_html=True)
+            
+        # Tokenizer verification
+        with st.spinner('Loading tokenizers...'):
+            try:
+                tokenizer1 = AutoTokenizer.from_pretrained(model_id1)
+                tokenizer2 = AutoTokenizer.from_pretrained(model_id2)
+                st.write(f"**{model_id1} Tokenizer Vocab Size**: {tokenizer1.vocab_size}")
+                st.write(f"**{model_id2} Tokenizer Vocab Size**: {tokenizer2.vocab_size}")
+            except Exception as e:
+                st.error(f"Error loading tokenizers: {e}")
+    else:
+        st.error("Please enter both model IDs.")