Update app.py

app.py

@@ -13,50 +13,50 @@ import time
 
 class DrugGENConfig:
     # Inference configuration
-    submodel='DrugGEN'
-    inference_model="/home/user/app/experiments/models/DrugGEN/"
-    sample_num=100
+    submodel = 'DrugGEN'
+    inference_model = "/home/user/app/experiments/models/DrugGEN/"
+    sample_num = 100
 
     # Data configuration
-    inf_smiles='/home/user/app/data/chembl_test.smi'
-    train_smiles='/home/user/app/data/chembl_train.smi'
-    inf_batch_size=1
-    mol_data_dir='/home/user/app/data'
-    features=False
+    inf_smiles = '/home/user/app/data/chembl_test.smi'
+    train_smiles = '/home/user/app/data/chembl_train.smi'
+    inf_batch_size = 1
+    mol_data_dir = '/home/user/app/data'
+    features = False
 
     # Model configuration
-    act='relu'
-    max_atom=45
-    dim=128
-    depth=1
-    heads=8
-    mlp_ratio=3
-    dropout=0.
+    act = 'relu'
+    max_atom = 45
+    dim = 128
+    depth = 1
+    heads = 8
+    mlp_ratio = 3
+    dropout = 0.
 
     # Seed configuration
-    set_seed=True
-    seed=10
+    set_seed = True
+    seed = 10
 
-    disable_correction=False
+    disable_correction = False
 
 
 class DrugGENAKT1Config(DrugGENConfig):
-    submodel='DrugGEN'
-    inference_model="/home/user/app/experiments/models/DrugGEN-akt1/"
-    train_drug_smiles='/home/user/app/data/akt_train.smi'
-    max_atom=45
+    submodel = 'DrugGEN'
+    inference_model = "/home/user/app/experiments/models/DrugGEN-akt1/"
+    train_drug_smiles = '/home/user/app/data/akt_train.smi'
+    max_atom = 45
 
 
 class DrugGENCDK2Config(DrugGENConfig):
-    submodel='DrugGEN'
-    inference_model="/home/user/app/experiments/models/DrugGEN-cdk2/"
-    train_drug_smiles='/home/user/app//data/cdk2_train.smi'
-    max_atom=38
+    submodel = 'DrugGEN'
+    inference_model = "/home/user/app/experiments/models/DrugGEN-cdk2/"
+    train_drug_smiles = '/home/user/app/data/cdk2_train.smi'
+    max_atom = 38
 
 
 class NoTargetConfig(DrugGENConfig):
