Update app.py

app.py

@@ -68,10 +68,10 @@ model_configs = {
 
 
 
-def function(model_name: str, num_molecules: int, seed_num: int) -> tuple[PIL.Image, pd.DataFrame, str]:
+def function(model_name: str, num_molecules: int, seed_num: int):
     '''
     Returns:
-    image, score_df, file path
+    image, score_df, file_path, and individual metrics
     '''
     if model_name == "DrugGEN-NoTarget":
         model_name = "NoTarget"
@@ -112,6 +112,20 @@ def function(model_name: str, num_molecules: int, seed_num: int) -> tuple[PIL.Im
         "SA Score": [scores["sa"].iloc[0]]
     })
 
+    # Extract individual metrics
+    validity = scores["validity"].iloc[0]
+    uniqueness = scores["uniqueness"].iloc[0]
+    novelty_train = scores["novelty"].iloc[0]
+    novelty_test = scores["novelty_test"].iloc[0]
+    drug_novelty = scores["drug_novelty"].iloc[0]
+    runtime = et
+    qed = scores["qed"].iloc[0]
+    sa = scores["sa"].iloc[0]
+    int_div = scores["IntDiv"].iloc[0]
+    snn_chembl = scores["snn_chembl"].iloc[0]
+    snn_drug = scores["snn_drug"].iloc[0]
+    max_len = scores["max_len"].iloc[0]
+
     output_file_path = f'experiments/inference/{model_name}/inference_drugs.txt'
 
     new_path = f'{model_name}_denovo_mols.smi'
@@ -146,7 +160,7 @@ def function(model_name: str, num_molecules: int, seed_num: int) -> tuple[PIL.Im
         highlightBondLists=None,
     )
 
-    return molecule_image, score_df, new_path
+    return molecule_image, score_df, new_path, validity, uniqueness, novelty_train, novelty_test, drug_novelty, runtime, qed, sa, int_div, snn_chembl, snn_drug, max_len
 
 
 
@@ -154,9 +168,21 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("# DrugGEN: Target Centric De Novo Design of Drug Candidate Molecules with Graph Generative Deep Adversarial Networks")
-            with gr.Row():
-                gr.Markdown("[![arXiv](https://img.shields.io/badge/arXiv-2302.07868-b31b1b.svg)](https://arxiv.org/abs/2302.07868)")
-                gr.Markdown("[![github-repository](https://img.shields.io/badge/GitHub-black?logo=github)](https://github.com/HUBioDataLab/DrugGEN)")
+            
+            gr.HTML("""
+            <div style="display: flex; gap: 10px; margin-bottom: 15px;">
+                <a href="https://arxiv.org/abs/2302.07868" target="_blank" style="text-decoration: none;">
+                    <div style="display: inline-block; background-color: #b31b1b; color: white; padding: 5px 10px; border-radius: 5px; font-size: 14px;">
+                        <span style="font-weight: bold;">arXiv</span> 2302.07868
+                    </div>
+                </a>
+                <a href="https://github.com/HUBioDataLab/DrugGEN" target="_blank" style="text-decoration: none;">
+                    <div style="display: inline-block; background-color: #24292e; color: white; padding: 5px 10px; border-radius: 5px; font-size: 14px;">
+                        <span style="font-weight: bold;">GitHub</span> Repository
+                    </div>
+                </a>
+            </div>
+            """)
             
             with gr.Accordion("About DrugGEN Models", open=False):
                 gr.Markdown("""
@@ -234,29 +260,67 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
             )
 
         with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.TabItem("Generated Molecules"):
-                    image_output = gr.Image(
-                        label="Sample of Generated Molecules",
-                        elem_id="molecule_display"
-                    )
-                    file_download = gr.File(
-                        label="Download All Generated Molecules (SMILES format)",
-                    )
+            image_output = gr.Image(
+                label="Sample of Generated Molecules",
+                elem_id="molecule_display"
+            )
+            
+            file_download = gr.File(
+                label="Download All Generated Molecules (SMILES format)",
+            )
+            
+            with gr.Box():
+                gr.Markdown("### Performance Metrics")
                 
-                with gr.TabItem("Performance Metrics"):
-                    scores_df = gr.Dataframe(
-                        label="Model Performance Metrics",
-                        headers=["Runtime (seconds)", "Validity", "Uniqueness", "Novelty (Train)", "Novelty (Test)", 
-                                "Drug Novelty", "Max Length", "Mean Atom Type", "SNN ChEMBL", "SNN Drug", 
-                                "Internal Diversity", "QED", "SA Score"]
-                    )
-
-
-    gr.Markdown("### Created by the HU BioDataLab | [GitHub](https://github.com/HUBioDataLab/DrugGEN) | [Paper](https://arxiv.org/abs/2302.07868)")
+                with gr.Row():
+                    with gr.Column():
+                        validity = gr.Number(label="Validity", precision=3)
+                        uniqueness = gr.Number(label="Uniqueness", precision=3)
+                        novelty_train = gr.Number(label="Novelty (Train)", precision=3)
+                        novelty_test = gr.Number(label="Novelty (Test)", precision=3)
+                        drug_novelty = gr.Number(label="Drug Novelty", precision=3)
+                        runtime = gr.Number(label="Runtime (seconds)", precision=2)
+                    
+                    with gr.Column():
+                        qed = gr.Number(label="QED Score", precision=3, info="Higher is more drug-like (0-1)")
+                        sa = gr.Number(label="SA Score", precision=3, info="Lower is easier to synthesize (1-10)")
+                        int_div = gr.Number(label="Internal Diversity", precision=3)
+                        snn_chembl = gr.Number(label="SNN ChEMBL", precision=3)
+                        snn_drug = gr.Number(label="SNN Drug", precision=3)
+                        max_len = gr.Number(label="Max Length", precision=3)
+            
+            with gr.Accordion("All Metrics (Table View)", open=False):
+                scores_df = gr.Dataframe(
+                    headers=["Runtime (seconds)", "Validity", "Uniqueness", "Novelty (Train)", "Novelty (Test)", 
+                            "Drug Novelty", "Max Length", "Mean Atom Type", "SNN ChEMBL", "SNN Drug", 
+                            "Internal Diversity", "QED", "SA Score"]
+                )
+            
+    gr.Markdown("### Created by the HUBioDataLab | [GitHub](https://github.com/HUBioDataLab/DrugGEN) | [Paper](https://arxiv.org/abs/2302.07868)")
+
+    submit_button.click(
+        function, 
+        inputs=[model_name, num_molecules, seed_num], 
+        outputs=[
+            image_output, 
+            scores_df, 
+            file_download,
+            validity,
+            uniqueness,
+            novelty_train,
+            novelty_test,
+            drug_novelty,
+            runtime,
+            qed,
+            sa,
+            int_div,
+            snn_chembl,
+            snn_drug,
+            max_len
+        ], 
+        api_name="inference"
+    )
 
-    submit_button.click(function, inputs=[model_name, num_molecules, seed_num], outputs=[image_output, scores_df, file_download], api_name="inference")
-#demo.queue(concurrency_count=1)
 demo.queue()
 demo.launch()