Update Optimize a batch page, fix UI

app.py

@@ -101,6 +101,8 @@ def mol_passes_filters_custom(mol,
         return 'Isomeric'
     if Chem.MolFromSmiles(smiles) is None:
         return 'Isomeric'
+    if not check_vocab(Chem.MolToSmiles(mol)):
+        return 'NoVocab'
     return 'YES'
 
 def penalized_logp_standard(mol):
@@ -169,12 +171,26 @@ def set_step(id):
 def get_step():
     return st.session_state['current_step']
 
+
+vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
+vocab_set = set(vocab)
+vocab = Vocab(vocab)
+
+def check_vocab(smiles):
+    cset = set()
+    mol = MolTree(smiles)
+    for c in mol.nodes:
+        cset.add(c.smiles)
+    return cset.issubset(vocab_set)
+
+
 @st.cache_resource
 def load_model():
-    vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
-    vocab = Vocab(vocab)
     model = JTPropVAE(vocab, 450, 56, 20, 3)
-    model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
+    if torch.cuda.is_available():
+        model.load_state_dict(torch.load(model_path))
+    else:
+        model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
     return model
 
 from streamlit_lottie import st_lottie
@@ -211,6 +227,33 @@ def oam_sidebar(step):
     st.markdown('|')
     st.markdown("<h4 style='color: "+color_ls[3]+"'>Finished</h4>",unsafe_allow_html=True)
 
+def oab_sidebar(step):
+    st.title('**Optimize a batch**')
+    prog_bar = st.progress(0)
+    # cur_step = get_step()
+    if step == 0: prog_bar.progress(0)
+    if step == 1: prog_bar.progress(20)
+    if step == 2: prog_bar.progress(40)
+    if step == 3: prog_bar.progress(60)
+    if step == 4: prog_bar.progress(80)
+    if step == 5: prog_bar.progress(100)
+    st.markdown('\n')
+    
+    # st.markdown(get_step())
+    color_ls = colorize_step(6,step)
+    
+    st.markdown("<h4 style='color: "+color_ls[0]+"'>Upload SMILES file</h4>",unsafe_allow_html=True)
+    st.markdown('|')
+    st.markdown("<h4 style='color: "+color_ls[1]+"'>Checking SMILES</h4>",unsafe_allow_html=True)
+    st.markdown('|')
+    st.markdown("<h4 style='color: "+color_ls[2]+"'>Select scores</h4>",unsafe_allow_html=True)
+    st.markdown('|')
+    st.markdown("<h4 style='color: "+color_ls[3]+"'>Choose settings</h4>",unsafe_allow_html=True)
+    st.markdown('|')
+    st.markdown("<h4 style='color: "+color_ls[4]+"'>Optimizing a batch</h4>",unsafe_allow_html=True)
+    st.markdown('|')
+    st.markdown("<h4 style='color: "+color_ls[5]+"'>Finished</h4>",unsafe_allow_html=True)
+
 # @st.cache_data(experimental_allow_widgets=True)
 
 # if 'sidebar_con' not in st.session_state:
@@ -299,24 +342,37 @@ def Optimize_a_molecule():
     with sidebar_con.container():
             set_step(0)
             oam_sidebar(0)
-    oab_sel_container = st.empty()
+    # oab_sel_container = st.empty()
     if 'checked_single' not in st.session_state:
         st.session_state.checked_single = 'NO'
-    with oab_sel_container.container():
-            oam_sel_col = st.columns([3,7])
-            oam_sel_col[0].selectbox("Select an example",options=['-','Sorafenib','Pazopanib','Sunitinib'],key='mode',on_change=reset_oam_state)
-            if st.session_state.mode == '-':
-                st.session_state.smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',max_chars=200)
-            elif st.session_state.mode == 'Sorafenib':
-                st.session_state.smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
-                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',max_chars=200,disabled=True)
-            elif st.session_state.mode == 'Pazopanib':
-                st.session_state.smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
-                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',max_chars=200,disabled=True)
-            elif st.session_state.mode == 'Sunitinib':
-                st.session_state.smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
-                oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C',max_chars=200,disabled=True)
-            st.columns([4,1,4])[1].button('Check SMILES',on_click=check_single,args=[st.session_state.smiles])
+    if 'mode' not in st.session_state:
+        st.session_state.mode = 0
+    if 'single_optimized' not in st.session_state:
+        st.session_state.single_optimized = False
+    if 'smiles_checked' not in st.session_state:
+        st.session_state.smiles_checked = False
+    # with oab_sel_container.container():
+    oam_sel_col = st.columns([3,7])
+    mode = oam_sel_col[0].selectbox("Select an example",options=['-','Sorafenib','Pazopanib','Sunitinib'],on_change=reset_oam_state,index=st.session_state.mode)
+    if mode == '-':
+        smiles = oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):',max_chars=200)
+        st.session_state.mode = 0
+    elif mode == 'Sorafenib':
+        smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
+        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',max_chars=200,disabled=True)
+        st.session_state.mode = 1
+    elif mode == 'Pazopanib':
+        smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
+        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N',max_chars=200,disabled=True)
+        st.session_state.mode = 2
+    elif mode == 'Sunitinib':
+        smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
+        oam_sel_col[1].text_input('Enter a SMILES string (max 200 chars):','CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C',max_chars=200,disabled=True)
+        st.session_state.mode = 3
+    st.columns([4,1,4])[1].button('Check SMILES',on_click=check_single,args=[smiles])
+    
+    if 'optim_single_butt' not in locals():
+        optim_single_butt = False
     
     check_single_con = st.empty()
     if 'smiles_selected' in st.session_state:
@@ -325,19 +381,19 @@ def Optimize_a_molecule():
                 if 'checked_single' in st.session_state:
                     if st.session_state.checked_single == 'EnterError':
                         st.markdown("<p style='text-align: center; color: red;'><b>Please enter a SMILES string.</b></p>",unsafe_allow_html=True)
-                        sidebar_con.empty()
+                        # sidebar_con.empty()
                         with sidebar_con.container():
                                 set_step(0)
                                 oam_sidebar(0)
                     elif st.session_state.checked_single == 'MolError':
                         st.markdown("<p style='text-align: center; color: red;'><b>SMILES is invalid. Please enter a valid SMILES.</b></p>",unsafe_allow_html=True)
