inference.py 19.7 KB
Newer Older
Shenggan's avatar
Shenggan committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import random
import sys
import time
21
from datetime import date
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
22
23
import tempfile
import contextlib
zhuww's avatar
zhuww committed
24
import logging
Shenggan's avatar
Shenggan committed
25
26
27

import numpy as np
import torch
28
import torch.multiprocessing as mp
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
29
import pickle
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
30
import shutil
31
from fastfold.model.hub import AlphaFold
Shenggan's avatar
Shenggan committed
32

33
import fastfold
34
35
36
import fastfold.relax.relax as relax
from fastfold.common import protein, residue_constants
from fastfold.config import model_config
37
from fastfold.model.fastnn import set_chunk_size
38
from fastfold.data import data_pipeline, feature_pipeline, templates
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
39
from fastfold.data.tools import hhsearch, hmmsearch
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
40
41
from fastfold.workflow.template import FastFoldDataWorkFlow, FastFoldMultimerDataWorkFlow

42
from fastfold.utils.inject_fastnn import inject_fastnn
43
from fastfold.data.parsers import parse_fasta
44
45
46
from fastfold.utils.import_weights import import_jax_weights_
from fastfold.utils.tensor_utils import tensor_tree_map

zhuww's avatar
zhuww committed
47
48
49
50
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)

oahzxl's avatar
oahzxl committed
51
52
53
if int(torch.__version__.split(".")[0]) >= 1 and int(torch.__version__.split(".")[1]) > 11:
    torch.backends.cuda.matmul.allow_tf32 = True

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
54
55
56
57
58
59
@contextlib.contextmanager
def temp_fasta_file(fasta_str: str):
    with tempfile.NamedTemporaryFile('w', suffix='.fasta') as fasta_file:
        fasta_file.write(fasta_str)
        fasta_file.seek(0)
        yield fasta_file.name
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

def add_data_args(parser: argparse.ArgumentParser):
    parser.add_argument(
        '--uniref90_database_path',
        type=str,
        default=None,
    )
    parser.add_argument(
        '--mgnify_database_path',
        type=str,
        default=None,
    )
    parser.add_argument(
        '--pdb70_database_path',
        type=str,
        default=None,
    )
    parser.add_argument(
        '--uniclust30_database_path',
        type=str,
        default=None,
    )
    parser.add_argument(
        '--bfd_database_path',
        type=str,
        default=None,
    )
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
87
88
89
90
91
92
93
94
95
96
    parser.add_argument(
        "--pdb_seqres_database_path",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--uniprot_database_path",
        type=str,
        default=None,
    )
97
98
99
100
    parser.add_argument('--jackhmmer_binary_path', type=str, default='/usr/bin/jackhmmer')
    parser.add_argument('--hhblits_binary_path', type=str, default='/usr/bin/hhblits')
    parser.add_argument('--hhsearch_binary_path', type=str, default='/usr/bin/hhsearch')
    parser.add_argument('--kalign_binary_path', type=str, default='/usr/bin/kalign')
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
101
102
    parser.add_argument("--hmmsearch_binary_path", type=str, default="hmmsearch")
    parser.add_argument("--hmmbuild_binary_path", type=str, default="hmmbuild")
103
104
105
106
107
108
109
    parser.add_argument(
        '--max_template_date',
        type=str,
        default=date.today().strftime("%Y-%m-%d"),
    )
    parser.add_argument('--obsolete_pdbs_path', type=str, default=None)
    parser.add_argument('--release_dates_path', type=str, default=None)
110
    parser.add_argument('--chunk_size', type=int, default=None)
LuGY's avatar
LuGY committed
111
    parser.add_argument('--enable_workflow', default=False, action='store_true', help='run inference with ray workflow or not')
oahzxl's avatar
oahzxl committed
112
    parser.add_argument('--inplace', default=False, action='store_true')
Shenggan's avatar
Shenggan committed
113

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
114

115
116
117
118
def inference_model(rank, world_size, result_q, batch, args):
    os.environ['RANK'] = str(rank)
    os.environ['LOCAL_RANK'] = str(rank)
    os.environ['WORLD_SIZE'] = str(world_size)
