Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-#from src.utils import plogp, sf_decode, sim
 import pandas as pd
 from rdkit import Chem
 from rdkit.Chem import AllChem
@@ -59,12 +58,57 @@ def sim(input_smile, output_smile):
     else: return None 
 
 
-def greet(name):
+def gen_opt(gen_input):
+    tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen")
+    model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen")
+    
+    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 = name
+    input = opt_input
+
+    smis_input = [sf.decoder(i) for i in input]
+    mols_input = []
+    for smi in smis_input:
+        mol = Chem.MolFromSmiles(smi)
+        mols_input.append(mol)
+    
+    opt_input_img = Draw.MolsToGridImage(
+        mols_input,
+        molsPerRow=4,
+        subImgSize=(200,200),
+        legends=['' for x in mols]
+    )
     
     sf_input = tokenizer(input, return_tensors="pt")
     molecules = model.generate(
@@ -83,7 +127,6 @@ def greet(name):
     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]
     
@@ -93,20 +136,86 @@ def greet(name):
     candidate_selfies = {"candidates": sf_output, "improvement": plogp_improve, "sim": simm}
     data = pd.DataFrame(candidate_selfies)
     
-    return data[(data['improvement']> 0) & (data['sim']>0.4)]
+    results = data[(data['improvement']> 0) & (data['sim']>0.4)]
+    opt_output  = results["candidates"].tolist()
+    opt_output_imp = results["improvement"].tolist()
+    opt_output_sim = results["sim"].tolist()
 
-  
+    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
+# 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]']
 
+# ]
 
-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]']
 
-]
 
+# iface = gr.Interface(fn=greet, inputs="text", outputs="numpy", title="Molecular Language Model as Multi-task Generator",examples=examples)
+# iface.launch()
+
+with gr.Blocks() as demo:
+    init_triple_input()
+    gr.Markdown("# MolGen: Molecular Language Model as Multi-task Generator")
+
+    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.Textbox(label="Visualization", lines=3, placeholder="")
+       
+            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.Textbox(label="Input Visualization", lines=3, placeholder="")
+                    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.Textbox(label="Output Visualization", lines=3, placeholder="")
+
+            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])
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="numpy", title="Molecular Language Model as Multi-task Generator",examples=examples)
-iface.launch()
+demo.launch()