app.py

import streamlit as st
from trainer import Trainer


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=False
    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(),
    "CrossLoss": CrossLossConfig(),
    "NoTarget": NoTargetConfig(),
}


with st.sidebar:
    st.title("DrugGEN")
    with st.form("model_selection_from"):

        model_name = st.radio(
        "Select a model to make inference",
        ('Prot', 'CrossLoss', 'NoTarget'))

        submitted = st.form_submit_button("Start Computing")

    if submitted:

        st.write(model_name)
    
# with st.spinner('Setting up the trainer class...'):
#     trainer = Trainer(ProtConfig())

# with st.spinner('Generating Molecules...'):
#     trainer.inference()

st.title("text")

if submitted:
    trainer = Trainer(model_configs[model_name])
    trainer.inference()
    st.success(f"Success with the inference of {model_name}")