SBVS/args.py

import argparse
import torch
from chemprop.features import get_available_features_generators

def add_args():
    """
    Adds predict arguments to an ArgumentParser.

    :param parser: An ArgumentParser.
    """
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # general arguments
    parser.add_argument('--gpu', type=int, default=1,
                        # choices=list(range(torch.cuda.device_count())),
                        help='Which GPU to use')
    parser.add_argument('--no_cuda', action='store_true', default=False,
                        help='Turn off cuda')
    parser.add_argument('--mode', type=str, default='train',
                        choices=['train', 'inference', 'retrain', 'finetune',
                                 'baseline_QSAR', 'baseline_CPI', 'baseline_inference'],
                        help='Mode to run script in')
    parser.add_argument('--print', action='store_true', default=False,
                        help='Print log')
    parser.add_argument('--data_path', type=str,
                        help='Path to CSV file containing training data',
                        default=None)
    parser.add_argument('--ref_path', type=str,
                        help='Path to CSV file containing refence data for scaling in inference mode',
                        default=None)
    parser.add_argument('--lig_file', type=str, 
                        help='Path to H5 file containing ligand data',
                        default=None)
    parser.add_argument('--model_path', type=str, default=None,
                        help='Path to model checkpoint (.pt file) for inference, retrain, or finetune')
    parser.add_argument('--dataset_type', type=str, choices=['classification', 'regression', 'joint'],
                        help='Type of dataset. \
                        Joint mode is only for GGAP-CPI, data should include "Activity_Type" columns')
    parser.add_argument('--save_dir', type=str, default=None,
                        help='dir name in exp_results folder where predictions will be saved')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed')
    parser.add_argument('--split_type', type=str, default='moleculeACE',
                        choices=['random', 'scaffold_balanced', 'moleculeACE', 'predetermined'],
                        help='Method of splitting the data into training, validation, and test')
    parser.add_argument('--split_sizes', type=float, nargs='+',
                        default=[0.8, 0.1, 0.1],
                        help='Proportions of data to use for training, validation, and test')
    parser.add_argument('--features_scaling', action='store_true', default=False,
                        help='Turn on scaling of features')                  
    parser.add_argument('--features_generator', type=str, nargs='*',
                        choices=get_available_features_generators(),
                        help='Method of generating additional features')
    parser.add_argument('--features_path', type=str, nargs='*', default=None,
                        help='Path to features to use in FNN (instead of features_generator)')
    parser.add_argument('--no_features_scaling', action='store_true', default=False,
                        help='Turn off scaling of features')
    parser.add_argument('--max_data_size', type=int, default=None,
                        help='Maximum number of data points to load')
    
    # training arguments
    parser.add_argument('--checkpoint_path', type=str,
                        default='KANO_model/dumped/pretrained_graph_encoder/original_CMPN_0623_1350_14000th_epoch.pkl',
                        help='Path to model checkpoint (.pt file)')
    # parser.add_argument('--loss', type=str, default='MSE CLS CL',
    #                     help='Loss function seperated with space. MSE: mean squared error, CLS: cross entropy loss, CL: contrastive loss')
    parser.add_argument('--loss_weights', type=str, default='1 1',
                        help='Weights for MSE, CLS, and CL loss functions seperated with space'
                        'Note: MSE: mean squared error, CLS: cross entropy loss, CL: contrastive loss'
                        'Set 0 to ignore the specific loss function')
    parser.add_argument('--type_thre', type=int, default=4, 
                        help='Which type of bioassay data to use for training, only for joint mode')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Batch size')
    parser.add_argument('--epochs', type=int, default=100,
                        help='Number of epochs')
    parser.add_argument('--lr', type=float, default=1e-4,
                        help='Learning rate')
    parser.add_argument('--temperature', type=float, default=0.1)
    parser.add_argument('--warmup_epochs', type=int, default=2)
    parser.add_argument('--encoder_name', type=str, default='CMPNN',
                        help='selected molecule encoder')
    parser.add_argument('--metric', type=str, default='rmse',
                        choices=['auc', 'prc-auc', 'rmse', 'mae', 'r2', 'accuracy', 'cross_entropy'],
                        help='Metric to optimize during training')
    # model arguments
    # you may not able to change most of these arguments if you use a pretrained model
    parser.add_argument('--train_model', type=str, default='GGAP_CPI', 
                        choices=['GGAP_CPI', 'KANO_ESM'], 
                        help='KANO_Prot for CPI-type model, KANO_ESM as a baseline model')
    parser.add_argument('--ablation', type=str, default='none', 
                        choices=['none', 'KANO', 'GCN', 'Attn', 'ESM'], 
                        help='Ablation study for KANO_Prot model')
    parser.add_argument('--baseline_model', type=str, default=None,
                        choices=['MLP', 'SVM', 'RF', 'GBM', 'KNN',
                                 'GAT', 'GCN', 'AFP', 'MPNN', 'CNN',
                                 'Transformer','LSTM', 'KANO',
                                 'DeepDTA', 'GraphDTA', 'MolTrans',
                                 'HyperAttentionDTI', 'PerceiverCPI',
                                 'ECFP_ESM_GBM', 'ECFP_ESM_RF', 
                                 'KANO_ESM_GBM', 'KANO_ESM_RF'],
                        help='Type of baseline model to train if select mode as baseline_QSAR or baseline_CPI')
    parser.add_argument('--hidden_size', type=int, default=300)
    parser.add_argument('--ffn_hidden_size', type=int, default=300)
    parser.add_argument('--ffn_num_layers', type=int, default=2)
    parser.add_argument('--dropout', type=float, default=0.0)
    parser.add_argument('--activation', type=str, default='ReLU')
    parser.add_argument('--depth', type=int, default=3)
    parser.add_argument('--step', type=str, default='functional_prompt')
    parser.add_argument('--num_heads', type=int, default=5)
    parser.add_argument('--pooling', type=str, default='cross_attn', choices=['cross_attn', 'mean'])

    args = parser.parse_args()
    # add and modify some args
    if '.csv' not in args.data_path:
        args.data_path += '.csv'
    args.endpoint_type = args.data_path.split('/')[1]
    args.data_name = args.data_path.split('/')[-1].split('.')[0]
    if not args.no_cuda and torch.cuda.is_available():
        args.cuda = True
    else:
        args.cuda = False
    args.prompt = False
    args.atom_messages = False
    args.use_input_features = None
    args.bias = False
    args.undirected = False
    args.features_only = False
    args.max_lr = args.lr * 10
    args.init_lr = args.lr
    args.final_lr = args.lr
    args.num_lrs = 1
    args.num_runs = 1
    args.smiles_columns = 'smiles'
    args.target_columns = 'y'
    args.output_size = args.num_tasks = 1
    args.prot_dir = 'data/Protein_pretrained_feat'

    loss_func = ['MSE', 'CLS']
    loss_wt = args.loss_weights.split(' ')
    args.loss_func_wt = dict(zip(loss_func, loss_wt))
    print(args.loss_func_wt)
    if args.baseline_model in ['DeepDTA', 'GraphDTA', 'MolTrans']:
        args.mode == 'baseline_CPI'
    if args.metric in ['auc', 'prc-auc', 'accuracy', 'r2']:
        args.minimize_score = False
    elif args.metric in ['rmse', 'mae', 'cross_entropy']:
        args.minimize_score = True

    return args