Update UI

app.py

@@ -1,5 +1,9 @@
+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")
+
 import sys, os
 import rdkit
 import rdkit.Chem as Chem
@@ -7,6 +11,8 @@ from rdkit.Chem.Draw import MolToImage
 from rdkit.Chem import Descriptors
 import sascorer
 import networkx as nx
+from stqdm import stqdm
+import base64, io
 
 os.environ['KMP_DUPLICATE_LIB_OK']='True'
 
@@ -15,21 +21,34 @@ from mol_tree import Vocab, MolTree
 from jtprop_vae import JTPropVAE
 from molbloom import buy
 
-
-lg = rdkit.RDLogger.logger() 
-lg.setLevel(rdkit.RDLogger.CRITICAL)
-
-st.header('Junction Tree Variational Autoencoder for Molecular Graph Generation (JTVAE)')
-st.subheader('Wengong Jin, Regina Barzilay, Tommi Jaakkola')
-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)'''
-
-with st.expander('About'):
-    st.markdown(descrip)
-
-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')
+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;
+}
+'''
+
+st.markdown(f'<style>{css}</style>',unsafe_allow_html=True)
+
+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
 
 def penalized_logp_standard(mol):
 
@@ -61,51 +80,97 @@ def penalized_logp_standard(mol):
     standardized_cycle = (cycle_score - cycle_mean) / cycle_std
     return standardized_log_p + standardized_SA + standardized_cycle
 
