Update streamlit_app.py

streamlit_app.py

@@ -65,7 +65,13 @@ def load_stability_model():
 
 embeddings_dict = load_embeddings()
 final_df = load_results()
-stability_model = load_stability_model()
+
+def get_stability_model():
+    if "stability_model" not in st.session_state:
+        with st.spinner("Loading stability model..."):
+            st.session_state.stability_model = load_stability_model()
+    return st.session_state.stability_model
+
 # --------------------------------------------------
 # TOP: Molecule of Interest (Primary Input)
 # --------------------------------------------------
@@ -87,8 +93,10 @@ if top_predict and top_smiles.strip():
     if mol:
         mw = Descriptors.MolWt(mol)
         emb = embeddings_dict.get(top_smiles, torch.randn(256)).unsqueeze(0)
+        model = get_stability_model()
         with torch.no_grad():
-            stab = stability_model(emb).item()
+            stab = model(emb).item()
+
 
         res1, res2 = st.columns(2)
         res1.metric("Zone-IV Shelf-Life (days)", f"{stab:.2f}")
@@ -146,8 +154,10 @@ if st.sidebar.button("Predict"):
         if mol:
             mw = Descriptors.MolWt(mol)
             emb = embeddings_dict.get(user_smiles, torch.randn(256)).unsqueeze(0)
+            model = get_stability_model()
             with torch.no_grad():
-                stab = stability_model(emb).item()
+                stab = model(emb).item()
+
             st.sidebar.success(f"Zone-IV Shelf-life: {stab:.2f} days")
             st.sidebar.info(f"Molecular Weight: {mw:.2f}")
         else: