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="experiments/models/DrugGEN/"
    sample_num=100
    disable_correction=False  # corresponds to correct=True in old config

    # Data configuration
    inf_smiles='data/chembl_test.smi'  # corresponds to inf_raw_file in old config
    train_smiles='data/chembl_train.smi'
    train_drug_smiles='data/akt1_train.smi'
    inf_batch_size=1
    mol_data_dir='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


class DrugGENAKT1Config(DrugGENConfig):
    submodel='DrugGEN'
    inference_model="experiments/models/DrugGEN-AKT1/"
    train_drug_smiles='data/akt1_train.smi'
    max_atom=45


class DrugGENCDK2Config(DrugGENConfig):
    submodel='DrugGEN'
    inference_model="experiments/models/DrugGEN-CDK2/"
    train_drug_smiles='data/cdk2_train.smi'
    max_atom=38


class NoTargetConfig(DrugGENConfig):
    submodel="NoTarget"
    inference_model="experiments/models/NoTarget/"
    train_drug_smiles='data/chembl_train.smi'  # No specific target, use general ChEMBL data


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



def function(model_name: str, num_molecules: int, seed_num: int) -> tuple[PIL.Image, pd.DataFrame, str]:
    '''
    Returns:
    image, score_df, file path
    '''
    if model_name == "DrugGEN-NoTarget":
        model_name = "NoTarget"
    
    config = model_configs[model_name]
    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!")


    inferer = Inference(config)
    start_time = time.time()
    scores = inferer.inference() # create scores_df out of this
    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]]
    })

    output_file_path = f'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 num_molecules > 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,
    )

    return molecule_image, score_df, new_path



with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) 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)")
            
            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), a serine/threonine-protein kinase that plays a key role in regulating cell survival, metabolism, and growth. AKT1 is a significant target in cancer therapy, particularly for breast, colorectal, and ovarian cancers.

The model learns from:
- General drug-like molecules from ChEMBL database
- Known AKT1 inhibitors
- Maximum atom count: 45

### DrugGEN-CDK2
This model targets the human CDK2 protein (UniProt ID: P24941), a cyclin-dependent kinase involved in cell cycle regulation. CDK2 inhibitors are being investigated for treating various cancers, particularly those with dysregulated cell cycle control.

The model learns from:
- General drug-like molecules from ChEMBL database
- Known CDK2 inhibitors
- Maximum atom count: 38

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

## How It Works
DrugGEN uses a graph-based generative adversarial network (GAN) architecture where:
1. The generator creates molecular graphs
2. The discriminator evaluates them against real molecules
3. The model learns to generate increasingly realistic and target-specific molecules

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 drugs

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

            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_num = gr.Textbox(
                label="Random Seed (Optional)",
                value="",
                info="Set a specific seed for reproducible results, or leave empty for random generation"
            )

            submit_button = gr.Button(
                value="Generate Molecules",
                variant="primary",
                size="lg"
            )

        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)",
                    )
                
                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"]
                    )
            
            with gr.Accordion("Generation Settings", open=False):
                gr.Markdown("""
                ## Technical Details
                
                - This demo runs on CPU which limits generation speed
                - Generating 200 molecules takes approximately 6 minutes
                - For faster generation or larger batches, run the model on GPU using our GitHub repository
                - The model uses a graph-based representation of molecules
                - Maximum atom count varies by model (AKT1: 45, CDK2: 38)
                """)

    gr.Markdown("### Created by the HU BioDataLab | [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], api_name="inference")
#demo.queue(concurrency_count=1)
demo.queue()
demo.launch()