add override flag

PaddlePaddle · Aug 17, 2024 · d528300 · d528300
1 parent d2aa1b9
commit d528300
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml b/apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml
@@ -18,6 +18,7 @@ use_small_bfd: false # Corresponds to --use_small_bfd
 
 input: null  # Corresponds to --input_json, required field
 output: null  # Corresponds to --output_dir, required field
+override: false # Set true to override existing msa output directory
 
 
 # Binary tool paths, leave them as null to find proper ones under PATH or conda bin path
@@ -56,4 +57,3 @@ preset:
 # Other configurations
 other:
   maxit_binary: /mnt/data/yinying/software/maxit/maxit-v11.100-prod-src/bin/maxit  # Corresponds to --maxit_binary
-  no_msa_templ_feats: false  # Corresponds to --no_msa_templ_feats
diff --git a/apps/protein_folding/helixfold3/helixfold/inference.py b/apps/protein_folding/helixfold3/helixfold/inference.py
@@ -487,7 +487,7 @@ def main(cfg: DictConfig):
         assert cfg.db.uniclust30 is not None
 
     logger.info('Getting MSA/Template Pipelines...')
-    msa_templ_data_pipeline_dict = get_msa_templates_pipeline(args)
+    msa_templ_data_pipeline_dict = get_msa_templates_pipeline(cfg=cfg)
 
     ### Create model
     model_config = config.model_config(cfg.job_id)
@@ -515,8 +515,9 @@ def main(cfg: DictConfig):
     msa_output_dir.mkdir(parents=True, exist_ok=True)
 
     features_pkl = output_dir_base.joinpath('final_features.pkl')
-    if features_pkl.exists():
+    if features_pkl.exists() and not cfg.override:
         with open(features_pkl, 'rb') as f:
+            logging.info(f'Load features from precomputed {features_pkl}')
             feature_dict = pickle.load(f)
     else:
         feature_dict = feature_processing_aa.process_input_json(