From 4fc8977c94193a99f44196d9fa2e802472e408b6 Mon Sep 17 00:00:00 2001 From: zhuww Date: Thu, 23 Feb 2023 09:21:40 +0800 Subject: [PATCH 1/2] save prediction output pickle file of monomer --- inference.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/inference.py b/inference.py index 2978478c..9fa8f24e 100644 --- a/inference.py +++ b/inference.py @@ -478,6 +478,14 @@ def inference_monomer_model(args): f'{tag}_{args.model_name}_relaxed.pdb') with open(relaxed_output_path, 'w') as f: f.write(relaxed_pdb_str) + + if(args.save_prediction_result): + # Save the prediction result .pkl + prediction_result_path = os.path.join( + args.output_dir, f'{tag}_{args.model_name}.pkl' + ) + with open(prediction_result_path, "wb") as fp: + pickle.dump(out, fp) if __name__ == "__main__": From 5f08e3f3b161ba83d8155ff3a82659b2e1893ff5 Mon Sep 17 00:00:00 2001 From: zhuwenwen <1536492181@qq.com> Date: Sat, 2 Dec 2023 16:48:35 +0800 Subject: [PATCH 2/2] save features.pkl of the monomer --- inference.py | 37 +++++++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/inference.py b/inference.py index 9fa8f24e..f6edde21 100644 --- a/inference.py +++ b/inference.py @@ -341,15 +341,6 @@ def inference_monomer_model(args): print("running in monomer mode...") config = model_config(args.model_name) - 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, - obsolete_pdbs_path=args.obsolete_pdbs_path - ) - 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 @@ -357,8 +348,6 @@ def inference_monomer_model(args): assert args.bfd_database_path is not None assert args.uniref30_database_path is not None - data_processor = data_pipeline.DataPipeline(template_featurizer=template_featurizer,) - output_dir_base = args.output_dir random_seed = args.data_random_seed @@ -423,11 +412,31 @@ def inference_monomer_model(args): use_small_bfd=use_small_bfd, no_cpus=args.cpus, ) + t = time.perf_counter() alignment_runner.run(fasta_path, local_alignment_dir) + print(f"Alignment data time: {time.perf_counter() - t}") - feature_dict = data_processor.process_fasta(fasta_path=fasta_path, - alignment_dir=local_alignment_dir) - + features_output_path = os.path.join(local_alignment_dir, 'features.pkl') + if os.path.exists(features_output_path): + feature_dict = pickle.load(open(features_output_path, 'rb')) + + else: + 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, + obsolete_pdbs_path=args.obsolete_pdbs_path + ) + + data_processor = data_pipeline.DataPipeline(template_featurizer=template_featurizer,) + + feature_dict = data_processor.process_fasta(fasta_path=fasta_path, + alignment_dir=local_alignment_dir) + with open(features_output_path, 'wb') as f: + pickle.dump(feature_dict, f, protocol=4) + # Remove temporary FASTA file os.remove(fasta_path)