Update app.py

app.py

@@ -99,23 +99,43 @@ 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
-
-    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)
+if "compare_button_clicked" not in st.session_state:
+    st.session_state.compare_button_clicked = False
+
+if st.session_state.compare_button_clicked:
+    with st.spinner('Comparing models and loading tokenizers...'):
+        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 col2:
+                st.write("### Model 2")
+                st.markdown(right_html, unsafe_allow_html=True)
+                
+            # Tokenizer verification
+            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.")
+        st.session_state.compare_button_clicked = False
+else:
+    if st.button("Compare Models"):
+        st.session_state.compare_button_clicked = True