119
    # init distributed for Dynamic Axial Parallelism
120
    fastfold.distributed.init_dap()
121
    torch.cuda.set_device(rank)
Shenggan's avatar
Shenggan committed
122
    config = model_config(args.model_name)
123
124
    if args.chunk_size:
        config.globals.chunk_size = args.chunk_size
oahzxl's avatar
oahzxl committed
125
    config.globals.inplace = args.inplace
126
    config.globals.is_multimer = args.model_preset == 'multimer'
Shenggan's avatar
Shenggan committed
127
128
129
    model = AlphaFold(config)
    import_jax_weights_(model, args.param_path, version=args.model_name)

130
    model = inject_fastnn(model)
Shenggan's avatar
Shenggan committed
131
    model = model.eval()
132
    model = model.cuda()
Shenggan's avatar
Shenggan committed
133

134
135
    set_chunk_size(model.globals.chunk_size)

136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    with torch.no_grad():
        batch = {k: torch.as_tensor(v).cuda() for k, v in batch.items()}

        t = time.perf_counter()
        out = model(batch)
        print(f"Inference time: {time.perf_counter() - t}")

    out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

    result_q.put(out)

    torch.distributed.barrier()
    torch.cuda.synchronize()


def main(args):
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
152
153
154
155
156
157
158
159
    if args.model_preset == "multimer":
        inference_multimer_model(args)
    else:
        inference_monomer_model(args)


