run_pretrained_openfold.py 20.1 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 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.
15
import argparse
16
from copy import deepcopy
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
17
from datetime import date
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
18
import logging
19
import math
20
import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
21
import os
22
23
24
25

logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
26

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
27
import pickle
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
28
29
30
from pytorch_lightning.utilities.deepspeed import (
    convert_zero_checkpoint_to_fp32_state_dict
)
31
32
import random
import sys
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
33
34
import time
import torch
35
import re
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
36

37
38
39
40
41
42
43
44
45
46
47
48
49
torch_versions = torch.__version__.split(".")
torch_major_version = int(torch_versions[0])
torch_minor_version = int(torch_versions[1])
if(
    torch_major_version > 1 or 
    (torch_major_version == 1 and torch_minor_version >= 12)
):
    # Gives a large speedup on Ampere-class GPUs
    torch.set_float32_matmul_precision("high")

torch.set_grad_enabled(False)

from openfold.config import model_config, NUM_RES
50
from openfold.data import templates, feature_pipeline, data_pipeline
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
51
from openfold.model.model import AlphaFold
52
from openfold.model.torchscript import script_preset_
53
from openfold.np import residue_constants, protein
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
54
55
import openfold.np.relax.relax as relax
from openfold.utils.import_weights import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
56
57
    import_jax_weights_,
)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
58
from openfold.utils.tensor_utils import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
59
60
    tensor_tree_map,
)
61
62
63
64
from openfold.utils.trace_utils import (
    pad_feature_dict_seq,
    trace_model_,
)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
65
from scripts.utils import add_data_args
66

67

68
TRACING_INTERVAL = 50
69
70


Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
71
72
73
74
75
76
77
def precompute_alignments(tags, seqs, alignment_dir, args):
    for tag, seq in zip(tags, seqs):
        tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
78
        if(args.use_precomputed_alignments is None and not os.path.isdir(local_alignment_dir)):
79
            logger.info(f"Generating alignments for {tag}...")
80
81
                
            os.makedirs(local_alignment_dir)
82

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
83
84
85
86
87
88
89
90
91
92
93
94
            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,
                no_cpus=args.cpus,
            )
            alignment_runner.run(
95
                tmp_fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
96
            )
