Add option to select chains in PDB files

openfold/data/data_pipeline.py

@@ -472,7 +472,8 @@ class DataPipeline:
         self,
         pdb_path: str,
         alignment_dir: str,
-        is_distillation: bool = True
+        is_distillation: bool = True,
+        chain_id: Optional[str] = None,
     ) -> FeatureDict:
         """
             Assembles features for a protein in a PDB file.
@@ -480,7 +481,7 @@ class DataPipeline:
         with open(pdb_path, 'r') as f:
             pdb_str = f.read()
 
-        protein_object = protein.from_pdb_string(pdb_str)
+        protein_object = protein.from_pdb_string(pdb_str, chain_id)
         input_sequence = _aatype_to_str_sequence(protein_object.aatype) 
         description = os.path.splitext(os.path.basename(pdb_path))[0].upper()
         pdb_feats = make_pdb_features(