-                        sidebar_con.empty()
+                        # sidebar_con.empty()
                         with sidebar_con.container():
                                 set_step(0)
                                 oam_sidebar(0)
                     elif st.session_state.checked_single == 'YES':
                         st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
-                        st.code(st.session_state.smiles)
+                        st.code(Chem.CanonSmiles(st.session_state.smiles))
                         st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
                         mol = Chem.MolFromSmiles(st.session_state.smiles)
                         imgByteArr = io.BytesIO()
@@ -355,7 +411,7 @@ def Optimize_a_molecule():
                         st.session_state.smiles_checked = True
                         # render_sidebar()
                         # col1, col2, col3 = st.columns(3)
-                        sidebar_con.empty()
+                        # sidebar_con.empty()
                         with sidebar_con.container():
                                 set_step(1)
                                 oam_sidebar(1)
@@ -363,41 +419,45 @@ def Optimize_a_molecule():
                                 st.slider('Choose learning rate: ',0.0,5.0,0.4,key='lr_s')
                                 st.slider('Choose similarity cutoff: ',0.0,1.0,0.4,key='sim_cutoff_s')
                                 st.slider('Choose number of iterations: ',1,100,80,key='n_iter_s')
-                                st.session_state.optim_single_butt = st.form_submit_button("Optimize")
+                                optim_single_butt = st.form_submit_button("Optimize")
                     else:
                         st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
                         st.code(st.session_state.smiles)
-                        st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed. Please use another molecule.</b></p>",unsafe_allow_html=True)
-                        sidebar_con.empty()
+                        if st.session_state.checked_single == 'NoVocab':
+                            st.markdown("<p style='text-align: center; color: red;'><b>The molecule contains unavailable vocab(s). Please use another molecule.</b></p>",unsafe_allow_html=True)
+                        else:
+                            st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed. Please use another molecule.</b></p>",unsafe_allow_html=True)
+                        # sidebar_con.empty()
                         with sidebar_con.container():
                                 set_step(0)
                                 oam_sidebar(0)
         else: check_single_con.empty()
     
     optim_single_con = st.empty()
-    if 'optim_single_butt' in st.session_state:
-        if st.session_state.optim_single_butt and st.session_state.smiles_checked:
-            sidebar_con.empty()
+    if st.session_state.smiles_checked:
+        if optim_single_butt:
+            # sidebar_con.empty()
             with sidebar_con.container():
                     set_step(2)
                     oam_sidebar(2)
+            
+            ani_con = st.empty()
+            with ani_con.container():
+                st.markdown('Operation in progress. Please wait...')
+                render_animation()
+                model = load_model()
+                st.session_state.new_smiles,st.session_state.sim = optim_single(st.session_state.smiles,model,st.session_state.lr_s,st.session_state.sim_cutoff_s,st.session_state.n_iter_s)
+                st.session_state.single_optimized = True
+            ani_con.empty()
+                # sidebar_con.empty()
+        if st.session_state.single_optimized:
             with optim_single_con.container():
-                ani_con = st.empty()
-                with ani_con.container():
-                    st.markdown('Operation in progress. Please wait...')
-                    render_animation()
-                    new_smiles,sim = optim_single(st.session_state.smiles,st.session_state.lr_s,st.session_state.sim_cutoff_s,st.session_state.n_iter_s)
-                ani_con.empty()
-                sidebar_con.empty()
-                with sidebar_con.container():
-                        set_step(3)
-                        oam_sidebar(3)
-                if new_smiles is None:
+                if st.session_state.new_smiles is None:
                     st.markdown("<h4 style='text-align: center; color: red;'>Cannot optimize! Please choose another setting.</h4>",unsafe_allow_html=True)
                 else:
                     st.markdown("<b style='text-align: center;'>New SMILES</b>",unsafe_allow_html=True)
-                    st.code(new_smiles)
-                    new_mol = Chem.MolFromSmiles(new_smiles)
+                    st.code(st.session_state.new_smiles)
+                    new_mol = Chem.MolFromSmiles(st.session_state.new_smiles)
                     if new_mol is None:
                         st.markdown("<p style='text-align: center; color: red;'>New SMILES is invalid! Please choose another setting.</p>",unsafe_allow_html=True)
                         # st.write('New SMILES is invalid.')
@@ -422,59 +482,299 @@ def Optimize_a_molecule():
                         col32.metric('Cycle score', '%d' % (-st.session_state.new_cycle),'%d'%(-st.session_state.new_cycle+st.session_state.cycle),delta_color='inverse')
                         col42.metric('Penalized LogP', '%.5f' % (st.session_state.new_pen_p),'%.5f'%(st.session_state.new_pen_p-st.session_state.pen_p))
                         # st.metric('New penalized logP score','%.5f' % (new_score), '%.5f'%(new_score-score))
-                        st.metric('Similarity','%.5f' % (sim))
+                        st.metric('Similarity','%.5f' % (st.session_state.sim))
                         # st.write('Caching ZINC20 if necessary...')
                         with st.spinner("Caching ZINC20 if necessary..."):
-                            if buy(new_smiles, catalog='zinc20',canonicalize=True):
+                            if buy(st.session_state.new_smiles, catalog='zinc20',canonicalize=True):
                                 st.write('This molecule exists.')
                                 st.markdown("<h3 style='text-align: center; color: darkturquoise;'><b>This molecule exists.</h3>",unsafe_allow_html=True)
                             else:
                                 # st.write('THIS MOLECULE DOES NOT EXIST!')
                                 st.markdown("<h3 style='text-align: center; color: mediumseagreen;'>THIS MOLECULE DOES NOT EXIST!</h3>",unsafe_allow_html=True)
                         st.markdown("<p style='text-align: center; color: grey;'>Checked using molbloom</p>",unsafe_allow_html=True)
+            with sidebar_con.container():
+                    set_step(3)
+                    oam_sidebar(3)
         else: optim_single_con.empty()
+    else: optim_single_con.empty()
 
 def check_single(smiles):
     # render_view()
     st.session_state.smiles_selected = True
-    check_single_con = st.empty()
+    st.session_state.smiles = smiles
+    # check_single_con = st.empty()
     
     # optim = False
-    with check_single_con.container():
-        if len(smiles) == 0:
-            st.session_state.checked_single = 'EnterError'
+    # with check_single_con.container():
+    if len(smiles) == 0:
+        st.session_state.checked_single = 'EnterError'
+    else:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            st.session_state.checked_single = 'MolError'
         else:
