Skip to content

Commit

Permalink
Fix FASTA parsing bug in inference script, add option to save outputs
Browse files Browse the repository at this point in the history
  • Loading branch information
gahdritz committed Mar 29, 2022
1 parent 9f55817 commit 16e0082
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions run_pretrained_openfold.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,10 @@ def main(args):

# Gather input sequences
with open(args.fasta_path, "r") as fp:
lines = [l.strip() for l in fp.readlines()]
data = fp.read()

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

for tag, seq in zip(tags, seqs):
fasta_path = os.path.join(args.output_dir, "tmp.fasta")
Expand Down Expand Up @@ -179,6 +179,13 @@ def main(args):
with open(relaxed_output_path, 'w') as f:
f.write(relaxed_pdb_str)

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)


if __name__ == "__main__":
parser = argparse.ArgumentParser()
Expand Down Expand Up @@ -213,6 +220,10 @@ def main(args):
automatically according to the model name from
openfold/resources/params"""
)
parser.add_argument(
"--save_outputs", type=bool, default=False,
help="Whether to save all model outputs, including embeddings, etc."
)
parser.add_argument(
"--cpus", type=int, default=4,
help="""Number of CPUs with which to run alignment tools"""
Expand Down

0 comments on commit 16e0082

Please sign in to comment.