def inference_multimer_model(args):
    print("running in multimer mode...")
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
160
    config = model_config(args.model_name)
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
161
162
163
164
165
166
167
168
169
170
171
172
    
    predict_max_templates = 4

    template_featurizer = templates.HmmsearchHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
        max_hits=predict_max_templates,
        kalign_binary_path=args.kalign_binary_path,
        release_dates_path=args.release_dates_path,
        obsolete_pdbs_path=args.obsolete_pdbs_path,
    )

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
173
    if(not args.use_precomputed_alignments):
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
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
            if args.enable_workflow:
                print("Running alignment with ray workflow...")
                alignment_runner = FastFoldMultimerDataWorkFlow(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    hhblits_binary_path=args.hhblits_binary_path,
                    hmmsearch_binary_path=args.hmmsearch_binary_path,
                    hmmbuild_binary_path=args.hmmbuild_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
                    mgnify_database_path=args.mgnify_database_path,
                    bfd_database_path=args.bfd_database_path,
                    uniclust30_database_path=args.uniclust30_database_path,
                    uniprot_database_path=args.uniprot_database_path,
                    pdb_seqres_database_path=args.pdb_seqres_database_path,
                    use_small_bfd=(args.bfd_database_path is None),
                    no_cpus=args.cpus
                )
            else:
                alignment_runner = data_pipeline.AlignmentRunnerMultimer(
                    jackhmmer_binary_path=args.jackhmmer_binary_path,
                    hhblits_binary_path=args.hhblits_binary_path,
                    hmmsearch_binary_path=args.hmmsearch_binary_path,
                    hmmbuild_binary_path=args.hmmbuild_binary_path,
                    uniref90_database_path=args.uniref90_database_path,
                    mgnify_database_path=args.mgnify_database_path,
                    bfd_database_path=args.bfd_database_path,
                    uniclust30_database_path=args.uniclust30_database_path,
                    uniprot_database_path=args.uniprot_database_path,
                    pdb_seqres_database_path=args.pdb_seqres_database_path,
                    use_small_bfd=(args.bfd_database_path is None),
                    no_cpus=args.cpus
                )
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
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
    else:
        alignment_runner = None

    monomer_data_processor = data_pipeline.DataPipeline(
        template_featurizer=template_featurizer,
    )


    data_processor = data_pipeline.DataPipelineMultimer(
            monomer_data_pipeline=monomer_data_processor,
    )

    output_dir_base = args.output_dir
    random_seed = args.data_random_seed
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
    
    feature_processor = feature_pipeline.FeaturePipeline(
        config.data
    )

    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
    if(not args.use_precomputed_alignments):
        alignment_dir = os.path.join(output_dir_base, "alignments")
    else:
        alignment_dir = args.use_precomputed_alignments

    # Gather input sequences
    fasta_path = args.fasta_path
    with open(fasta_path, "r") as fp:
        data = fp.read()

    lines = [
        l.replace('\n', '') 
        for prot in data.split('>') for l in prot.strip().split('\n', 1)
    ][1:]
    tags, seqs = lines[::2], lines[1::2]


    for tag, seq in zip(tags, seqs):
        local_alignment_dir = os.path.join(alignment_dir, tag)
        if(args.use_precomputed_alignments is None):
            if not os.path.exists(local_alignment_dir):
                os.makedirs(local_alignment_dir)
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
250
251
252
            else:
                shutil.rmtree(local_alignment_dir)
                os.makedirs(local_alignment_dir)
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
253
254
255
            
            chain_fasta_str = f'>chain_{tag}\n{seq}\n'
            with temp_fasta_file(chain_fasta_str) as chain_fasta_path:
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
256
257
258
259
260
261
262
263
                if args.enable_workflow:
                    print("Running alignment with ray workflow...")
                    t = time.perf_counter()
                    alignment_runner.run(chain_fasta_path, alignment_dir=local_alignment_dir)
                    print(f"Alignment data workflow time: {time.perf_counter() - t}")
                else:
                    alignment_runner.run(chain_fasta_path, local_alignment_dir)
                
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
264
265
266
267
268
269
270
                print(f"Finished running alignment for {tag}")
                
    local_alignment_dir = alignment_dir

    feature_dict = data_processor.process_fasta(
        fasta_path=fasta_path, alignment_dir=local_alignment_dir
    )
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
271
    # feature_dict = pickle.load(open("/home/lcmql/data/features_pdb1o5d.pkl", "rb"))
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
272
273
274
275

    processed_feature_dict = feature_processor.process_features(
        feature_dict, mode='predict', is_multimer=True,
    )
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
276

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
    batch = processed_feature_dict

    manager = mp.Manager()
    result_q = manager.Queue()
    torch.multiprocessing.spawn(inference_model, nprocs=args.gpus, args=(args.gpus, result_q, batch, args))

    out = result_q.get()

    # Toss out the recycling dimensions --- we don't need them anymore
    batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
    
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)

    plddt_b_factors = np.repeat(plddt[..., None], residue_constants.atom_type_num, axis=-1)

    unrelaxed_protein = protein.from_prediction(features=batch,
                                                result=out,
                                                b_factors=plddt_b_factors)

    # Save the unrelaxed PDB.
    unrelaxed_output_path = os.path.join(args.output_dir,
                                            f'{tag}_{args.model_name}_unrelaxed.pdb')
    with open(unrelaxed_output_path, 'w') as f:
        f.write(protein.to_pdb(unrelaxed_protein))

    amber_relaxer = relax.AmberRelaxation(
        use_gpu=True,
        **config.relax,
    )

    # Relax the prediction.
    t = time.perf_counter()
    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
    print(f"Relaxation time: {time.perf_counter() - t}")

    # Save the relaxed PDB.
    relaxed_output_path = os.path.join(args.output_dir,
                                        f'{tag}_{args.model_name}_relaxed.pdb')
    with open(relaxed_output_path, 'w') as f:
        f.write(relaxed_pdb_str)


Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
320
321
def inference_monomer_model(args):
    print("running in monomer mode...")
322
323
    config = model_config(args.model_name)

Shenggan's avatar
Shenggan committed
324
325
326
327
328
329
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
        max_hits=config.data.predict.max_templates,
        kalign_binary_path=args.kalign_binary_path,
        release_dates_path=args.release_dates_path,
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
330
331
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
Shenggan's avatar
Shenggan committed
332

333
334
335
336
337
338
    use_small_bfd = args.preset == 'reduced_dbs'  # (args.bfd_database_path is None)
    if use_small_bfd:
        assert args.bfd_database_path is not None
    else:
        assert args.bfd_database_path is not None
        assert args.uniclust30_database_path is not None