-            mol = Chem.MolFromSmiles(smiles)
-            if mol is None:
-                st.session_state.checked_single = 'MolError'
-            else:
-                st.session_state.smiles = Chem.MolToSmiles(mol)
-                st.session_state.logp,st.session_state.sa,st.session_state.cycle,st.session_state.pen_p = penalized_logp_standard(mol)
-                moses_passed = mol_passes_filters_custom(mol)
-                st.session_state.checked_single = moses_passed
+            st.session_state.smiles = Chem.MolToSmiles(mol)
+            st.session_state.logp,st.session_state.sa,st.session_state.cycle,st.session_state.pen_p = penalized_logp_standard(mol)
+            moses_passed = mol_passes_filters_custom(mol)
+            st.session_state.checked_single = moses_passed
 
 
-def optim_single(smiles,lr,sim_cutoff,n_iter):
-    # render_view()
-    
-        
-        model = load_model()
+def optim_single(smiles,model,lr,sim_cutoff,n_iter):
 
         new_smiles,sim = model.optimize(smiles, sim_cutoff=sim_cutoff, lr=lr, num_iter=n_iter)
-        
+
         return new_smiles,sim
 
         
             
 def Optimize_a_batch():
-    st.empty()
+    st.session_state.sc_name = ['logp','sa','cycle','pen_logp']
+    st.session_state.new_sc_name = ['new_'+n for n in st.session_state.sc_name]
+    st.markdown("<h2 style='text-align: center;'>Optimize a batch</h2>",unsafe_allow_html=True)
+    with st.expander(':snowman: :blue[Instruction]'):
+        guide = """<h4 style='color:tomato;'>Steps to optimize a molucule</h4>
+        1. Upload a text file with SMILES string on each line.</br>
+        2. Check the SMILES strings to make sure that they are valid and pass MOSES filters.</br>
+        3. Select scores to calculate (penalized LogP included). Keep passed SMILES and calculate selected scores.</br>
+        4. Configure the settings to generate new molecules. The new molecules should have higher penalized LogP values.</br>
+            - Learning rate: How 'far' from each molecule that you want to search</br>
+            - Similarity cutoff: How 'similar' to each molecule that you want to search</br>
+            - Number of iterations: Number of generation trials per molecule</br>
+        5. <i>(Optional)</i> You can download the dataframe at any steps as *.csv file.</br>
+        <h4 style='color:darkturquoise;'>Annotation</h4>
+        <b>SMILES</b> - Simplified molecular-input line-entry system</br>
+        <b>LogP</b> - The log of the partition coefficient of a solute between octanol and water, at near infinite dilution</br>
+        <b>SA score</b> - Synthetic Accessibility Score (lower is better)</br>
+        <b>Cycle score</b> - A number of carbon rings of size larger than 6 (lower is better)</br>
+        <b>Penalized LogP</b> - Standardized score of <i>LogP - SA score - Cycle score</i></br>
+        <b>Similarity</b> - Molecular similarity is calculated via Morgan fingerprint of radius 2 with Tanimoto similarity</br>
+        """
+        st.markdown(guide,unsafe_allow_html=True)
+        
+    with st.sidebar:
+        sidebar_con = st.empty()
+    # sidebar_con.empty()
+    with sidebar_con.container():
+            set_step(0)
+            oab_sidebar(0)
+    oab_upl_container = st.empty()
+    if 'smiles_upload_change' not in st.session_state:
+        st.session_state.smiles_upload_change = False
+    if 'checked_batch' not in st.session_state:
+        st.session_state.checked_batch = False
+    if 'batch_left_checked' not in st.session_state:
+        st.session_state.batch_left_checked = False
+    if 'scores_calculated' not in st.session_state:
+        st.session_state.scores_calculated = False
+    if 'batch_optimized' not in st.session_state:
+        st.session_state.batch_optimized = False
+    
+    with oab_upl_container.container():
+        st.session_state['smiles_file'] = st.file_uploader("Upload a text file with SMILES on each line :sparkles:",on_change=reset_oab_state)
+    if 'check_batch_butt' not in locals():
+        check_batch_butt = False
+    
+    if st.session_state['smiles_file'] is not None:
+        if st.session_state.smiles_upload_change:
+            smiles_list = io.StringIO(st.session_state.smiles_file.getvalue().decode("utf-8"))
+            smiles_list = list(smiles_list.getvalue().rstrip().split('\n'))
+            st.markdown('Number of SMILES: '+str(len(smiles_list)))
+            if len(smiles_list) == 1:
+                    st.markdown("<p style='text-align: center; color: red;'><b>Please use <i>Optimize a molecule</i> tab.</b></p>",unsafe_allow_html=True)
+                    with sidebar_con.container():
+                        set_step(0)
+                        oab_sidebar(0)
+            else:
+                st.session_state['df'] = pd.DataFrame({'SMILES':smiles_list})
+                st.dataframe(st.session_state['df'],use_container_width=True)
+                check_batch_butt = st.button('Check SMILES')
+        else:
+            # if not st.session_state.checked_batch:
+            if st.session_state['smiles_file'] is not None:
+                st.dataframe(st.session_state['df'],use_container_width=True)
+                # st.button('Check SMILES',on_click=check_batch,args=[smiles_list],key='check_batch_butt')
+
+    if check_batch_butt:
+        if st.session_state.smiles_upload_change:
+            with sidebar_con.container():
+                    set_step(1)
+                    oab_sidebar(1)
+            check_batch(list(st.session_state['df'].SMILES))
+            st.session_state.smiles_upload_change = False
+
+    if 'calc_batch_butt' not in locals():
+        calc_batch_butt = False
+    check_batch_con = st.empty()
+    calc_batch_con = st.empty()
+    if st.session_state.checked_batch:
+        with check_batch_con.container():
+            passed_num = st.session_state['df'][st.session_state['df'].checked != 'invalid'].shape[0]
+            st.markdown('Number of passed SMILES: '+str(passed_num))
+            st.dataframe(st.session_state['df'],use_container_width=True)
+            if passed_num == 0:
+                st.markdown("<p style='text-align: center; color: red;'><b>The uploaded file contains no suitable SMILES string.</b></p>",unsafe_allow_html=True)
+                st.session_state.batch_left_checked = False
+                with sidebar_con.container():
+                        set_step(0)
+                        oab_sidebar(0)
+            else:
+                st.session_state.batch_left_checked = True
+            df = st.session_state['df']
+            download_df(df,0)
+            choose_score_con = st.empty()
+            if st.session_state.batch_left_checked:
+                with sidebar_con.container():
+                    set_step(2)
+                    oab_sidebar(2)
+                with choose_score_con.container():
+                    with st.form("Choose score to calculate"):
+                        st.markdown("<h4>Choose score(s) to calculate</h4>",unsafe_allow_html=True)
+                        st.caption('Penalized LogP is always calculated.')
+                        st.checkbox('LogP',key='logp_cal')
+                        st.checkbox('SA score',key='sa_cal')
+                        st.checkbox('Cycle score',key='cycle_cal')
+                        calc_batch_butt = st.form_submit_button("Keep passed SMILES and calculate scores")
+            else:
+                choose_score_con.empty()
+    else:
+        check_batch_con.empty()
+    
+    if 'optim_batch_butt' not in locals():
+        optim_batch_butt = False
+    # if 'calc_batch_butt' in st.session_state:
+    if calc_batch_butt and st.session_state.batch_left_checked:
+        # if not st.session_state.scores_calculated:
+        smiles_list = list(st.session_state['df'][st.session_state['df'].checked != 'invalid'].checked)
+        st.session_state.score_df = calc_scores(smiles_list)
+        st.session_state.batch_optimized = False
+    if st.session_state.scores_calculated:
+        calc_batch_con.empty()
+        with calc_batch_con.container():
+            st.dataframe(st.session_state.score_df,use_container_width=True)
+            score_df = st.session_state.score_df
+            download_df(score_df,1)
+            with sidebar_con.container():
+                set_step(3)
+                oab_sidebar(3)
+            with st.form(":gear: Settings"):
+                    st.slider('Choose learning rate: ',0.0,5.0,0.4,key='lr_b')
+                    st.slider('Choose similarity cutoff: ',0.0,1.0,0.4,key='sim_cutoff_b')
+                    st.slider('Choose number of iterations: ',1,100,80,key='n_iter_b')
+                    optim_batch_butt = st.form_submit_button("Optimize")
+    else:
+        calc_batch_con.empty()
+    
+    
+    optim_batch_con = st.empty()
+    ani_con = st.empty()
+    if optim_batch_butt and st.session_state.scores_calculated:
+        optim_batch_con.empty()
+        with sidebar_con.container():
+                set_step(4)
+                oab_sidebar(4)
+        with ani_con.container():
+            st.markdown('Operation in progress. Please wait...')
+            gen_results = []
+            render_animation()
+            st.markdown('Generating new SMILES string(s)...')
+            model = load_model()
+            for canon_smiles in stqdm(list(st.session_state.score_df.SMILES)):
+                gen_results.append(optim_single(canon_smiles,model,st.session_state.lr_b,st.session_state.sim_cutoff_b,st.session_state.n_iter_b))
+            st.markdown('Checking generated SMILES string(s) ...')
+            st.session_state.new_score_df = calc_scores_new(gen_results)
+        ani_con.empty()
+    if st.session_state.batch_optimized:
+        with sidebar_con.container():
+                set_step(5)
+                oab_sidebar(5)
+        with optim_batch_con.container():
+            new_score_df = st.session_state.new_score_df
+            # new_score_df.style.applymap(highlight_result, subset=pd.IndexSlice[:, ['new_smiles']])
+            st.markdown("<h3 style='text-align: center; color: mediumseagreen;'>RESULTS</h3>",unsafe_allow_html=True)
+            st.dataframe(new_score_df.style.applymap(highlight_result, subset=pd.IndexSlice[:, ['new_smiles']]),use_container_width=True)
+            download_df(new_score_df,3)
+    else:
+        optim_batch_con.empty()
+        
+
+
+def check_batch(smiles_list):
+    check = []
+    for smi in stqdm(smiles_list):
+            mol = Chem.MolFromSmiles(smi)
+            if (mol is not None) and (mol_passes_filters_custom(mol) == 'YES'):
+                check.append(Chem.MolToSmiles(mol))
+            else:
+                check.append('invalid')
+    st.session_state['df'] = pd.concat([st.session_state['df'],pd.DataFrame({'checked':check})],axis=1)
+    st.session_state.checked_batch = True
+    # return check
+
+def calc_scores(smiles_list):
+    score_df = pd.concat([st.session_state.df[st.session_state.df.checked != 'invalid'].SMILES,pd.DataFrame({'Checked_SMILES':smiles_list})],axis=1)
+    scores =[]
+    for smi in stqdm(smiles_list):
+        logp,sa,cycle,pen_logp = penalized_logp_standard(Chem.MolFromSmiles(smi))
+        scores+=[(logp,sa,cycle,pen_logp)]
+    s_df = pd.DataFrame(scores,columns=st.session_state.sc_name)
+    for n, checked in zip(st.session_state.sc_name,[st.session_state.logp_cal,st.session_state.sa_cal,st.session_state.cycle_cal,True]):
+        if checked:
+            score_df = pd.concat([score_df,s_df[n]],axis=1)
+    st.session_state.scores_calculated = True
+    return score_df
+
+def calc_scores_new(result):
+    new_scores =[]
+    for new_smiles,sim in stqdm(result):
+        if new_smiles is None:
+            new_scores+=[('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)]
+        else:
+            new_mol = Chem.MolFromSmiles(new_smiles)
+            if new_mol is None:
+                new_scores+=[('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)]
+            else:
+                logp,sa,cycle,pen_logp = penalized_logp_standard(new_mol)
+                new_scores+=[(new_smiles,sim,logp,sa,cycle,pen_logp)]
+    new_col = ['new_smiles','sim']+st.session_state.new_sc_name
+    s_df = pd.DataFrame(new_scores,columns=new_col)
+    new_score_df = st.session_state.score_df
+    for n, checked in zip(new_col,[True, True,st.session_state.logp_cal,st.session_state.sa_cal,st.session_state.cycle_cal,True]):
+        if checked:
+            new_score_df = pd.concat([new_score_df,s_df[n]],axis=1)
+    st.session_state.batch_optimized = True
+    return new_score_df
+
+def highlight_result(value):
+    if value == 'invalid': color = 'tomato'
+    else: color = 'mediumseagreen'
+    return 'color: %s' % color
+
+@st.cache_data(experimental_allow_widgets=True)
+def download_df(df,id):
+    with st.expander(':arrow_down: Download this dataframe'):
+        st.markdown("<h4 style='color:tomato;'>Select column(s) to save:</h4>",unsafe_allow_html=True)
+        for col in df.columns:
+            st.checkbox(col,key=str(id)+'_col_'+str(col),value=True)
+        st.text_input('File name (.csv):','dataframe',key=str(id)+'_file_name')
+        
+        ste.download_button('Download',df_to_file(df[[col for col in df.columns if st.session_state[str(id)+'_col_'+str(col)]]]),st.session_state[str(id)+'_file_name']+'.csv')
 
 def reset_oam_state():
     st.session_state.smiles_selected = False
     st.session_state.checked_single = 'NO'
     st.session_state.smiles_checked = False
+    st.session_state.single_optimized = False
     set_step(0)
     
+def reset_oab_state():
+    st.session_state.smiles_upload_change = True
+    st.session_state.smiles_uploaded = False
+    st.session_state.checked_batch = False
+    st.session_state.batch_left_checked = False
+    st.session_state.scores_calculated = False
+    st.session_state.batch_optimized = False
+    set_step(0)
+
 def rerun():
     st.experimental_rerun()
 def render_view():

app_backup.py

@@ -1,437 +0,0 @@
-from pathlib import Path
-import torch
-from st_on_hover_tabs import on_hover_tabs
-import streamlit as st
-st.set_page_config(layout="wide")
-
-model_path = './model.iter-685000'
-
-import sys, os
-import rdkit
-import rdkit.Chem as Chem
-from rdkit.Chem.Draw import MolToImage
-from rdkit.Chem import Descriptors
-import sascorer
-import networkx as nx
-from stqdm import stqdm
-import base64, io
-import pandas as pd
-import streamlit_ext as ste
-
-os.environ['KMP_DUPLICATE_LIB_OK']='True'
-
-sys.path.append('%s/fast_jtnn/' % os.path.dirname(os.path.realpath(__file__)))
-from mol_tree import Vocab, MolTree
-from jtprop_vae import JTPropVAE
-from molbloom import buy
-
-css='''
-[data-testid="metric-container"] {
-    width: fit-content;
-    margin: auto;
-}
-
-[data-testid="metric-container"] > div {
-    width: fit-content;
-    margin: auto;
-}
-
-[data-testid="metric-container"] label {
-    width: fit-content;
-    margin: auto;
-
-[data-testid="stDataFrameResizable"] {
-    width: fit-content;
-    margin: auto;
-}
-}
-'''
-
-st.markdown(f'<style>{css}</style>',unsafe_allow_html=True)
-
-s_buff = io.StringIO()
-def img_to_bytes(img_path):
-    img_bytes = Path(img_path).read_bytes()
-    encoded = base64.b64encode(img_bytes).decode()
-    return encoded
-def img_to_html(img_path):
-    img_html = "<img src='data:image/png;base64,{}' class='img-fluid' style='max-width: 500px;'>".format(
-      img_to_bytes(img_path)
-    )
-    return img_html
-
-_mcf = pd.read_csv('./mcf.csv')
-_pains = pd.read_csv('./wehi_pains.csv',
-                     names=['smarts', 'names'])
-_mcf_filters = [Chem.MolFromSmarts(x) for x in
-            _mcf['smarts'].values]
-_pains_filters = [Chem.MolFromSmarts(x) for x in
-            _pains['smarts'].values]
-
-def mol_passes_filters_custom(mol,
-                       allowed=None,
-                       isomericSmiles=False):
-    """
-    Checks if mol
-    * passes MCF and PAINS filters,
-    * has only allowed atoms
-    * is not charged
-    """
-    allowed = allowed or {'C', 'N', 'S', 'O', 'F', 'Cl', 'Br', 'H'}
-    if mol is None:
-        return 'NoMol'
-    ring_info = mol.GetRingInfo()
-    if ring_info.NumRings() != 0 and any(
-            len(x) >= 8 for x in ring_info.AtomRings()
-    ):
-        return 'ManyRings'
-    h_mol = Chem.AddHs(mol)
-    if any(atom.GetFormalCharge() != 0 for atom in mol.GetAtoms()):
-        return 'Charged'
-    if any(atom.GetSymbol() not in allowed for atom in mol.GetAtoms()):
-        return 'AtomNotAllowed'
-    if any(h_mol.HasSubstructMatch(smarts) for smarts in _mcf_filters):
-        return 'MCF'
-    if any(h_mol.HasSubstructMatch(smarts) for smarts in _pains_filters):
-        return 'PAINS'
-    smiles = Chem.MolToSmiles(mol, isomericSmiles=isomericSmiles)
-    if smiles is None or len(smiles) == 0:
-        return 'Isomeric'
-    if Chem.MolFromSmiles(smiles) is None:
-        return 'Isomeric'
-    return 'YES'
-
-def penalized_logp_standard(mol):
-
-    logP_mean = 2.4399606244103639873799239
-    logP_std = 0.9293197802518905481505840
-    SA_mean = -2.4485512208785431553792478
-    SA_std = 0.4603110476923852334429910
-    cycle_mean = -0.0307270378623088931402396
-    cycle_std = 0.2163675785228087178335699
-
-    log_p = Descriptors.MolLogP(mol)
-    SA = -sascorer.calculateScore(mol)
-
-    # cycle score
-    cycle_list = nx.cycle_basis(nx.Graph(Chem.rdmolops.GetAdjacencyMatrix(mol)))
-    if len(cycle_list) == 0:
-        cycle_length = 0
-    else:
-        cycle_length = max([len(j) for j in cycle_list])
-    if cycle_length <= 6:
-        cycle_length = 0
-    else:
-        cycle_length = cycle_length - 6
-    cycle_score = -cycle_length
-    # print(logP_mean)
-
-    standardized_log_p = (log_p - logP_mean) / logP_std
-    standardized_SA = (SA - SA_mean) / SA_std
-    standardized_cycle = (cycle_score - cycle_mean) / cycle_std
-    return log_p,SA,cycle_score,standardized_log_p + standardized_SA + standardized_cycle
-
-lg = rdkit.RDLogger.logger() 
-lg.setLevel(rdkit.RDLogger.CRITICAL)
-
-def About():
-    descrip = '''
-We seek to automate the design of molecules based on specific chemical properties. In computational terms, this task involves continuous embedding and generation of molecular graphs. Our primary contribution is the direct realization of molecular graphs, a task previously approached by generating linear SMILES strings instead of graphs. Our junction tree variational autoencoder generates molecular graphs in two phases, by first generating a tree-structured scaffold over chemical substructures, and then combining them into a molecule with a graph message passing network. This approach allows us to incrementally expand molecules while maintaining chemical validity at every step. We evaluate our model on multiple tasks ranging from molecular generation to optimization. Across these tasks, our model outperforms previous state-of-the-art baselines by a significant margin.
-
-[https://arxiv.org/abs/1802.04364](https://arxiv.org/abs/1802.04364)'''
-    st.markdown(descrip)
-    st.markdown("<p style='text-align: center;'>"+
-                    img_to_html('about.png')+
-                    "</p>", unsafe_allow_html=True)
-    
-def Optimize_a_molecule():
-    st.markdown("<h2 style='text-align: center;'>Optimize a molecule</h2>",unsafe_allow_html=True)
-    with st.expander(':snowman: :blue[Instruction]'):
-        guide = """<h4 style='color:tomato;'>Steps to optimize a molucule</h4>
-        1. Select from examples, or manually enter a valid SMILES string of a molecule.</br>
-        2. Configure the settings to generate a new molecule. The new molecule should have a higher penalized LogP value.</br>
-            - Learning rate: How 'far' from the molecule that you want to search.</br>
-            - Similarity cutoff: How 'similar' to the molecule that you want to search.</br>
-            - Number of iterations: Number of generation trials.</br>
-        <h4 style='color:darkturquoise;'>Annotation</h4>
-        <b>SMILES</b> - Simplified molecular-input line-entry system</br>
-        <b>LogP</b> - The log of the partition coefficient of a solute between octanol and water, at near infinite dilution</br>
-        <b>SA score</b> - Synthetic Accessibility Score (lower is better)</br>
-        <b>Cycle score</b> - A number of carbon rings of size larger than 6 (lower is better)</br>
-        <b>Penalized LogP</b> - Standardized score of <i>LogP - SA score - Cycle score</i></br>
-        <b>Similarity</b> - Molecular similarity is calculated via Morgan fingerprint of radius 2 with Tanimoto similarity</br>
-        """
-        st.markdown(guide,unsafe_allow_html=True)
-    
-    st.selectbox("Select an example",options=['-','Sorafenib','Pazopanib','Sunitinib'],key='mode')
-    if st.session_state.mode == '-':
-        smiles = st.text_input('Enter a SMILES string (max 200 chars):',max_chars=200)
-    elif st.session_state.mode == 'Sorafenib':
-        smiles = 'CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F'
-    elif st.session_state.mode == 'Pazopanib':
-        smiles = 'CC1=C(C=C(C=C1)NC2=NC=CC(=N2)N(C)C3=CC4=NN(C(=C4C=C3)C)C)S(=O)(=O)N'
-    elif st.session_state.mode == 'Sunitinib':
-        smiles = 'CCN(CC)CCNC(=O)C1=C(NC(=C1C)C=C2C3=C(C=CC(=C3)F)NC2=O)C'
-
-    if len(smiles) > 0:
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            st.markdown("<p style='text-align: center; color: red;'><b>SMILES is invalid. Please enter a valid SMILES.</b></p>",unsafe_allow_html=True)
-        else:
-            canon_smiles = Chem.MolToSmiles(mol)
-            st.markdown("<b>Canonicalized SMILES</b>",unsafe_allow_html=True)
-            st.code(canon_smiles)
-            logp,sa,cycle,pen_p = penalized_logp_standard(mol)
-            moses_passed = mol_passes_filters_custom(mol)
-            if moses_passed=='YES':
-                st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
-            else:
-                st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed. Please use another molecule.</b></p>",unsafe_allow_html=True)
-        # with st.columns(3)[1]:
-        # st.markdown("<style>{text-align: center;}</style>",unsafe_allow_html=True)
-            imgByteArr = io.BytesIO()
-            MolToImage(mol,size=(400,400)).save(imgByteArr,format='PNG')
-            st.markdown("<p style='text-align: center;'>"+
-                    f"<img src='data:image/png;base64,{base64.b64encode(imgByteArr.getvalue()).decode()}' class='img-fluid'>"+
-                    "</p>", unsafe_allow_html=True)
-        # st.image(MolToImage(mol,size=(300,300)))
-            col1, col2, col3, col4 = st.columns(4)
-            col1.metric('LogP', '%.5f' % (logp))
-            col2.metric('SA score', '%.5f' % (-sa))
-            col3.metric('Cycle score', '%d' % (-cycle))
-            col4.metric('Penalized LogP', '%.5f' % (pen_p))
-
-        if (mol is not None) and (moses_passed=='YES'):
-        # col1, col2, col3 = st.columns(3)
-            with st.form(":gear:Settings"):
-                    st.slider('Choose learning rate: ',0.0,5.0,0.4,key='lr')
-                    st.slider('Choose similarity cutoff: ',0.0,1.0,0.4,key='sim_cutoff')
-                    st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
-                    optim = st.form_submit_button("Optimize")
-            vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
-            vocab = Vocab(vocab)
-            if optim:
-            # st.write('Testing')
-                # canon_smiles = Chem.MolToSmiles(mol)
-            
-            # with st.columns(3)[1]:
-                with st.spinner("Operation in progress. Please wait."):
-
-                    model = JTPropVAE(vocab, 450, 56, 20, 3)
-
-                    model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
-        
-                    new_smiles,sim = model.optimize(canon_smiles, sim_cutoff=st.session_state.sim_cutoff, lr=st.session_state.lr, num_iter=st.session_state.n_iter)
-        
-                    del model
-                if new_smiles is None:
-                    st.markdown("<h4 style='text-align: center; color: red;'>Cannot optimize! Please choose another setting.</h4>",unsafe_allow_html=True)
-                else:
-                    st.markdown("<b style='text-align: center;'>New SMILES</b>",unsafe_allow_html=True)
-                    st.code(new_smiles)
-                    new_mol = Chem.MolFromSmiles(new_smiles)
-                    if new_mol is None:
-                        st.markdown("<p style='text-align: center; color: red;'>New SMILES is invalid! Please choose another setting.</p>",unsafe_allow_html=True)
-                        # st.write('New SMILES is invalid.')
-                    else:
-                        # st.write('New SMILES molecule:')
-                        imgByteArr = io.BytesIO()
-                        MolToImage(new_mol,size=(400,400)).save(imgByteArr,format='PNG')
-                        st.markdown("<p style='text-align: center;'>"+
-                                    f"<img src='data:image/png;base64,{base64.b64encode(imgByteArr.getvalue()).decode()}' class='img-fluid'>"+
-                                    "</p>", unsafe_allow_html=True)
-                        
-                        new_moses_passed = mol_passes_filters_custom(new_mol)
-                        if new_moses_passed=='YES':
-                            st.markdown("<p style='text-align: center; color: mediumseagreen'>MOSES filters passed successfully.</p>",unsafe_allow_html=True)
-                        else:
-                            st.markdown("<p style='text-align: center; color: red;'><b>MOSES filters passed failed.</b></p>",unsafe_allow_html=True)
-                        new_logp,new_sa,new_cycle,new_pen_p = penalized_logp_standard(new_mol)
-                    # st.write('New penalized logP score: %.5f' % (new_score))
-                        col12, col22, col32, col42 = st.columns(4)
-                        col12.metric('LogP', '%.5f' % (new_logp),'%.5f'%(new_logp-logp))
-                        col22.metric('SA score', '%.5f' % (-new_sa),'%.5f'%(-new_sa+sa),delta_color='inverse')
-                        col32.metric('Cycle score', '%d' % (-new_cycle),'%d'%(-new_cycle+cycle),delta_color='inverse')
-                        col42.metric('Penalized LogP', '%.5f' % (new_pen_p),'%.5f'%(new_pen_p-pen_p))
-                        # st.metric('New penalized logP score','%.5f' % (new_score), '%.5f'%(new_score-score))
-                        st.metric('Similarity','%.5f' % (sim))
-                        # st.write('Caching ZINC20 if necessary...')
-                        with st.spinner("Caching ZINC20 if necessary..."):
-                            if buy(new_smiles, catalog='zinc20',canonicalize=True):
-                                st.write('This molecule exists.')
-                                st.markdown("<h3 style='text-align: center; color: darkturquoise;'><b>This molecule exists.</h3>",unsafe_allow_html=True)
-                            else:
-                                # st.write('THIS MOLECULE DOES NOT EXIST!')
-                                st.markdown("<h3 style='text-align: center; color: mediumseagreen;'>THIS MOLECULE DOES NOT EXIST!</h3>",unsafe_allow_html=True)
-                        st.markdown("<p style='text-align: center; color: grey;'>Checked using molbloom</p>",unsafe_allow_html=True)
-
-def Optimize_a_batch():
-    st.markdown("<h2 style='text-align: center;'>Optimize a batch</h2>",unsafe_allow_html=True)
-    # st.text_input('Enter a SMILES string:','CNC(=O)C1=NC=CC(=C1)OC2=CC=C(C=C2)NC(=O)NC3=CC(=C(C=C3)Cl)C(F)(F)F',key='smiles')
-    with st.expander(':snowman: :blue[Instruction]'):
-        guide = """<h4 style='color:tomato;'>Steps to optimize a batch</h4>
-        1. Upload a text file with SMILES string on each line.</br>
-        2. Check the SMILES strings to make sure that they are valid and pass MOSES filters.</br>
-        3. Select scores to calculate (penalized LogP included). Keep passed SMILES and calculate selected scores.</br>
-        4. Configure the settings to generate new molecules. The new molecules should have higher penalized LogP values.</br>
-            - Learning rate: How 'far' from each molecule that you want to search</br>
-            - Similarity cutoff: How 'similar' to each molecule that you want to search</br>
-            - Number of iterations: Number of generation trials per molecule</br>
-        5. <i>(Optional)</i> You can download the dataframe at any steps as *.csv file.</br>
-        <h4 style='color:darkturquoise;'>Annotation</h4>
-        <b>SMILES</b> - Simplified molecular-input line-entry system</br>
-        <b>LogP</b> - The log of the partition coefficient of a solute between octanol and water, at near infinite dilution</br>
-        <b>SA score</b> - Synthetic Accessibility Score (lower is better)</br>
-        <b>Cycle score</b> - A number of carbon rings of size larger than 6 (lower is better)</br>
-        <b>Penalized LogP</b> - Standardized score of <i>LogP - SA score - Cycle score</i></br>
-        <b>Similarity</b> - Molecular similarity is calculated via Morgan fingerprint of radius 2 with Tanimoto similarity.</br>
-        """
-        st.markdown(guide,unsafe_allow_html=True)
-    st.session_state['smiles_file'] = st.file_uploader("Upload a text file with SMILES on each line :sparkles:")
-    if st.session_state.smiles_file is not None:
-            # smiles_list = [str(line).rstrip('\n') for line in smiles_file]
-        smiles_list = io.StringIO(st.session_state.smiles_file.getvalue().decode("utf-8"))
-        smiles_list = list(smiles_list.getvalue().rstrip().split('\n'))
-        st.markdown('Number of SMILES: '+str(len(smiles_list)))
-        if len(smiles_list) == 1:
-                st.markdown("<p style='text-align: center; color: red;'><b>Please use <i>Optimize a molecule</i> tab.</b></p>",unsafe_allow_html=True)
-        else:
-            def check_smiles():
-                    check = []
-                    for smi in stqdm(smiles_list):
-                            mol = Chem.MolFromSmiles(smi)
-                            if (mol is not None) and (mol_passes_filters_custom(mol) == 'YES'):
-                                check.append(Chem.MolToSmiles(mol))
-                            else:
-                                check.append('invalid')
-                    st.session_state['df'] = pd.concat([df,pd.DataFrame({'checked':check})],axis=1)
-            df = pd.DataFrame({'SMILES':smiles_list})
-            st.dataframe(df,use_container_width=True)
-            st.button('Check SMILES',key='check',on_click=check_smiles)
-                
-        if 'df' in st.session_state:
-        # st.markdown('Number of SMILES: '+str(len(st.session_state.smiles_list)))
-                df = st.session_state['df']
-                st.markdown('Number of passed SMILES: '+str(df[df.checked != 'invalid'].shape[0]))
-                st.dataframe(df,use_container_width=True)
-                download_df(df,0)
-                with st.form("Choose score to calculate"):
-                    st.caption('Penalized LogP is always calculated.')
-                    logp_cal= st.checkbox('LogP',key='logp_cal')
-                    sa_cal= st.checkbox('SA score',key='sa_cal')
-                    cycle_cal= st.checkbox('Cycle score',key='cycle_cal')
-                    calc = st.form_submit_button("Keep passed SMILES and calculate scores")
-                # st.session_state['pen_logp_cal'] = True
-                if calc:
-                    smiles_list = list(df[df.checked != 'invalid'].SMILES)
-                    # st.write(smiles_list)
-                    score_df = pd.DataFrame({'SMILES':smiles_list})
-                    scores =[]
-                    # pen_logp_ls = logp_ls = sa_ls = cycle_ls =[]
-                    st.session_state.sc_name = ['logp','sa','cycle','pen_logp']
-                    for smi in stqdm(smiles_list):
-                        logp,sa,cycle,pen_logp = penalized_logp_standard(Chem.MolFromSmiles(smi))
-                        # st.write(f"{logp}\t{sa}\t{cycle}\t{pen_logp}")
-                        scores+=[(logp,sa,cycle,pen_logp)]
-                    s_df = pd.DataFrame(scores,columns=st.session_state.sc_name)
-                    # st.write(st.session_state.log_ls)
-                    # to_concat = [score_df] + [d for d,checked in zip(dfs,[logp_cal,sa_cal,cycle_cal]) if checked] + [pd.DataFrame({'pen_logp':pen_logp_ls})]
-                    for n, checked in zip(st.session_state.sc_name,[logp_cal,sa_cal,cycle_cal,True]):
-                        if checked:
-                            score_df = pd.concat([score_df,s_df[n]],axis=1)
-                    # score_df['pen_logp'] = st.session_state.pen_logp_ls
-                    st.session_state.score_df = score_df
-                    st.dataframe(score_df,use_container_width=True)
-                    download_df(score_df,1)
-                    def batch_generate(smiles_list,df):
-                        vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
-                        vocab = Vocab(vocab)
-                        container = st.empty()
-                        with container.container():
-                            
-                            st_lottie(get_ani_json(), height=200, width=300)
-                            st.markdown('Please wait...')
-
-                            model = JTPropVAE(vocab, 450, 56, 20, 3)
-
-                            model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
-
-                            new_scores =[]
-                            for canon_smiles in stqdm(smiles_list):
-                                new_smiles,sim = model.optimize(canon_smiles, sim_cutoff=st.session_state.sim_cutoff, lr=st.session_state.lr, num_iter=st.session_state.n_iter)
-                                if new_smiles is None:
-                                    new_scores+=[('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)]
-                                else:
-                                    new_mol = Chem.MolFromSmiles(new_smiles)
-                                    if new_mol is None:
-                                        new_scores+=[('invalid',-100.0,-100.0,-100.0,-100.0,-100.0)]
-                                    else:
-                                        logp,sa,cycle,pen_logp = penalized_logp_standard(new_mol)
-                                        new_scores+=[(new_smiles,sim,logp,sa,cycle,pen_logp)]
-                
-                            del model
-                        with container.container():
-                            sc_name = ['new_'+n for n in st.session_state.sc_name]
-                            s_df = pd.DataFrame(new_scores,columns=['new_smiles','sim']+sc_name)
-                            new_score_df = df
-                            # st.write(st.session_state.log_ls)
-                            # to_concat = [score_df] + [d for d,checked in zip(dfs,[logp_cal,sa_cal,cycle_cal]) if checked] + [pd.DataFrame({'pen_logp':pen_logp_ls})]
-                            for n, checked in zip(['new_smiles','sim']+sc_name,[True, True,st.session_state.logp_cal,st.session_state.sa_cal,st.session_state.cycle_cal,True]):
-                                if checked:
-                                    new_score_df = pd.concat([new_score_df,s_df[n]],axis=1)
-                            st.markdown("<h3 style='text-align: center; color: mediumseagreen;'>RESULTS</h3>",unsafe_allow_html=True)
-                            st.dataframe(new_score_df,use_container_width=True)
-                            download_df(new_score_df,2)
-                    with st.form(":gear: Settings",clear_on_submit=False):
-                        st.slider('Choose learning rate: ',0.0,5.0,0.4,key='lr')
-                        st.slider('Choose similarity cutoff: ',0.0,1.0,0.4,key='sim_cutoff')
-                        st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
-                        optim = st.form_submit_button("Optimize",on_click=batch_generate,args=(smiles_list,score_df))
-                    
-
-from streamlit_lottie import st_lottie
-import requests
-
-def get_ani_json():
-    animation_response = requests.get('https://assets1.lottiefiles.com/packages/lf20_vykpwt8b.json')
-    animation_json = dict()
-    
-    if animation_response.status_code == 200:
-        animation_json = animation_response.json()
-    else:
-        print("Error in the URL")     
-                           
-    return animation_json
-    
-def download_df(df,id):
-    with st.expander(':arrow_down: Download this dataframe'):
-        st.markdown("<h4 style='color:tomato;'>Select column(s) to save:</h4>",unsafe_allow_html=True)
-        for col in df.columns:
-            st.checkbox(col,key=str(id)+'_col_'+str(col))
-        st.text_input('File name (.csv):','dataframe',key=str(id)+'_file_name')
-        
-        ste.download_button('Download',df_to_file(df[[col for col in df.columns if st.session_state[str(id)+'_col_'+str(col)]]]),st.session_state[str(id)+'_file_name']+'.csv')
-
-def df_to_file(df):
-    s_buff.seek(0)
-    df.to_csv(s_buff)
-    return s_buff.getvalue().encode()
-
-st.markdown("<h1 style='text-align: center;'>Junction Tree Variational Autoencoder for Molecular Graph Generation (JTVAE)</h1>",unsafe_allow_html=True)
-st.markdown("<h3 style='text-align: center;'>Wengong Jin, Regina Barzilay, Tommi Jaakkola</h3>",unsafe_allow_html=True)
-# st.markdown("<-- Use the sidebar to explore --")
-st.markdown("<h4 style='text-align: center;'>-- Use the sidebar to explore --</h4>",unsafe_allow_html=True)
-st.markdown('<style>' + open('./style.css').read() + '</style>', unsafe_allow_html=True)
-
-with st.sidebar:
-    # st.header('+')
-    st.markdown("<h5 style='text-align: center; color:grey;'>Explore</h5>",unsafe_allow_html=True)
-    tabs = on_hover_tabs(tabName=['Optimize a molecule', 'Optimize a batch', 'About'], 
-                         iconName=['science', 'batch_prediction', 'info'], default_choice=0)
-
-page_names_to_funcs = {
-    "About": About,
-    "Optimize a molecule":Optimize_a_molecule,
-    "Optimize a batch":Optimize_a_batch
-}
-page_names_to_funcs[tabs]()