Update app.py

app.py

@@ -112,19 +112,25 @@ def function(model_name: str, num_molecules: int, seed_num: int):
         "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]
+    # Create basic metrics dataframe
+    basic_metrics = pd.DataFrame({
+        "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 (s)": [round(et, 2)]
+    })
+    
+    # Create advanced metrics dataframe
+    advanced_metrics = pd.DataFrame({
+        "QED": [scores["qed"].iloc[0]],
+        "SA Score": [scores["sa"].iloc[0]],
+        "Internal Diversity": [scores["IntDiv"].iloc[0]],
+        "SNN ChEMBL": [scores["snn_chembl"].iloc[0]],
+        "SNN Drug": [scores["snn_drug"].iloc[0]],
+        "Max Length": [scores["max_len"].iloc[0]]
+    })
 
     output_file_path = f'experiments/inference/{model_name}/inference_drugs.txt'
 
@@ -160,7 +166,7 @@ def function(model_name: str, num_molecules: int, seed_num: int):
         highlightBondLists=None,
     )
 
-    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
+    return molecule_image, score_df, new_path, basic_metrics, advanced_metrics
 
 
 
@@ -200,6 +206,8 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
             
             with gr.Accordion("About DrugGEN Models", open=False):
                 gr.Markdown("""
+## Model Variations
+
 ### DrugGEN-AKT1
 This model is designed to generate molecules targeting the human AKT1 protein (UniProt ID: P31749).
 
@@ -227,7 +235,7 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
 ### Novelty Metrics
 - **Novelty (Train)**: Percentage of molecules not found in the training set
 - **Novelty (Test)**: Percentage of molecules not found in the test set
-- **Drug Novelty**: Percentage of molecules not found in known drugs
+- **Drug Novelty**: Percentage of molecules not found in known inhibitors of the target protein
 
 ### Structural Metrics
 - **Max Length**: Maximum component length in the generated molecules
@@ -272,6 +280,19 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
             )
 
         with gr.Column(scale=2):
+            with gr.Row():
+                with gr.Column():
+                    basic_metrics_df = gr.Dataframe(
+                        headers=["Validity", "Uniqueness", "Novelty (Train)", "Novelty (Test)", "Drug Novelty", "Runtime (s)"],
+                        elem_id="basic-metrics"
+                    )
+                
+                with gr.Column():
+                    advanced_metrics_df = gr.Dataframe(
+                        headers=["QED", "SA Score", "Internal Diversity", "SNN ChEMBL", "SNN Drug", "Max Length"],
+                        elem_id="advanced-metrics"
+                    )
+            
             image_output = gr.Image(
                 label="Sample of Generated Molecules",
                 elem_id="molecule_display"
@@ -300,14 +321,7 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
                         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(
@@ -332,5 +346,7 @@ For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
         ], 
         api_name="inference"
     )
+
 demo.queue()
-demo.launch()
+demo.launch()
+