Update src/streamlit_app.py

src/streamlit_app.py

@@ -301,26 +301,23 @@ tab1, tab2 = st.tabs(["🔬 Fingerprint Model", "🧬 GCN Model"])
 with tab1:
     st.subheader("Fingerprint-based Prediction")
 
-    # ✅ SMILES validation function
-    def is_valid_smiles(smiles):
-        mol = Chem.MolFromSmiles(smiles)
-        return mol is not None
-
     with st.form("fp_form"):
         smiles_fp = st.text_input("Enter SMILES", "CCO")
         show_debug_fp = st.checkbox("🐞 Show Debug Info (raw score/logit)", key="fp_debug")
         predict_btn = st.form_submit_button("🔍 Predict")
 
     if predict_btn:
-        with st.spinner("Predicting..."):
-            # ✅ Early validation
-            if not is_valid_smiles(smiles_fp):
-                st.error("❌ Invalid SMILES input. Please check your string.")
-            else:
+    with st.spinner("Predicting..."):
+        if not is_valid_smiles(smiles_fp):
+            st.error("❌ Invalid SMILES input. Please check your string.")
+        else:
+            try:
                 mol = Chem.MolFromSmiles(smiles_fp)
                 fp = fp_gen.GetFingerprint(mol)
                 arr = np.array(fp).reshape(1, -1)
-                tensor = torch.tensor(arr).float()
+                tensor = torch.tensor(arr).float().to("cpu")
+                fp_model.to("cpu")  # Ensure model is on CPU
+
                 with torch.no_grad():
                     output = fp_model(tensor)
                     prob = torch.sigmoid(output).item()
@@ -344,6 +341,10 @@ with tab1:
 
                 st.plotly_chart(plot_distribution(df, 'fp', prob), use_container_width=True)
 
+            except Exception as e:
+                st.error(f"Prediction error: {str(e)}")
+
+
     with st.expander("📌 Example SMILES to Try"):
         st.markdown("""
         - `CCO` (Ethanol)  
@@ -374,10 +375,6 @@ with tab2:
 
     SUPPORTED_ATOMS = {1, 6, 7, 8, 9, 16, 17, 35, 53}  # H, C, N, O, F, S, Cl, Br, I
 
-    # ✅ Molecule validation functions
-    def is_valid_smiles(smiles):
-        mol = Chem.MolFromSmiles(smiles)
-        return mol is not None
 
     def is_supported(mol):
         return all(atom.GetAtomicNum() in SUPPORTED_ATOMS for atom in mol.GetAtoms())