-    submodel="NoTarget"
-    inference_model="/home/user/app/experiments/models/NoTarget/"
+    submodel = "NoTarget"
+    inference_model = "/home/user/app/experiments/models/NoTarget/"
 
 
 model_configs = {
@@ -66,24 +66,34 @@ model_configs = {
 }
 
 
-
-def function(model_name: str, input_mode: str, num_molecules: int = None, seed_num: str = None, smiles_input: str = None):
-    '''
-    Returns:
-    image, metrics_df, file_path, basic_metrics, advanced_metrics
-    '''
-    if model_name == "DrugGEN-NoTarget":
-        model_name = "NoTarget"
+def run_inference(mode: str, model_name: str, num_molecules: int, seed_num: str, custom_smiles: str):
+    """
+    Depending on the selected mode, either generate new molecules or evaluate provided SMILES.
     
+    Returns:
+        image, file_path, basic_metrics, advanced_metrics
+    """
     config = model_configs[model_name]
-    
-    # Handle the input mode
-    if input_mode == "generate":
-        config.sample_num = num_molecules
-
-        if config.sample_num > 250:
-            raise gr.Error("You have requested to generate more than the allowed limit of 250 molecules. Please reduce your request to 250 or fewer.")
 
+    if mode == "Custom Input SMILES":
+        # Process the custom input SMILES
+        smiles_list = [s.strip() for s in custom_smiles.strip().splitlines() if s.strip() != ""]
+        if len(smiles_list) > 100:
+            raise gr.Error("You have provided more than the allowed limit of 100 molecules. Please provide 100 or fewer.")
+        # Write the custom SMILES to a temporary file and update config
+        temp_input_file = "custom_input.smi"
+        with open(temp_input_file, "w") as f:
+            for s in smiles_list:
+                f.write(s + "\n")
+        config.inf_smiles = temp_input_file
+        config.sample_num = len(smiles_list)
+        # Always use a random seed for custom mode
+        config.seed = random.randint(0, 10000)
+    else:
+        # Classical Generation mode
+        config.sample_num = num_molecules
+        if config.sample_num > 200:
+            raise gr.Error("You have requested to generate more than the allowed limit of 200 molecules. Please reduce your request to 200 or fewer.")
         if seed_num is None or seed_num.strip() == "":
             config.seed = random.randint(0, 10000)
         else:
@@ -91,70 +101,25 @@ def function(model_name: str, input_mode: str, num_molecules: int = None, seed_n
                 config.seed = int(seed_num)
             except ValueError:
                 raise gr.Error("The seed must be an integer value!")
-    else:  # input_mode == "smiles"
-        if not smiles_input or smiles_input.strip() == "":
-            raise gr.Error("Please enter at least one SMILES string.")
-            
-        # Split by newlines and filter empty lines
-        smiles_list = [s.strip() for s in smiles_input.strip().split('\n') if s.strip()]
-        
-        if len(smiles_list) > 100:
-            raise gr.Error("You have entered more than the allowed limit of 100 SMILES. Please reduce your input.")
-        
-        # Validate all SMILES
-        invalid_smiles = []
-        for i, smi in enumerate(smiles_list):
-            mol = Chem.MolFromSmiles(smi)
-            if mol is None:
-                invalid_smiles.append((i+1, smi))
-        
-        if invalid_smiles:
-            invalid_str = "\n".join([f"Line {i}: {smi}" for i, smi in invalid_smiles])
-            raise gr.Error(f"The following SMILES are invalid:\n{invalid_str}")
-        
-        # Save SMILES to a temporary file that matches the expected input format
-        temp_smiles_file = f'/home/user/app/data/temp_input.smi'
-        with open(temp_smiles_file, 'w') as f:
-            f.write('\n'.join(smiles_list))
-        
-        # Update config to use this file
-        config.inf_smiles = temp_smiles_file
-        config.sample_num = len(smiles_list)
-        
-        # Always use a fixed seed for SMILES mode
-        config.seed = 42
 
-    if model_name != "NoTarget":
-        model_name = "DrugGEN"
+    # Adjust model name for the inference if not using NoTarget
+    if model_name != "DrugGEN-NoTarget":
+        target_model_name = "DrugGEN"
+    else:
+        target_model_name = "NoTarget"
 
     inferer = Inference(config)
     start_time = time.time()
     scores = inferer.inference()  # This returns a DataFrame with specific columns
     et = time.time() - start_time
 
-    score_df = pd.DataFrame({
-        "Runtime (seconds)": [et],
-        "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]],
-        "Max Length": [scores["max_len"].iloc[0]],
-        "Mean Atom Type": [scores["mean_atom_type"].iloc[0]],
-        "SNN ChEMBL": [scores["snn_chembl"].iloc[0]],
-        "SNN Drug": [scores["snn_drug"].iloc[0]],
-        "Internal Diversity": [scores["IntDiv"].iloc[0]],
-        "QED": [scores["qed"].iloc[0]],
-        "SA Score": [scores["sa"].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]],
+        "Novelty (Inference)": [scores["novelty_test"].iloc[0]],
+        "Novelty (Real Inhibitors)": [scores["drug_novelty"].iloc[0]],
         "Runtime (s)": [round(et, 2)]
     })
     
@@ -164,13 +129,13 @@ def function(model_name: str, input_mode: str, num_molecules: int = None, seed_n
         "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]]
+        "SNN Real Inhibitors": [scores["snn_drug"].iloc[0]],
+        "Average Length": [scores["max_len"].iloc[0]]
     })
 
