app.py

import gradio as gr
from inference import Inference
import PIL
from PIL import Image
import pandas as pd
import random
from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem.Draw import IPythonConsole
import shutil
import os
import time

class DrugGENConfig:
    # Inference configuration
    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

    # Model configuration
    act='relu'
    max_atom=45
    dim=128
    depth=1
    heads=8
    mlp_ratio=3
    dropout=0.

    # Seed configuration
    set_seed=True
    seed=10

    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


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


class NoTargetConfig(DrugGENConfig):
    submodel="NoTarget"
    inference_model="/home/user/app/experiments/models/NoTarget/"


model_configs = {
    "DrugGEN-AKT1": DrugGENAKT1Config(),
    "DrugGEN-CDK2": DrugGENCDK2Config(),
    "DrugGEN-NoTarget": NoTargetConfig(),
}



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"
    
    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 seed_num is None or seed_num.strip() == "":
            config.seed = random.randint(0, 10000)
        else:
            try:
                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"

    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]],
        "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'/home/user/app/experiments/inference/{model_name}/inference_drugs.txt'

    new_path = f'{model_name}_denovo_mols.smi'
    os.rename(output_file_path, new_path)

    with open(new_path) as f:
        inference_drugs = f.read()

    generated_molecule_list = inference_drugs.split("\n")[:-1]

    rng = random.Random(config.seed)
    if len(generated_molecule_list) > 12:
        selected_molecules = 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]

    drawOptions = Draw.rdMolDraw2D.MolDrawOptions()
    drawOptions.prepareMolsBeforeDrawing = False
    drawOptions.bondLineWidth = 0.5

    molecule_image = Draw.MolsToGridImage(
        selected_molecules,
        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("""
    <style>
    #metrics-container {
        border: 1px solid rgba(128, 128, 128, 0.3);
        border-radius: 8px;
        padding: 15px;
        margin-top: 15px;
        margin-bottom: 15px;
        background-color: rgba(255, 255, 255, 0.05);
    }
    </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>
            </div>
            """)
            
            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).

### DrugGEN-CDK2
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

For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
                """)
            
            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

### 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

### Structural Metrics
- **Max Length**: Maximum 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)

### 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"
                    )
                    
                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"
            )


    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], 
        outputs=[
            image_output, 
            file_download,
            basic_metrics_df,
            advanced_metrics_df
        ], 
        api_name="inference"
    )
#demo.queue(concurrency_count=1)
demo.queue()
demo.launch()