Update app.py

app.py

@@ -86,30 +86,38 @@ 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 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
+compare_button_clicked = st.button("Compare Models", key="compare_button")
 
-        with col1:
-            st.write("### Model 1")
-            st.markdown(left_html, unsafe_allow_html=True)
+if compare_button_clicked:
+    st.session_state["compare_button_disabled"] = True
 
-        with col2:
-            st.write("### Model 2")
-            st.markdown(right_html, unsafe_allow_html=True)
+if "compare_button_disabled" not in st.session_state:
+    st.session_state["compare_button_disabled"] = False
+
+if st.session_state["compare_button_disabled"]:
+    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.")
             
-        # Tokenizer verification
-        with st.spinner('Loading tokenizers...'):
+            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)
@@ -117,5 +125,8 @@ if st.button("Compare Models"):
                 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.")
+        else:
+            st.error("Please enter both model IDs.")
+        st.session_state["compare_button_disabled"] = False
+else:
+    st.button("Compare Models", key="compare_button_disabled", disabled=True)