-mol = Chem.MolFromSmiles(st.session_state.smiles)
-if mol is None:
-    st.write('SMILES is invalid. Please enter a valid SMILES.')
-else:
-    st.write('Molecule:')
-    st.image(MolToImage(mol,size=(300,300)))
-    score = penalized_logp_standard(mol)
-    st.write('Penalized logP score: %.5f' % (score))
-
-if mol is not None:
-    st.slider('Choose learning rate: ',0.0,10.0,0.4,key='lr')
-    st.slider('Choose similarity cutoff: ',0.0,3.0,0.4,key='sim_cutoff')
-    st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
-    vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
-    vocab = Vocab(vocab)
-    if st.button('Optimize'):
-        st.write('Testing')
-
-        model = JTPropVAE(vocab, 450, 56, 20, 3)
-
-        model.load_state_dict(torch.load('./model.iter-685000',map_location=torch.device('cpu')))
+lg = rdkit.RDLogger.logger() 
+lg.setLevel(rdkit.RDLogger.CRITICAL)
+
+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('<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 batch', 'About'], 
+                         iconName=['science', 'batch_prediction', 'info'], default_choice=0)
+
+if tabs == '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)
+elif tabs == 'Optimize a molecule':
+    st.markdown("<h2 style='text-align: center;'>Optimize a molecule</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')
+
+    mol = Chem.MolFromSmiles(st.session_state.smiles)
+    if mol is None:
+        st.markdown("<p style='text-align: center; color: red;'>SMILES is invalid. Please enter a valid SMILES.</p>",unsafe_allow_html=True)
+    else:
+        score = penalized_logp_standard(mol)
+        # 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)))
+        st.metric('Penalized logP score', '%.5f' % (score))
+
+    if mol is not None:
+        # col1, col2, col3 = st.columns(3)
+        st.slider('Choose learning rate: ',0.0,10.0,0.4,key='lr')
+        st.slider('Choose similarity cutoff: ',0.0,3.0,0.4,key='sim_cutoff')
+        st.slider('Choose number of iterations: ',1,100,80,key='n_iter')
+        vocab = [x.strip("\r\n ") for x in open('./vocab.txt')] 
+        vocab = Vocab(vocab)
+        if st.button('Optimize'):
+            # st.write('Testing')
+            
+            # 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.iter-685000',map_location=torch.device('cpu')))
         
-        new_smiles,sim = model.optimize(st.session_state.smiles, sim_cutoff=st.session_state.sim_cutoff, lr=st.session_state.lr, num_iter=st.session_state.n_iter)
+                    new_smiles,sim = model.optimize(st.session_state.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.write('Cannot optimize.')
-        else:
-            st.write('New SMILES:')
-            st.code(new_smiles)
-            new_mol = Chem.MolFromSmiles(new_smiles)
-            if new_mol is None:
-                st.write('New SMILES is invalid.')
-            else:
-                st.write('New SMILES molecule:')
-                st.image(MolToImage(new_mol,size=(300,300)))
-                new_score = penalized_logp_standard(new_mol)
-                st.write('New penalized logP score: %.5f' % (new_score))
-                st.write('Caching ZINC20 if necessary...')
-                if buy(new_smiles, catalog='zinc20',canonicalize=True):
-                    st.write('This molecule exists.')
-                    st.caption('Checked by molbloom.')
+                    del model
+                if new_smiles is None:
+                    st.markdown("<p style='text-align: center; color: red;'>Cannot optimize! Please choose another setting.</p>",unsafe_allow_html=True)
                 else:
-                    st.write('THIS MOLECULE DOES NOT EXIST!')
-                    st.caption('Checked by molbloom.')
-        
-    
-    
+                    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_score = penalized_logp_standard(new_mol)
+                    # st.write('New penalized logP score: %.5f' % (new_score))
+                        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: cyan;'><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: lightgreen;'>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)
+elif tabs == 'Optimize batch':
+    st.write('Incoming...')
+
+

fast_jtnn/mol_tree.py

@@ -2,7 +2,7 @@ import rdkit
 import rdkit.Chem as Chem
 from chemutils import get_clique_mol, tree_decomp, get_mol, get_smiles, set_atommap, enum_assemble, decode_stereo
 from vocab import *
-import argparse
+# import argparse
 
 class MolTreeNode(object):
 
@@ -124,45 +124,45 @@ def data_process_chunk(smiles_list):
         #     print(i, ' / 1584663')
     return list(cset)
 
-if __name__ == "__main__":
-    import sys
-    lg = rdkit.RDLogger.logger() 
-    lg.setLevel(rdkit.RDLogger.CRITICAL)
+# if __name__ == "__main__":
+#     import sys
+#     lg = rdkit.RDLogger.logger() 
+#     lg.setLevel(rdkit.RDLogger.CRITICAL)
     
-    i = 0
+#     i = 0
     
-    import os
-    from joblib import Parallel,delayed
-    from tqdm import tqdm
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--smiles_path', type=str,required=True)
-    parser.add_argument('--vocab_path', type=str,required=True)
-    parser.add_argument('--prop', type=bool,default=False)
-    parser.add_argument('--ncpu', default=8,type=int)
-    args = parser.parse_args()
-
-    if args.prop:
-        import pandas as pd
-        smiles_list = pd.read_csv(args.smiles_path,usecols=['SMILES'])
-        smiles_list = list(smiles_list.SMILES)
-    else:
-        with open(args.smiles_path,'r') as f:
-            smiles_list = [line.split()[0] for line in f]
-    print('Total smiles = ',len(smiles_list))
-
-    # moses: 1584663
+#     import os
+#     from joblib import Parallel,delayed
+#     from tqdm import tqdm
+#     parser = argparse.ArgumentParser()
+#     parser.add_argument('--smiles_path', type=str,required=True)
+#     parser.add_argument('--vocab_path', type=str,required=True)
+#     parser.add_argument('--prop', type=bool,default=False)
+#     parser.add_argument('--ncpu', default=8,type=int)
+#     args = parser.parse_args()
+
+#     if args.prop:
+#         import pandas as pd
+#         smiles_list = pd.read_csv(args.smiles_path,usecols=['SMILES'])
+#         smiles_list = list(smiles_list.SMILES)
+#     else:
+#         with open(args.smiles_path,'r') as f:
+#             smiles_list = [line.split()[0] for line in f]
+#     print('Total smiles = ',len(smiles_list))
+
+#     # moses: 1584663
     
-    chunk_size = 10000
-    vocab_set_list = Parallel(n_jobs=args.ncpu)(
-        delayed(data_process_chunk)(smiles_list[start: start + chunk_size])
-        for start in tqdm(range(0, len(smiles_list), chunk_size))
-    )
-    vocab_list =[]
-    for set_list in vocab_set_list:
-        vocab_list.extend(set_list)
-
-    cset = set(vocab_list)
-    with open(args.vocab_path,'w') as f:
-        for x in cset:
-            f.write(''.join([x,'\n']))
+#     chunk_size = 10000
+#     vocab_set_list = Parallel(n_jobs=args.ncpu)(
+#         delayed(data_process_chunk)(smiles_list[start: start + chunk_size])
+#         for start in tqdm(range(0, len(smiles_list), chunk_size))
+#     )
+#     vocab_list =[]
+#     for set_list in vocab_set_list:
+#         vocab_list.extend(set_list)
+
+#     cset = set(vocab_list)
+#     with open(args.vocab_path,'w') as f:
+#         for x in cset:
+#             f.write(''.join([x,'\n']))
 

requirements.txt

@@ -1,8 +1,9 @@
 rdkit
 numpy
 torch
-argparse
 tqdm
 networkx
 scipy
-molbloom
+molbloom
+st_on_hover_tabs
+stqdm

style.css

@@ -0,0 +1,58 @@
+section[data-testid='stSidebar'] {
+    background-color: #111;
+    min-width:unset !important;
+    width: unset !important;
+    flex-shrink: unset !important;
+    
+}
+
+button[kind="header"] {
+    background-color: transparent;
+    color:rgb(180, 167, 141)
+}
+
+@media(hover){
+    /* header element to be removed */
+    header[data-testid="stHeader"] {
+        display:none;
+    }
+
+    /* The navigation menu specs and size */
+    section[data-testid='stSidebar'] > div {
+        height: 100%;
+        width: 95px;
+        position: relative;
+        z-index: 1;
+        top: 0;
+        left: 0;
+        background-color: #111;
+        overflow-x: hidden;
+        transition: 0.5s ease;
+        padding-top: 60px;
+        white-space: nowrap;
+    }
+
+    /* The navigation menu open and close on hover and size */
+    /* section[data-testid='stSidebar'] > div {
+    height: 100%;
+    width: 75px; /* Put some width to hover on. */
+    /* } 
+
+    /* ON HOVER */
+    section[data-testid='stSidebar'] > div:hover{
+    width: 300px;
+    }
+
+    /* The button on the streamlit navigation menu - hidden */
+    /* button[kind="header"] {
+        display: none;
+    } */
+}
+     
+@media(max-width: 272px){
+    
+    section[data-testid='stSidebar'] > div {
+        width:15rem;
+    }
+}
+