File size: 13,046 Bytes
02e480f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
import argparse


def get_parser(parser=None):
    if parser is None:
        parser = argparse.ArgumentParser()

    # Model
    #model_arg = parser.add_argument_group('Model')
    parser.add_argument('--n_head',
                           type=int, default=8,
                           help='GPT number of heads')
    parser.add_argument('--n_layer',
                           type=int, default=12,
                           help='GPT number of layers')
    parser.add_argument('--q_dropout',
                           type=float, default=0.5,
                           help='Encoder layers dropout')
    parser.add_argument('--d_dropout',
                           type=float, default=0.1,
                           help='Decoder layers dropout')
    parser.add_argument('--n_embd',
                           type=int, default=768,
                           help='Latent vector dimensionality')
    parser.add_argument('--fc_h',
                           type=int, default=512,
                           help='Fully connected hidden dimensionality')


    # Train
    #train_arg = parser.add_argument_group('Train')
    parser.add_argument('--n_batch',
                           type=int, default=512,
                           help='Batch size')
    parser.add_argument('--unlike_alpha',
                           type=float, default=1.0,
                           help='unlikelihood loss alpha weight')
    parser.add_argument('--from_scratch',
                           action='store_true', default=False,
                           help='train on qm9 from scratch')
    parser.add_argument('--unlikelihood',
                           action='store_true', default=False,
                           help='use unlikelihood loss with gpt pretrain')
    parser.add_argument('--grad_acc',
                           type=int, default=1,
                           help='number of batches to accumulate gradients')
    parser.add_argument('--checkpoint_every',
                           type=int, default=1000,
                           help='save checkpoint every x iterations')
    parser.add_argument('--clip_grad',
                           type=int, default=50,
                           help='Clip gradients to this value')
    parser.add_argument('--lr_start',
                           type=float, default=3 * 1e-4,
                           help='Initial lr value')
    parser.add_argument('--lr_end',
                           type=float, default=3 * 1e-4,
                           help='Maximum lr weight value')
    parser.add_argument('--lr_multiplier',
                           type=int, default=1,
                           help='lr weight multiplier')
    parser.add_argument('--n_last',
                           type=int, default=1000,
                           help='Number of iters to smooth loss calc')
    parser.add_argument('--n_jobs',
                           type=int, default=1,
                           help='Number of threads')
    parser.add_argument('--accelerator',
                        type=str, default='ddp',
                        help='The accelerator backend to use (previously known as distributed_backend)')
    parser.add_argument('--num_nodes',
                        type=int, default=1,
                        help='number of GPU nodes for distributed training')
    parser.add_argument('--device',
                        type=str, default='cuda',
                        help='Device to run: "cpu" or "cuda:<device number>"')
    parser.add_argument('--seed',
                        type=int, default=12345,
                        help='Seed')
    parser.add_argument('--init_params_from',
                        type=str, default='',
                        help='Path to a ckpt used to initialize the parameters if no restart_path is provided')
    parser.add_argument('--train_decoder_every',
                        type=int, default=10,
                        help='Optimize decoder params every n batches')
    parser.add_argument('--lr_decoder',
                        type=float, default=1e-4,
                        help='Learning rate for decoder part')
    parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus")
    parser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')
    parser.add_argument('--dist-backend', default='nccl', type=str, help='distributed backend')
    parser.add_argument('--save_checkpoint_path', default='/data', help='checkpoint saving path')
    parser.add_argument('--load_checkpoint_path', default='', help='checkpoint loading path')

    #common_arg = parser.add_argument_group('Common')
    parser.add_argument('--vocab_load',
                            type=str, required=False,
                            help='Where to load the vocab')
    parser.add_argument('--n_samples',
                            type=int, required=False,
                            help='Number of samples to sample')
    parser.add_argument('--gen_save',
                            type=str, required=False,
                            help='Where to save the gen molecules')
    parser.add_argument("--max_len",
                            type=int, default=100,
                            help="Max of length of SMILES")
    parser.add_argument('--train_load',
                            type=str, required=False,
                            help='Where to load the model')
    parser.add_argument('--val_load',
                            type=str, required=False,
                            help='Where to load the model')
    parser.add_argument('--n_workers',
                            type=int, required=False, default=1,
                            help='Where to load the model')
    #beam search hyper parameters
    parser.add_argument('--beam_size', type=int, default=0,
                            help="Number of beams to generate")
    parser.add_argument('--num_seq_returned', type=int, default=0,
                            help="number of beams to be returned (must be <= beam_size")
    parser.add_argument('--min_len', type=int, default=1,
                            help="minimum length to be generated")
    parser.add_argument('--nucleus_thresh', type=float, default=.9,
                            help="nucleus sampling threshold")
    parser.add_argument('--finetune_path',
                           type=str, default="",
                           help='path to  trainer file to continue training')
    parser.add_argument('--restart_path',
                           type=str, default="",
                           help='path to  trainer file to continue training')
    parser.add_argument('--data_path',
                           type=str, default="",
                           help='path to pubchem file')
    parser.add_argument('--pretext_size',
                           type=int, default=0,
                           help='number of k-mers to pretext')
    parser.add_argument('--model_save_dir',
                            type=str, required=False, default='./models_dump/',
                            help='Where to save the models/log/config/vocab')
    parser.add_argument('--model_save',
                            type=str, required=False, default='model.pt',
                            help='Where to save the model')
    #parser.add_argument('--save_frequency',
    #                        type=int, default=20,
    #                        help='How often to save the model')
    parser.add_argument('--num_epoch',
                            type=int, default=1,
                            help='number of epochs to train')
    #parser.add_argument('--num_iter',
    #                        type=int, default=-1,
    #                        help='how many itersations per epoch (for unlikelihood tuning)')
    parser.add_argument('--log_file',
                            type=str, required=False,
                            help='Where to save the log')
    parser.add_argument('--tb_loc',
                            type=str, required=False,
                            help='Where to save the tensorflow location')
    parser.add_argument('--config_save',
                            type=str, required=False,
                            help='Where to save the config')
    parser.add_argument('--vocab_save',
                            type=str,
                            help='Where to save the vocab')

   # resume_arg = parser.add_argument_group('Resume')
    parser.add_argument('--debug',
                           default=False, action='store_true',
                           help='do not erase cache at end of program')
    parser.add_argument('--fast_dev_run',
                           default=False,
                           help='This flag runs a “unit test” by running n if set to n (int) else 1 if set to True training and validation batch(es).')
    parser.add_argument('--freeze_model',
                           default=False, action='store_true',
                           help='freeze weights of bert model during fine tuning')
    parser.add_argument('--resume',
                           default=False, action='store_true',
                           help='Resume from a saved model')
    parser.add_argument('--rotate',
                           default=False, action='store_true',
                           help='use rotational relative embedding')
    parser.add_argument('--model_load',
                            type=str, required=False,
                            help='Where to load the model')
    parser.add_argument('--root_dir',
                            type=str, required=False, default='.',
                            help='location of root dir')
    parser.add_argument('--config_load',
                            type=str, required=False,
                            help='Where to load the config')
    parser.add_argument('--gpus',
                            type=int, required=False, default=1,
                            help='number of gpus to use')
    #parser.add_argument('--start_epoch',
    #                        type=int, required=False, default=0,
    #                        help='Where to load the config')

    parser.add_argument('--model_arch',
                            type=str, required=False,
                            help='used to teack model arch in params')
    parser.add_argument('--eval_every',
                            type=int, default=50000,
                            help='run evaluation every x iterations')
    parser.add_argument('--num_feats',
                            type=int, required=False, default=32,
                            help='number of random reatures for FAVOR+')
    parser.add_argument('--max_epochs',
                            type=int, required=False, default=1,
                            help='max number of epochs')

    # debug() FINE TUNEING
    # parser.add_argument('--save_dir', type=str, required=True)
    parser.add_argument('--mode',
                           type=str, default='cls',
                           help='type of pooling to use')
    parser.add_argument("--dataset_length", type=int, default=None, required=False)
    parser.add_argument("--num_workers", type=int, default=0, required=False)
    parser.add_argument("--dropout", type=float, default=0.1, required=False)
    #parser.add_argument("--dims", type=int, nargs="*", default="", required=False)
    parser.add_argument(
        "--smiles_embedding",
        type=str,
        default="/dccstor/medscan7/smallmolecule/runs/ba-predictor/small-data/embeddings/protein/ba_embeddings_tanh_512_2986138_2.pt",
    )
    # parser.add_argument("--train_pct", type=str, required=False, default="95")
    #parser.add_argument("--aug", type=int, required=True)
    parser.add_argument("--dataset_name", type=str, required=False, default="sol")
    parser.add_argument("--measure_name", type=str, required=False, default="measure")
    parser.add_argument("--smi_ted_version", type=str, required=True, default="v1")
    #parser.add_argument("--emb_type", type=str, required=True)
    #parser.add_argument("--checkpoints_folder", type=str, required=True)
    #parser.add_argument("--results_dir", type=str, required=True)
    #parser.add_argument("--patience_epochs", type=int, required=True)

    parser.add_argument(
        "--data_root",
        type=str,
        required=False,
        default="/dccstor/medscan7/smallmolecule/runs/ba-predictor/small-data/affinity",
    )
    # parser.add_argument("--use_bn", type=int, default=0)
    parser.add_argument("--use_linear", type=int, default=0)

    parser.add_argument("--lr", type=float, default=0.001)
    # parser.add_argument("--weight_decay", type=float, default=5e-4)
    # parser.add_argument("--val_check_interval", type=float, default=1.0)
    parser.add_argument("--batch_size", type=int, default=64)

    return parser
def parse_args():
    parser = get_parser()
    args = parser.parse_args()
    return args