import gradio as gr from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import pandas as pd from rdkit import Chem from rdkit.Chem import AllChem from rdkit import DataStructs from rdkit.Chem import Descriptors from rdkit.Chem import Draw import selfies as sf from rdkit.Chem import RDConfig import os import sys sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score')) import sascorer def get_largest_ring_size(mol): cycle_list = mol.GetRingInfo().AtomRings() if cycle_list: cycle_length = max([len(j) for j in cycle_list]) else: cycle_length = 0 return cycle_length def plogp(smile): if smile: mol = Chem.MolFromSmiles(smile) if mol: log_p = Descriptors.MolLogP(mol) sas_score = sascorer.calculateScore(mol) largest_ring_size = get_largest_ring_size(mol) cycle_score = max(largest_ring_size - 6, 0) if log_p and sas_score and largest_ring_size: p_logp = log_p - sas_score - cycle_score return p_logp else: return -100 else: return -100 else: return -100 def sf_decode(selfies): try: decode = sf.decoder(selfies) return decode except sf.DecoderError: return '' def sim(input_smile, output_smile): if input_smile and output_smile: input_mol = Chem.MolFromSmiles(input_smile) output_mol = Chem.MolFromSmiles(output_smile) if input_mol and output_mol: input_fp = AllChem.GetMorganFingerprint(input_mol, 2) output_fp = AllChem.GetMorganFingerprint(output_mol, 2) sim = DataStructs.TanimotoSimilarity(input_fp, output_fp) return sim else: return None else: return None def gen_process(gen_input): tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen-large") model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen-large") sf_input = tokenizer(gen_input, return_tensors="pt") # beam search molecules = model.generate(input_ids=sf_input["input_ids"], attention_mask=sf_input["attention_mask"], max_length=15, min_length=5, num_return_sequences=4, num_beams=5) gen_output = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True).replace(" ","") for g in molecules] smis = [sf.decoder(i) for i in gen_output] mols = [] for smi in smis: mol = Chem.MolFromSmiles(smi) mols.append(mol) gen_output_image = Draw.MolsToGridImage( mols, molsPerRow=4, subImgSize=(200,200), legends=['' for x in mols] ) return "\n".join(gen_output), gen_output_image def opt_process(opt_input): tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen-large-opt") model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen-large-opt") input = opt_input smis_input = sf.decoder(input) mol_input = [] mol = Chem.MolFromSmiles(smis_input) mol_input.append(mol) opt_input_img = Draw.MolsToGridImage( mol_input, molsPerRow=4, subImgSize=(200,200), legends=['' for x in mol_input] ) sf_input = tokenizer(input, return_tensors="pt") molecules = model.generate( input_ids=sf_input["input_ids"], attention_mask=sf_input["attention_mask"], do_sample=True, max_length=100, min_length=5, top_k=30, top_p=1, num_return_sequences=10 ) sf_output = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True).replace(" ","") for g in molecules] sf_output = list(set(sf_output)) input_sm = sf_decode(input) sm_output = [sf_decode(sf) for sf in sf_output] input_plogp = plogp(input_sm) plogp_improve = [plogp(i)-input_plogp for i in sm_output] simm = [sim(i,input_sm) for i in sm_output] candidate_selfies = {"candidates": sf_output, "improvement": plogp_improve, "sim": simm} data = pd.DataFrame(candidate_selfies) results = data[(data['improvement']> 0) & (data['sim']>0.4)] opt_output = results["candidates"].tolist() opt_output_imp = results["improvement"].tolist() opt_output_imp = [str(i) for i in opt_output_imp] opt_output_sim = results["sim"].tolist() opt_output_sim = [str(i) for i in opt_output_sim] smis = [sf.decoder(i) for i in opt_output] mols = [] for smi in smis: mol = Chem.MolFromSmiles(smi) mols.append(mol) opt_output_img = Draw.MolsToGridImage( mols, molsPerRow=4, subImgSize=(200,200), legends=['' for x in mols] ) return opt_input_img, "\n".join(opt_output), "\n".join(opt_output_imp), "\n".join(opt_output_sim), opt_output_img with gr.Blocks() as demo: gr.Markdown("# MolGen: Domain-Agnostic Molecular Generation with Self-feedback") with gr.Tabs(): with gr.TabItem("Molecular Generation"): with gr.Row(): with gr.Column(): gen_input = gr.Textbox(label="Input", lines=1, placeholder="SELFIES Input") gen_button = gr.Button("Generate") with gr.Column(): gen_output = gr.Textbox(label="Generation Results", lines=5, placeholder="") gen_output_image = gr.Image(label="Visualization") gr.Examples( examples=[["[C][=C][C][=C][C][=C][Ring1][=Branch1]"], ["[C]"] ], inputs=[gen_input], outputs=[gen_output, gen_output_image], fn=gen_process, cache_examples=True, ) with gr.TabItem("Constrained Molecular Property Optimization"): with gr.Row(): with gr.Column(): opt_input = gr.Textbox(label="Input", lines=1, placeholder="SELFIES Input") opt_button = gr.Button("Optimize") with gr.Column(): opt_input_img = gr.Image(label="Input Visualization") opt_output = gr.Textbox(label="Optimization Results", lines=3, placeholder="") opt_output_imp = gr.Textbox(label="Optimization Property Improvements", lines=3, placeholder="") opt_output_sim = gr.Textbox(label="Similarity", lines=3, placeholder="") opt_output_img = gr.Image(label="Output Visualization") gr.Examples( examples=[["[C][C][=Branch1][C][=O][N][C][C][O][C][C][O][C][C][O][C][C][Ring1][N]"], ["[C][C][S][C][C][S][C][C][C][S][C][C][S][C][Ring1][=C]"], ["[N][#C][C][C][C@@H1][C][C][C][C][C][C][C][C][C][C][C][Ring1][N][=O]"] ], inputs=[opt_input], outputs=[opt_input_img, opt_output, opt_output_imp, opt_output_sim, opt_output_img], fn=opt_process, cache_examples=True, ) gen_button.click(fn=gen_process, inputs=[gen_input], outputs=[gen_output, gen_output_image]) opt_button.click(fn=opt_process, inputs=[opt_input], outputs=[opt_input_img, opt_output, opt_output_imp, opt_output_sim, opt_output_img]) demo.launch()