Spaces:
Sleeping
Sleeping
| import torch | |
| import streamlit as st | |
| 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 | |
| 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 | |
| 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') | |
| 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 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'))) | |
| 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.') | |
| else: | |
| st.write('THIS MOLECULE DOES NOT EXIST!') | |
| st.caption('Checked by molbloom.') | |