Update app.py

app.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import streamlit_ext as ste
+
 from trainer import Trainer
 import random
 from rdkit.Chem import Draw
@@ -118,8 +120,9 @@ with st.sidebar:
 
 
 
-if submitted or ("submitted" in st.session_state):
-    st.session_state["submitted"] = True
+if submitted:
+# if submitted or ("submitted" in st.session_state):
+    # st.session_state["submitted"] = True
     config = model_configs[model_name]
 
     config.inference_sample_num = molecule_num_input
@@ -139,7 +142,8 @@ if submitted or ("submitted" in st.session_state):
 
     with open(f'experiments/inference/{model_name}/inference_drugs.txt') as f:
         inference_drugs = f.read()
-    st.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'DrugGEN-{model_name}_denovo_mols.smi', mime="text/plain")
+    # st.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'DrugGEN-{model_name}_denovo_mols.smi', mime="text/plain")
+    ste.download_button(label="Click to download generated molecules", data=inference_drugs, file_name=f'DrugGEN-{model_name}_denovo_mols.smi', mime="text/plain")
 
     
     st.write("Structures of randomly selected 12 de novo molecules from the inference set:")