97
98
99
100
        else:
            logger.info(
                f"Using precomputed alignments for {tag} at {alignment_dir}..."
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
101
102
103
104
105

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


106
107
108
109
def round_up_seqlen(seqlen):
    return int(math.ceil(seqlen / TRACING_INTERVAL)) * TRACING_INTERVAL


Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
110
def run_model(model, batch, tag, args):
111
    with torch.no_grad(): 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
112
        # Disable templates if there aren't any in the batch
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
113
        model.config.template.enabled = model.config.template.enabled and any([
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
114
115
116
            "template_" in k for k in batch
        ])

117
        logger.info(f"Running inference for {tag}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
118
119
        t = time.perf_counter()
        out = model(batch)
120
        logger.info(f"Inference time: {time.perf_counter() - t}")
121
   
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
    return out


def prep_output(out, batch, feature_dict, feature_processor, args):
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)
    
    plddt_b_factors = np.repeat(
        plddt[..., None], residue_constants.atom_type_num, axis=-1
    )

    # Prep protein metadata
    template_domain_names = []
    template_chain_index = None
136
    if(feature_processor.config.common.use_templates and "template_domain_names" in feature_dict):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        template_domain_names = [
            t.decode("utf-8") for t in feature_dict["template_domain_names"]
        ]

        # This works because templates are not shuffled during inference
        template_domain_names = template_domain_names[
            :feature_processor.config.predict.max_templates
        ]

        if("template_chain_index" in feature_dict):
            template_chain_index = feature_dict["template_chain_index"]
            template_chain_index = template_chain_index[
                :feature_processor.config.predict.max_templates
            ]

    no_recycling = feature_processor.config.common.max_recycling_iters
    remark = ', '.join([
        f"no_recycling={no_recycling}",
        f"max_templates={feature_processor.config.predict.max_templates}",
156
        f"config_preset={args.config_preset}",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
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
    ])

    # For multi-chain FASTAs
    ri = feature_dict["residue_index"]
    chain_index = (ri - np.arange(ri.shape[0])) / args.multimer_ri_gap
    chain_index = chain_index.astype(np.int64)
    cur_chain = 0
    prev_chain_max = 0
    for i, c in enumerate(chain_index):
        if(c != cur_chain):
            cur_chain = c
            prev_chain_max = i + cur_chain * args.multimer_ri_gap

        batch["residue_index"][i] -= prev_chain_max

    unrelaxed_protein = protein.from_prediction(
        features=batch,
        result=out,
        b_factors=plddt_b_factors,
        chain_index=chain_index,
        remark=remark,
        parents=template_domain_names,
        parents_chain_index=template_chain_index,
    )

    return unrelaxed_protein


185
def parse_fasta(data):
186
    data = re.sub('>$', '', data, flags=re.M)
187
    lines = [
188
189
190
        l.replace('\n', '')
        for prot in data.split('>') for l in prot.strip().split('\n', 1)
    ][1:]
191
    tags, seqs = lines[::2], lines[1::2]
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
192

193
    tags = [t.split()[0] for t in tags]
194

195
    return tags, seqs
196

197

198
199
200
201
202
203
204
def generate_feature_dict(
    tags,
    seqs,
    alignment_dir,
    data_processor,
    args,
):
205
206
    tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
    if len(seqs) == 1:
207
        tag = tags[0]
208
209
210
211
212
213
214
        seq = seqs[0]
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
        feature_dict = data_processor.process_fasta(
            fasta_path=tmp_fasta_path, alignment_dir=local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
215
        )
216
217
218
219
220
221
222
223
224
225
226
227
    else:
        with open(tmp_fasta_path, "w") as fp:
            fp.write(
                '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
            )
        feature_dict = data_processor.process_multiseq_fasta(
            fasta_path=tmp_fasta_path, super_alignment_dir=alignment_dir,
        )

    # Remove temporary FASTA file
    os.remove(tmp_fasta_path)

228
    return feature_dict
229

230

231
232
233
234
235
236
237
def get_model_basename(model_path):
    return os.path.splitext(
                os.path.basename(
                    os.path.normpath(model_path)
                )
            )[0]

238

239
240
241
242
243
244
245
246
def make_output_directory(output_dir, model_name, multiple_model_mode):
    if multiple_model_mode:
        prediction_dir = os.path.join(output_dir, "predictions", model_name)
    else:
        prediction_dir = os.path.join(output_dir, "predictions")
    os.makedirs(prediction_dir, exist_ok=True)
    return prediction_dir

247

248
249
250
251
252
253
254
def count_models_to_evaluate(openfold_checkpoint_path, jax_param_path):
    model_count = 0
    if openfold_checkpoint_path:
        model_count += len(openfold_checkpoint_path.split(","))
    if jax_param_path:
        model_count += len(jax_param_path.split(","))
    return model_count
255

256

257
258
def load_models_from_command_line(args, config):
    # Create the output directory
259
260
261
262
263

    multiple_model_mode = count_models_to_evaluate(args.openfold_checkpoint_path, args.jax_param_path) > 1
    if multiple_model_mode:
        logger.info(f"evaluating multiple models")

264
265
    if args.jax_param_path:
        for path in args.jax_param_path.split(","):
266
267
            model_basename = get_model_basename(path)
            model_version = "_".join(model_basename.split("_")[1:])
268
269
270
            model = AlphaFold(config)
            model = model.eval()
            import_jax_weights_(
271
                model, path, version=model_version
272
273
            )
            model = model.to(args.model_device)
274
            logger.info(
275
                f"Successfully loaded JAX parameters at {path}..."
276
            )
277
278
            output_directory = make_output_directory(args.output_dir, model_basename, multiple_model_mode)
            yield model, output_directory
279
    
280
    if args.openfold_checkpoint_path:
Sam DeLuca's avatar
wip  
Sam DeLuca committed
281
        for path in args.openfold_checkpoint_path.split(","):
282
283
            model = AlphaFold(config)
            model = model.eval()
284
            checkpoint_basename = get_model_basename(path)
285
            if os.path.isdir(path):
286
                # A DeepSpeed checkpoint
287
288
289
                ckpt_path = os.path.join(
                    args.output_dir,
                    checkpoint_basename + ".pt",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
290
291
                )

292
293
                if not os.path.isfile(ckpt_path):
                    convert_zero_checkpoint_to_fp32_state_dict(
Sam DeLuca's avatar
wip  
Sam DeLuca committed
294
                        path,
295
296
                        ckpt_path,
                    )
297
298
                d = torch.load(ckpt_path)
                model.load_state_dict(d["ema"]["params"])
299
300
            else:
                ckpt_path = path
301
                d = torch.load(ckpt_path)
302

303
                if "ema" in d:
304
305
306
                    # The public weights have had this done to them already
                    d = d["ema"]["params"]
                model.load_state_dict(d)
307
            
308
            model = model.to(args.model_device)
309
            logger.info(
310
                f"Loaded OpenFold parameters at {path}..."
311
            )
312
313
            output_directory = make_output_directory(args.output_dir, checkpoint_basename, multiple_model_mode)
            yield model, output_directory
314
    
Sam DeLuca's avatar
wip  
Sam DeLuca committed
315
    if not args.jax_param_path and not args.openfold_checkpoint_path:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
316
317
318
319
320
        raise ValueError(
            "At least one of jax_param_path or openfold_checkpoint_path must "
            "be specified."
        )

321

322
323
def list_files_with_extensions(dir, extensions):
    return [f for f in os.listdir(dir) if f.endswith(extensions)]
324

325

326
327
328
329
def main(args):
    # Create the output directory
    os.makedirs(args.output_dir, exist_ok=True)

330
    config = model_config(args.config_preset)
331
332
333
334
335
336
337
    
    if(args.trace_model):
        if(not config.data.predict.fixed_size):
            raise ValueError(
                "Tracing requires that fixed_size mode be enabled in the config"
            )
    
338
339
340
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
341
        max_hits=config.data.predict.max_templates,
342
        kalign_binary_path=args.kalign_binary_path,
343
        release_dates_path=args.release_dates_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
344
345
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
346

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
347
    data_processor = data_pipeline.DataPipeline(
348
349
350
351
        template_featurizer=template_featurizer,
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
352
    random_seed = args.data_random_seed
353
    if random_seed is None:
354
355
356
357
358
        random_seed = random.randrange(2**32)
    
    np.random.seed(random_seed)
    torch.manual_seed(random_seed + 1)
    
359
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
360
361
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
362
    if args.use_precomputed_alignments is None:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
363
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
364
365
    else:
        alignment_dir = args.use_precomputed_alignments
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
366

367
368
    tag_list = []
    seq_list = []
369
    for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
370
        # Gather input sequences
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
371
        with open(os.path.join(args.fasta_dir, fasta_file), "r") as fp:
372
373
374
375
            data = fp.read()
    
        tags, seqs = parse_fasta(data)
        # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
376
        tag = '-'.join(tags)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
377

378
379
380
381
382
383
384
385
386
387
388
389
        tag_list.append(tag)
        seq_list.append(seqs)

    seq_sort_fn = lambda target: sum([len(s) for s in target[1]])
    sorted_targets = sorted(zip(tag_list, seq_list), key=seq_sort_fn)
    feature_dicts = {}
    for model, output_directory in load_models_from_command_line(args, config): 
        cur_tracing_interval = 0
        for tag, seqs in sorted_targets:
            output_name = f'{tag}_{args.config_preset}'
            if args.output_postfix is not None:
                output_name = f'{output_name}_{args.output_postfix}'
390
    
391
392
393
394
395
396
397
398
399
400
401
402
            # Does nothing if the alignments have already been computed
            precompute_alignments(tags, seqs, alignment_dir, args)
        
            feature_dict = feature_dicts.get(tag, None)
            if(feature_dict is None):
                feature_dict = generate_feature_dict(
                    tags,
                    seqs,
                    alignment_dir,
                    data_processor,
                    args,
                )
Sam DeLuca's avatar
Sam DeLuca committed
403

404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
                if(args.trace_model):
                    n = feature_dict["aatype"].shape[-2]
                    rounded_seqlen = round_up_seqlen(n)
                    feature_dict = pad_feature_dict_seq(
                        feature_dict, rounded_seqlen,
                    )

                feature_dicts[tag] = feature_dict

            processed_feature_dict = feature_processor.process_features(
                feature_dict, mode='predict',
            )

            processed_feature_dict = {
                k:torch.as_tensor(v, device=args.model_device) 
                for k,v in processed_feature_dict.items()
            }

            if(args.trace_model):
                if(rounded_seqlen > cur_tracing_interval):
                    logger.info(
                        f"Tracing model at {rounded_seqlen} residues..."
                    )
                    t = time.perf_counter()
                    trace_model_(model, processed_feature_dict)
                    logger.info(
                        f"Tracing time: {time.perf_counter() - t}"
                    )
                    cur_tracing_interval = rounded_seqlen
Sam DeLuca's avatar
Sam DeLuca committed
433

434
            out = run_model(model, processed_feature_dict, tag, args)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
435

436
            # Toss out the recycling dimensions --- we don't need them anymore
437
438
439
440
            processed_feature_dict = tensor_tree_map(
                lambda x: np.array(x[..., -1].cpu()), 
                processed_feature_dict
            )
441
442
443
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            unrelaxed_protein = prep_output(
444
445
446
447
448
                out, 
                processed_feature_dict, 
                feature_dict, 
                feature_processor, 
                args
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
449
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
450

451
452
453
454
            unrelaxed_output_path = os.path.join(
                output_directory, f'{output_name}_unrelaxed.pdb'
            )

455
456
            with open(unrelaxed_output_path, 'w') as fp:
                fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
457

458
            logger.info(f"Output written to {unrelaxed_output_path}...")
459
            
460
461
462
463
464
            if not args.skip_relaxation:
                amber_relaxer = relax.AmberRelaxation(
                    use_gpu=(args.model_device != "cpu"),
                    **config.relax,
                )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
465

466
                # Relax the prediction.
467
                logger.info(f"Running relaxation on {unrelaxed_output_path}...")
468
469
470
471
472
473
474
                t = time.perf_counter()
                visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", default="")
                if "cuda" in args.model_device:
                    device_no = args.model_device.split(":")[-1]
                    os.environ["CUDA_VISIBLE_DEVICES"] = device_no
                relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
                os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
475
                logger.info(f"Relaxation time: {time.perf_counter() - t}")
476
477
478

                # Save the relaxed PDB.
                relaxed_output_path = os.path.join(
479
                    output_directory, f'{output_name}_relaxed.pdb'
480
481
482
                )
                with open(relaxed_output_path, 'w') as fp:
                    fp.write(relaxed_pdb_str)
483
                
484
                logger.info(f"Relaxed output written to {relaxed_output_path}...")
485

486
487
            if args.save_outputs:
                output_dict_path = os.path.join(
488
                    output_directory, f'{output_name}_output_dict.pkl'
489
490
491
                )
                with open(output_dict_path, "wb") as fp:
                    pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
492

Sam DeLuca's avatar
Sam DeLuca committed
493
                logger.info(f"Model output written to {output_dict_path}...")
494

495

496
497
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
498
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
499
500
        "fasta_dir", type=str,
        help="Path to directory containing FASTA files, one sequence per file"
501
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
502
503
504
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
505
506
507
508
    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."""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
509
    )
510
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
511
512
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
513
514
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
515
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
516
517
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
518
519
    )
    parser.add_argument(
520
        "--config_preset", type=str, default="model_1",
521
        help="""Name of a model config preset defined in openfold/config.py"""
522
523
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
524
525
526
527
528
529
530
531
532
        "--jax_param_path", type=str, default=None,
        help="""Path to JAX model parameters. If None, and openfold_checkpoint_path
             is also None, parameters are selected automatically according to 
             the model name from openfold/resources/params"""
    )
    parser.add_argument(
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
533
    )
534
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
535
        "--save_outputs", action="store_true", default=False,
536
537
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
538
539
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
540
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
541
    )
542
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
543
        "--preset", type=str, default='full_dbs',
544
545
        choices=('reduced_dbs', 'full_dbs')
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
546
547
548
549
    parser.add_argument(
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
550
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
551
552
553
554
        "--data_random_seed", type=str, default=None
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
555
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
556
557
558
559
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
560
561
562
563
564
565
    parser.add_argument(
        "--trace_model", action="store_true", default=False,
        help="""Whether to convert parts of each model to TorchScript.
                Significantly improves runtime at the cost of lengthy
                'compilation.' Useful for large batch jobs."""
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
566
    add_data_args(parser)
567
568
    args = parser.parse_args()

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
569
570
    if(args.jax_param_path is None and args.openfold_checkpoint_path is None):
        args.jax_param_path = os.path.join(
571
            "openfold", "resources", "params", 
572
            "params_" + args.config_preset + ".npz"
573
574
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
575
576
577
578
579
580
    if(args.model_device == "cpu" and torch.cuda.is_available()):
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )

581
    main(args)