Shenggan's avatar
Shenggan committed
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355

    data_processor = data_pipeline.DataPipeline(template_featurizer=template_featurizer,)

    output_dir_base = args.output_dir
    random_seed = args.data_random_seed
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
    if (args.use_precomputed_alignments is None):
        alignment_dir = os.path.join(output_dir_base, "alignments")
    else:
        alignment_dir = args.use_precomputed_alignments

    # Gather input sequences
    with open(args.fasta_path, "r") as fp:
356
357
        fasta = fp.read()
    seqs, tags = parse_fasta(fasta)
LuGY's avatar
LuGY committed
358
    seq, tag = seqs[0], tags[0]
Shenggan's avatar
Shenggan committed
359

LuGY's avatar
LuGY committed
360
361
362
363
364
365
    print(f"tag:{tag}\nseq[{len(seq)}]:{seq}")
    batch = [None]
    
    fasta_path = os.path.join(args.output_dir, "tmp.fasta")
    with open(fasta_path, "w") as fp:
        fp.write(f">{tag}\n{seq}")
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
366

LuGY's avatar
LuGY committed
367
368
    print("Generating features...")
    local_alignment_dir = os.path.join(alignment_dir, tag)
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
369

LuGY's avatar
LuGY committed
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
    if (args.use_precomputed_alignments is None):
        if not os.path.exists(local_alignment_dir):
            os.makedirs(local_alignment_dir)
        if args.enable_workflow:
            print("Running alignment with ray workflow...")
            alignment_data_workflow_runner = FastFoldDataWorkFlow(
                jackhmmer_binary_path=args.jackhmmer_binary_path,
                hhblits_binary_path=args.hhblits_binary_path,
                hhsearch_binary_path=args.hhsearch_binary_path,
                uniref90_database_path=args.uniref90_database_path,
                mgnify_database_path=args.mgnify_database_path,
                bfd_database_path=args.bfd_database_path,
                uniclust30_database_path=args.uniclust30_database_path,
                pdb70_database_path=args.pdb70_database_path,
                use_small_bfd=use_small_bfd,
                no_cpus=args.cpus,
            )
            t = time.perf_counter()
            alignment_data_workflow_runner.run(fasta_path, alignment_dir=local_alignment_dir)
            print(f"Alignment data workflow time: {time.perf_counter() - t}")
        else:
            alignment_runner = data_pipeline.AlignmentRunner(
                jackhmmer_binary_path=args.jackhmmer_binary_path,
                hhblits_binary_path=args.hhblits_binary_path,
                hhsearch_binary_path=args.hhsearch_binary_path,
                uniref90_database_path=args.uniref90_database_path,
                mgnify_database_path=args.mgnify_database_path,
                bfd_database_path=args.bfd_database_path,
                uniclust30_database_path=args.uniclust30_database_path,
                pdb70_database_path=args.pdb70_database_path,
                use_small_bfd=use_small_bfd,
                no_cpus=args.cpus,
            )
            alignment_runner.run(fasta_path, local_alignment_dir)
            
    feature_dict = data_processor.process_fasta(fasta_path=fasta_path,
                                            alignment_dir=local_alignment_dir)

    # Remove temporary FASTA file
    os.remove(fasta_path)
Shenggan's avatar
Shenggan committed
410

LuGY's avatar
LuGY committed
411
412
413
414
    processed_feature_dict = feature_processor.process_features(
        feature_dict,
        mode='predict',
    )
415

LuGY's avatar
LuGY committed
416
417
418
    batch = processed_feature_dict

    manager = mp.Manager()
zhuww's avatar
zhuww committed
419
    result_q = manager.Queue()                         
LuGY's avatar
LuGY committed
420
    torch.multiprocessing.spawn(inference_model, nprocs=args.gpus, args=(args.gpus, result_q, batch, args))
Shenggan's avatar
Shenggan committed
421

LuGY's avatar
LuGY committed
422
    out = result_q.get()
Shenggan's avatar
Shenggan committed
423

LuGY's avatar
LuGY committed
424
425
426
427
428
    # Toss out the recycling dimensions --- we don't need them anymore
    batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
    
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)
429

LuGY's avatar
LuGY committed
430
    plddt_b_factors = np.repeat(plddt[..., None], residue_constants.atom_type_num, axis=-1)
Shenggan's avatar
Shenggan committed
431