-    output_file_path = f'/home/user/app/experiments/inference/{model_name}/inference_drugs.txt'
-
-    new_path = f'{model_name}_denovo_mols.smi'
+    # Process the output file from inference
+    output_file_path = f'/home/user/app/experiments/inference/{target_model_name}/inference_drugs.txt'
+    new_path = f'{target_model_name}_denovo_mols.smi'
     os.rename(output_file_path, new_path)
 
     with open(new_path) as f:
@@ -178,13 +143,14 @@ def function(model_name: str, input_mode: str, num_molecules: int = None, seed_n
 
     generated_molecule_list = inference_drugs.split("\n")[:-1]
 
+    # Randomly select up to 12 molecules for display
     rng = random.Random(config.seed)
     if len(generated_molecule_list) > 12:
-        selected_molecules = rng.choices(generated_molecule_list, k=12)
+        selected_smiles = rng.choices(generated_molecule_list, k=12)
     else:
-        selected_molecules = generated_molecule_list
-    
-    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules if Chem.MolFromSmiles(mol) is not None]
+        selected_smiles = generated_molecule_list
+
+    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_smiles if Chem.MolFromSmiles(mol) is not None]
 
     drawOptions = Draw.rdMolDraw2D.MolDrawOptions()
     drawOptions.prepareMolsBeforeDrawing = False
@@ -195,21 +161,15 @@ def function(model_name: str, input_mode: str, num_molecules: int = None, seed_n
         molsPerRow=3,
         subImgSize=(400, 400),
         maxMols=len(selected_molecules),
-        # legends=None,
         returnPNG=False,
         drawOptions=drawOptions,
         highlightAtomLists=None,
         highlightBondLists=None,
     )
 
-    # Clean up the temporary file if it was created
-    if input_mode == "smiles" and os.path.exists(temp_smiles_file):
-        os.remove(temp_smiles_file)
-
     return molecule_image, new_path, basic_metrics, advanced_metrics
 
 
-
 with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     # Add custom CSS for styling
     gr.HTML("""
@@ -225,44 +185,40 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     </style>
     """)
 
