gradio_app.py

import gradio as gr
from trainer import Trainer
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

class DrugGENConfig:
    submodel='CrossLoss'
    act='relu'
    z_dim=16
    max_atom=45
    lambda_gp=1
    dim=128
    depth=1
    heads=8
    dec_depth=1
    dec_heads=8
    dec_dim=128
    mlp_ratio=3
    warm_up_steps=0
    dis_select='mlp'
    init_type='normal'
    batch_size=128
    epoch=50
    g_lr=0.00001
    d_lr=0.00001
    g2_lr=0.00001
    d2_lr=0.00001
    dropout=0.
    dec_dropout=0.
    n_critic=1
    beta1=0.9
    beta2=0.999
    resume_iters=None
    clipping_value=2
    features=False
    test_iters=10_000
    num_test_epoch=30_000
    inference_sample_num=1000
    num_workers=1
    mode="inference"
    inference_iterations=100
    inf_batch_size=1
    protein_data_dir='data/akt'
    drug_index='data/drug_smiles.index'
    drug_data_dir='data/akt'
    mol_data_dir='data'
    log_dir='experiments/logs'
    model_save_dir='experiments/models'
    # inference_model=""
    sample_dir='experiments/samples'
    result_dir="experiments/tboard_output"
    dataset_file="chembl45_train.pt"
    drug_dataset_file="akt_train.pt"
    raw_file='data/chembl_train.smi'
    drug_raw_file="data/akt_train.smi"
    inf_dataset_file="chembl45_test.pt"
    inf_drug_dataset_file='akt_test.pt'
    inf_raw_file='data/chembl_test.smi'
    inf_drug_raw_file="data/akt_test.smi"
    log_sample_step=1000
    set_seed=True
    seed=1
    resume=False
    resume_epoch=None
    resume_iter=None
    resume_directory=None
    
class ProtConfig(DrugGENConfig):
    submodel="Prot"
    inference_model="experiments/models/Prot"

class CrossLossConfig(DrugGENConfig):
    submodel="CrossLoss"
    inference_model="experiments/models/CrossLoss"

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


model_configs = {
    "Prot": ProtConfig(),
    "DrugGEN": CrossLossConfig(),
    "DrugGEN-NoTarget": NoTargetConfig(),
}



def function(model_name: str, mol_num: int, seed: int) -> tuple[PIL.Image, pd.DataFrame, str]:
    '''
    Returns:
    image, score_df, file path
    '''

    config = model_configs[model_name]
    config.inference_sample_num = mol_num
    config.seed = seed

    trainer = Trainer(config)
    scores = trainer.inference() # create scores_df out of this

    score_df = pd.DataFrame(scores, index=[0])
    old_model_names = {
        "DrugGEN": "CrossLoss",
        "DrugGEN-NoTarget": "NoTarget",
    }
    output_file_path = f'experiments/inference/{old_model_names[model_name]}/inference_drugs.txt'
    
    import os
    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")
    
    rng = random.Random(seed)
    
    selected_molecules = rng.choices(generated_molecule_list,k=12)
    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules]
    
    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() 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("Expand to display information about models", open=False):
                gr.Markdown("""
### Model Variations
- **DrugGEN**: composed of one GAN, the input of the GAN1 generator is the real molecules dataset and the GAN1 discriminator compares the generated molecules with the real inhibitors of the given target.
- **DrugGEN-NoTarget**: composed of one GAN, focuses on learning the chemical properties from the ChEMBL training dataset, no target-specific generation.
        """)
            model_name = gr.Radio(
                choices=("DrugGEN", "DrugGEN-NoTarget"),
                value="DrugGEN",
                label="Select a model to make inference",
                info=" DrugGEN-Prot and DrugGEN-CrossLoss models design molecules to target the AKT1 protein"
            )

            num_molecules = gr.Number(
                label="Number of molecules to generate",
                precision=0, # integer input
                minimum=1,
                value=1000,
                maximum=10_000,
            )
            seed_num = gr.Number(
                label="RNG seed value (can be used for reproducibility):",
                precision=0, # integer input
                minimum=0,
                value=42,
            )

            submit_button = gr.Button(
                value="Start Generating"
            )

        with gr.Column(scale=2):
            scores_df = gr.Dataframe(
                label="Scores",
                headers=["Runtime (seconds)", "Validity", "Uniqueness", "Novelty (Train)", "Novelty (Inference)"],
            )
            file_download = gr.File(
                label="Click to download generated molecules",
            )
            image_output = gr.Image(
                label="Structures of randomly selected 12 de novo molecules from the inference set:"
            )
            # ).style(
            #     height=200*4,
            #     width=200*3,
            # )

    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.launch()