LuGY's avatar
LuGY committed
432
433
434
    unrelaxed_protein = protein.from_prediction(features=batch,
                                                result=out,
                                                b_factors=plddt_b_factors)
Shenggan's avatar
Shenggan committed
435

LuGY's avatar
LuGY committed
436
437
438
439
440
    # Save the unrelaxed PDB.
    unrelaxed_output_path = os.path.join(args.output_dir,
                                            f'{tag}_{args.model_name}_unrelaxed.pdb')
    with open(unrelaxed_output_path, 'w') as f:
        f.write(protein.to_pdb(unrelaxed_protein))
Shenggan's avatar
Shenggan committed
441

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
    # amber_relaxer = relax.AmberRelaxation(
    #     use_gpu=True,
    #     **config.relax,
    # )

    # # Relax the prediction.
    # t = time.perf_counter()
    # relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
    # print(f"Relaxation time: {time.perf_counter() - t}")

    # # Save the relaxed PDB.
    # relaxed_output_path = os.path.join(args.output_dir,
    #                                     f'{tag}_{args.model_name}_relaxed.pdb')
    # with open(relaxed_output_path, 'w') as f:
    #     f.write(relaxed_pdb_str)
zhuww's avatar
zhuww committed
457
458
459
460
461
462
463
464
465
    
    if(args.save_outputs):
            output_dict_path = os.path.join(
                args.output_dir, f'{tag}_{args.model_name}_output_dict.pkl'
            )
            with open(output_dict_path, "wb") as fp:
                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)

            logger.info(f"Model output written to {output_dict_path}...")
Shenggan's avatar
Shenggan committed
466

Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
467

Shenggan's avatar
Shenggan committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "fasta_path",
        type=str,
    )
    parser.add_argument(
        "template_mmcif_dir",
        type=str,
    )
    parser.add_argument("--use_precomputed_alignments",
                        type=str,
                        default=None,
                        help="""Path to alignment directory. If provided, alignment computation 
                is skipped and database path arguments are ignored.""")
    parser.add_argument(
        "--output_dir",
        type=str,
        default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
    )
    parser.add_argument("--model_name",
                        type=str,
                        default="model_1",
                        help="""Name of a model config. Choose one of model_{1-5} or 
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
493
             model_{1-5}_ptm or model_{1-5}_multimer, as defined on the AlphaFold GitHub.""")
Shenggan's avatar
Shenggan committed
494
495
496
497
498
    parser.add_argument("--param_path",
                        type=str,
                        default=None,
                        help="""Path to model parameters. If None, parameters are selected
             automatically according to the model name from 
499
             ./data/params""")
zhuww's avatar
zhuww committed
500
501
502
503
    parser.add_argument(
        "--save_outputs", action="store_true", default=False,
        help="Whether to save all model outputs, including embeddings, etc."
    )
Shenggan's avatar
Shenggan committed
504
505
506
507
    parser.add_argument("--cpus",
                        type=int,
                        default=12,
                        help="""Number of CPUs with which to run alignment tools""")
508
509
510
511
    parser.add_argument("--gpus",
                        type=int,
                        default=1,
                        help="""Number of GPUs with which to run inference""")
Shenggan's avatar
Shenggan committed
512
513
    parser.add_argument('--preset',
                        type=str,
514
                        default='full_dbs',
Shenggan's avatar
Shenggan committed
515
516
                        choices=('reduced_dbs', 'full_dbs'))
    parser.add_argument('--data_random_seed', type=str, default=None)
Fazzie-Maqianli's avatar
Fazzie-Maqianli committed
517
518
519
520
521
522
523
524
    parser.add_argument(
        "--model_preset",
        type=str,
        default="monomer",
        choices=["monomer", "multimer"],
        help="Choose preset model configuration - the monomer model, the monomer model with "
        "extra ensembling, monomer model with pTM head, or multimer model",
    )
Shenggan's avatar
Shenggan committed
525
526
527
528
    add_data_args(parser)
    args = parser.parse_args()

    if (args.param_path is None):
529
        args.param_path = os.path.join("data", "params", "params_" + args.model_name + ".npz")
Shenggan's avatar
Shenggan committed
530
531

    main(args)