Update app.py

app.py

@@ -86,15 +86,10 @@ 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")
 
-compare_button_clicked = st.button("Compare Models", key="compare_button")
+if "compare_button_clicked" not in st.session_state:
+    st.session_state.compare_button_clicked = False
 
-if compare_button_clicked:
-    st.session_state["compare_button_disabled"] = True
-
-if "compare_button_disabled" not in st.session_state:
-    st.session_state["compare_button_disabled"] = False
-
-if st.session_state["compare_button_disabled"]:
+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)
@@ -127,6 +122,7 @@ if st.session_state["compare_button_disabled"]:
                 st.error(f"Error loading tokenizers: {e}")
         else:
             st.error("Please enter both model IDs.")
-        st.session_state["compare_button_disabled"] = False
+        st.session_state.compare_button_clicked = False
 else:
-    st.button("Compare Models", key="compare_button_disabled", disabled=True)
+    if st.button("Compare Models"):
+        st.session_state.compare_button_clicked = True