-    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")
-            
-            gr.HTML("""
-            <div style="display: flex; gap: 10px; margin-bottom: 15px;">
-                <!-- arXiv badge -->
-                <a href="https://arxiv.org/abs/2302.07868" target="_blank" style="text-decoration: none;">
-                    <div style="
-                        display: inline-block; 
-                        background-color: #b31b1b; 
-                        color: #ffffff !important;  /* Force white text */
-                        padding: 5px 10px; 
-                        border-radius: 5px; 
-                        font-size: 14px;"
-                    >
-                        <span style="font-weight: bold;">arXiv</span> 2302.07868
-                    </div>
-                </a>
-                
-                <!-- GitHub badge -->
-                <a href="https://github.com/HUBioDataLab/DrugGEN" target="_blank" style="text-decoration: none;">
-                    <div style="
-                        display: inline-block; 
-                        background-color: #24292e; 
-                        color: #ffffff !important;  /* Force white text */
-                        padding: 5px 10px; 
-                        border-radius: 5px; 
-                        font-size: 14px;"
-                    >
-                        <span style="font-weight: bold;">GitHub</span> Repository
-                    </div>
-                </a>
+    gr.Markdown("# DrugGEN: Target Centric De Novo Design of Drug Candidate Molecules with Graph Generative Deep Adversarial Networks")
+    
+    gr.HTML("""
+    <div style="display: flex; gap: 10px; margin-bottom: 15px;">
+        <!-- arXiv badge -->
+        <a href="https://arxiv.org/abs/2302.07868" target="_blank" style="text-decoration: none;">
+            <div style="
+                display: inline-block; 
+                background-color: #b31b1b; 
+                color: #ffffff !important; 
+                padding: 5px 10px; 
+                border-radius: 5px; 
+                font-size: 14px;">
+                <span style="font-weight: bold;">arXiv</span> 2302.07868
+            </div>
+        </a>
+        
+        <!-- GitHub badge -->
+        <a href="https://github.com/HUBioDataLab/DrugGEN" target="_blank" style="text-decoration: none;">
+            <div style="
+                display: inline-block; 
+                background-color: #24292e; 
+                color: #ffffff !important; 
+                padding: 5px 10px; 
+                border-radius: 5px; 
+                font-size: 14px;">
+                <span style="font-weight: bold;">GitHub</span> Repository
             </div>
-            """)
-            
-            with gr.Accordion("About DrugGEN Models", open=False):
-                gr.Markdown("""
+        </a>
+    </div>
+    """)
+
+    with gr.Accordion("About DrugGEN Models", open=False):
+        gr.Markdown("""
 ## Model Variations
 
 ### DrugGEN-AKT1
@@ -272,256 +228,159 @@ This model is designed to generate molecules targeting the human AKT1 protein (U
 This model is designed to generate molecules targeting the human CDK2 protein (UniProt ID: P24941).
 
 ### DrugGEN-NoTarget
-This is a general-purpose model that generates diverse drug-like molecules without targeting a specific protein. It's useful for:
-- Exploring chemical space
-- Generating diverse scaffolds
-- Creating molecules with drug-like properties
+This is a general-purpose model that generates diverse drug-like molecules without targeting a specific protein.
+- Useful for exploring chemical space, generating diverse scaffolds, and creating molecules with drug-like properties.
 
 For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
-                """)
-            
-            with gr.Accordion("Understanding the Metrics", open=False):
-                gr.Markdown("""
+        """)
+    
+    with gr.Accordion("Understanding the Metrics", open=False):
+        gr.Markdown("""
 ## Evaluation Metrics
 
 ### Basic Metrics
 - **Validity**: Percentage of generated molecules that are chemically valid
 - **Uniqueness**: Percentage of unique molecules among valid ones
-- **Runtime**: Time taken to generate the requested molecules
+- **Runtime**: Time taken to generate or evaluate the molecules
 
 ### 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 inhibitors of the target protein
+- **Novelty (Inference)**: Percentage of molecules not found in the test set
+- **Novelty (Real Inhibitors)**: Percentage of molecules not found in known inhibitors of the target protein
 
 ### Structural Metrics
-- **Max Length**: Maximum component length in the generated molecules
+- **Average Length**: Average component length in the generated molecules
 - **Mean Atom Type**: Average distribution of atom types
 - **Internal Diversity**: Diversity within the generated set (higher is more diverse)
 
 ### Drug-likeness Metrics
 - **QED (Quantitative Estimate of Drug-likeness)**: Score from 0-1 measuring how drug-like a molecule is (higher is better)
-- **SA Score (Synthetic Accessibility)**: Score from 1-10 indicating ease of synthesis (lower is easier)
+- **SA Score (Synthetic Accessibility)**: Score from 1-10 indicating ease of synthesis (lower is better)
 
 ### Similarity Metrics
 - **SNN ChEMBL**: Similarity to ChEMBL molecules (higher means more similar to known drug-like compounds)
-- **SNN Drug**: Similarity to known drugs (higher means more similar to approved drugs)
-                """)
-            
-            model_name = gr.Radio(
-                choices=("DrugGEN-AKT1", "DrugGEN-CDK2", "DrugGEN-NoTarget"),
-                value="DrugGEN-AKT1",
-                label="Select Target Model",
-                info="Choose which protein target or general model to use for molecule generation"
-            )
-            
-            # Add a separator between model selection and input mode
-            gr.Markdown("---")
-            gr.Markdown("## Input Settings")
-            
-            # Replace radio with switch using a better layout
-            with gr.Row(equal_height=True):
-                with gr.Column(scale=1, min_width=150):
-                    gr.Markdown("### Classic Generation", elem_id="generate-mode-label")
-                
-                with gr.Column(scale=1, min_width=150):
-                    input_mode_switch = gr.Checkbox(
-                        value=False,
-                        label="Switch Input Mode",
-                        elem_id="input-mode-switch"
+- **SNN Real Inhibitors**: Similarity to known drugs (higher means more similar to approved drugs)
+        """)
+
+    # Use Gradio Tabs to separate the two modes.
+    with gr.Tabs():
+        with gr.TabItem("Classical Generation"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    model_name = gr.Radio(
+                        choices=("DrugGEN-AKT1", "DrugGEN-CDK2", "DrugGEN-NoTarget"),
+                        value="DrugGEN-AKT1",
+                        label="Select Target Model",
+                        info="Choose which protein target or general model to use for molecule generation"
+                    )
+
+                    num_molecules = gr.Slider(
+                        minimum=10,
+                        maximum=200,
+                        value=100,
+                        step=10,
+                        label="Number of Molecules to Generate",
+                        info="This space runs on a CPU, which may result in slower performance. Generating 100 molecules takes approximately 6 minutes. Therefore, we set a 200-molecule cap."
+                    )
+
+                    seed_num = gr.Textbox(
+                        label="Random Seed (Optional)",
+                        value="",
+                        info="Set a specific seed for reproducible results, or leave empty for random generation"
+                    )
+
+                    classical_submit = gr.Button(
+                        value="Generate Molecules",
+                        variant="primary",
+                        size="lg"
+                    )
+                with gr.Column(scale=2):
+                    basic_metrics_df = gr.Dataframe(
+                        headers=["Validity", "Uniqueness", "Novelty (Train)", "Novelty (Inference)", "Novelty (Real Inhibitors)", "Runtime (s)"],
+                        elem_id="basic-metrics"
+                    )
+                        
+                    advanced_metrics_df = gr.Dataframe(
+                        headers=["QED", "SA Score", "Internal Diversity", "SNN (ChEMBL)", "SNN (Real Inhibitors)", "Average Length"],
+                        elem_id="advanced-metrics"
+                    )
+
+                    file_download = gr.File(
+                        label="Download All Generated Molecules (SMILES format)"
                     )
                     
-                with gr.Column(scale=1, min_width=150):
-                    gr.Markdown("### Custom SMILES Input", elem_id="smiles-mode-label")
-            
-            # Add custom CSS and JavaScript for better styling
-            gr.HTML("""
-            <style>
-            #input-mode-switch {
-                margin: 20px auto;
-                display: flex;
-                justify-content: center;
-            }
-            
-            #generate-mode-label, #smiles-mode-label {
-                text-align: center;
-                margin-top: 10px;
-                font-weight: bold;
-                transition: opacity 0.3s ease;
-            }
-            
-            /* Make the inactive mode label more subtle */
-            #generate-mode-label {
-                opacity: 1;
-                color: #4CAF50;
-            }
-            
-            #smiles-mode-label {
-                opacity: 0.5;
-                color: #2196F3;
-            }
-            
-            .active-mode {
-                text-decoration: underline;
-                font-size: 1.1em;
-            }
-            
-            /* Style for the input boxes */
-            .input-box {
-                border: 2px solid rgba(128, 128, 228, 0.3);
-                border-radius: 10px;
-                padding: 15px;
-                margin-top: 15px;
-                background-color: rgba(32, 36, 45, 0.7);
-                box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-                transition: all 0.3s ease;
-            }
-            
-            .input-box:hover {
-                border-color: rgba(128, 128, 228, 0.6);
-                box-shadow: 0 6px 8px rgba(0, 0, 0, 0.15);
-            }
-            
-            /* Style the checkbox */
-            #input-mode-switch label {
-                font-weight: bold;
-                font-size: 1.1em;
-                color: rgba(128, 128, 228, 0.9);
-            }
-            
-            /* Add a hint to indicate the toggle functionality */
-            #input-mode-switch::after {
-                content: 'Click to toggle between modes';
-                display: block;
-                text-align: center;
-                font-size: 0.8em;
-                opacity: 0.7;
-                margin-top: 5px;
-            }
-            </style>
-            
-            <script>
-            // Add JavaScript to enhance the mode switching UI
-            document.addEventListener('DOMContentLoaded', function() {
-                // Get references to elements
-                const checkbox = document.querySelector('#input-mode-switch input[type="checkbox"]');
-                const generateLabel = document.querySelector('#generate-mode-label');
-                const smilesLabel = document.querySelector('#smiles-mode-label');
-                
-                // Add initial active class
-                generateLabel.classList.add('active-mode');
-                
-                // Add event listener to checkbox
-                if (checkbox) {
-                    checkbox.addEventListener('change', function() {
-                        if (this.checked) {
-                            // SMILES mode is active
-                            generateLabel.style.opacity = '0.5';
-                            smilesLabel.style.opacity = '1';
-                            generateLabel.classList.remove('active-mode');
-                            smilesLabel.classList.add('active-mode');
-                        } else {
-                            // Generate mode is active
-                            generateLabel.style.opacity = '1';
-                            smilesLabel.style.opacity = '0.5';
-                            generateLabel.classList.add('active-mode');
-                            smilesLabel.classList.remove('active-mode');
-                        }
-                    });
-                }
-            });
-            </script>
-            """)
-            
-            # Create container for generation mode inputs
-            with gr.Group(visible=True, elem_id="generate-box", elem_classes="input-box") as generate_group:
-                num_molecules = gr.Slider(
-                    minimum=10,
-                    maximum=250,
-                    value=100,
-                    step=10,
-                    label="Number of Molecules to Generate",
-                    info="This space runs on a CPU, which may result in slower performance. Generating 200 molecules takes approximately 6 minutes. Therefore, We set a 250-molecule cap. On a GPU, the model can generate 10,000 molecules in the same amount of time. Please check our GitHub repo for running our models on GPU."
-                )
-                
-                # Seed input used in generate mode
-                seed_num_generate = gr.Textbox(
-                    label="Random Seed (Optional)",
-                    value="",
-                    info="Set a specific seed for reproducible results, or leave empty for random generation"
-                )
-            
-            # Create container for SMILES input mode
-            with gr.Group(visible=False, elem_id="smiles-box", elem_classes="input-box") as smiles_group:
-                smiles_input = gr.Textbox(
-                    label="Input SMILES",
-                    info="Enter up to 100 SMILES strings, one per line",
-                    lines=10,
-                    placeholder="CC(=O)OC1=CC=CC=C1C(=O)O\nCCO\nC1=CC=C(C=C1)C(=O)O\n...",
-                )
-
-            # Handle visibility toggling between the two input modes
-            def toggle_visibility(checkbox_value):
-                return not checkbox_value, checkbox_value
-                
-            input_mode_switch.change(
-                fn=toggle_visibility,
-                inputs=[input_mode_switch],
-                outputs=[generate_group, smiles_group]
-            )
-
-            submit_button = gr.Button(
-                value="Generate Molecules",
-                variant="primary",
-                size="lg"
-            )
-            
-            # Helper function to determine which mode is active and which seed to use
-            def get_inputs(checkbox_value, num_mols, seed_gen, smiles):
-                mode = "smiles" if checkbox_value else "generate"
-                seed = "42" if checkbox_value else seed_gen  # Use default seed 42 for SMILES mode
-                return [mode, num_mols, seed, smiles]
-
-        with gr.Column(scale=2):
-            basic_metrics_df = gr.Dataframe(
-                headers=["Validity", "Uniqueness", "Novelty (Train)", "Novelty (Test)", "Novelty (Drug)", "Runtime (s)"],
-                elem_id="basic-metrics"
-            )
-                
-            advanced_metrics_df = gr.Dataframe(
-                headers=["QED", "SA Score", "Internal Diversity", "SNN (ChEMBL)", "SNN (Drug)", "Max Length"],
-                elem_id="advanced-metrics"
-                )
-
-            file_download = gr.File(
-                label="Download All Generated Molecules (SMILES format)",
-            )
-            
-            image_output = gr.Image(
-                label="Structures of Randomly Selected Generated Molecules",
-                elem_id="molecule_display"
-            )
+                    image_output = gr.Image(
+                        label="Structures of Randomly Selected Generated Molecules",
+                        elem_id="molecule_display"
+                    )
+
+        with gr.TabItem("Custom Input SMILES"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    # Reuse model selection for custom input
+                    model_name_custom = gr.Radio(
+                        choices=("DrugGEN-AKT1", "DrugGEN-CDK2", "DrugGEN-NoTarget"),
+                        value="DrugGEN-AKT1",
+                        label="Select Target Model",
+                        info="Choose which protein target or general model to use for evaluation"
+                    )
+                    custom_smiles = gr.Textbox(
+                        label="Input SMILES (one per line, maximum 100 molecules)",
+                        placeholder="C(C(=O)O)N\nCCO\n...",
+                        lines=10
+                    )
+                    custom_submit = gr.Button(
+                        value="Evaluate Custom SMILES",
+                        variant="primary",
+                        size="lg"
+                    )
+                with gr.Column(scale=2):
+                    basic_metrics_df_custom = gr.Dataframe(
+                        headers=["Validity", "Uniqueness", "Novelty (Train)", "Novelty (Inference)", "Novelty (Real Inhibitors)", "Runtime (s)"],
+                        elem_id="basic-metrics-custom"
+                    )
+                        
+                    advanced_metrics_df_custom = gr.Dataframe(
+                        headers=["QED", "SA Score", "Internal Diversity", "SNN (ChEMBL)", "SNN (Real Inhibitors)", "Average Length"],
+                        elem_id="advanced-metrics-custom"
+                    )
 
+                    file_download_custom = gr.File(
+                        label="Download All Molecules (SMILES format)"
+                    )
+                    
+                    image_output_custom = gr.Image(
+                        label="Structures of Randomly Selected Molecules",
+                        elem_id="molecule_display_custom"
+                    )
 
     gr.Markdown("### Created by the HUBioDataLab | [GitHub](https://github.com/HUBioDataLab/DrugGEN) | [Paper](https://arxiv.org/abs/2302.07868)")
 
-    submit_button.click(
-        fn=lambda model, checkbox, num_mols, seed_gen, smiles: function(
-            model, 
-            "smiles" if checkbox else "generate", 
-            num_mols, 
-            "42" if checkbox else seed_gen,  # Use default seed 42 for SMILES mode
-            smiles
-        ),
-        inputs=[model_name, input_mode_switch, num_molecules, seed_num_generate, smiles_input], 
+    # Set up the click actions for each tab.
+    classical_submit.click(
+        run_inference, 
+        inputs=[gr.State("Generate Molecules"), model_name, num_molecules, seed_num, gr.State("")],
         outputs=[
             image_output, 
             file_download,
             basic_metrics_df,
             advanced_metrics_df
-        ], 
-        api_name="inference"
+        ],
+        api_name="inference_classical"
+    )
+    
+    custom_submit.click(
+        run_inference, 
+        inputs=[gr.State("Custom Input SMILES"), model_name_custom, gr.State(0), gr.State(""), custom_smiles],
+        outputs=[
+            image_output_custom, 
+            file_download_custom,
+            basic_metrics_df_custom,
+            advanced_metrics_df_custom
+        ],
+        api_name="inference_custom"
     )
-#demo.queue(concurrency_count=1)
+
 demo.queue()
-demo.launch()
+demo.launch()
+