fix: output logs

PaddlePaddle · Aug 19, 2024 · dce79f8 · dce79f8
1 parent 330e455
commit dce79f8
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/apps/protein_folding/helixfold3/data/demo_3fap_protein_sm.json b/apps/protein_folding/helixfold3/data/demo_3fap_protein_sm.json
@@ -0,0 +1,19 @@
+{
+    "entities": [
+        {
+            "type": "protein",
+            "sequence": "GVQVETISPGDGRTFPKRGQTCVVHYTGMLEDGKKFDSSRDRNKPFKFMLGKQEVIRGWEEGVAQMSVGQRAKLTISPDYAYGATGHPGIIPPHATLVFDVELLKLE",
+            "count": 1
+        },
+        {
+            "type": "protein",
+            "sequence": "VAILWHEMWHEGLEEASRLYFGERNVKGMFEVLEPLHAMMERGPQTLKETSFNQAYGRDLMEAQEWCRKYMKSGNVKDLTQAWDLYYHVFRRIS",
+            "count": 1
+        },
+        {
+            "type": "ligand",
+            "smiles": "Cc1ccc(s1)[C@@H]\\2C[C@@H]3CC[C@H]([C@@](O3)(C(=O)C(=O)N4CCCC[C@H]4C(=O)O[C@@H](CC(=O)[C@@H](/C=C(/[C@H]([C@H](C(=O)[C@@H](C[C@@H](/C=C/C=C/C=C2\\C)C)C)OC)O)\\C)C)[C@H](C)C[C@@H]5CC[C@H]([C@@H](C5)OC)O)O)C",
+            "count": 1
+        }
+    ]
+}
diff --git a/apps/protein_folding/helixfold3/helixfold/infer_scripts/preprocess.py b/apps/protein_folding/helixfold3/helixfold/infer_scripts/preprocess.py
@@ -22,6 +22,7 @@
 from rdkit import Chem
 from rdkit.Chem import AllChem
 from helixfold.common import residue_constants
+from helixfold.data.tools import utils
 
 
 ## NOTE: this mapping is only useful for standard dna/rna/protein sequence input.
@@ -165,18 +166,19 @@ def _get_supported_formats(self) -> Tuple[str]:
 
     def _perform_conversion(self, input_type: str, input_value: str) -> Chem.Mol:
         with tempfile.NamedTemporaryFile(suffix=".mol2") as temp_file:
-            print(f"[OBABEL] Temporary file created: {temp_file.name}")
-            if input_type == 'smiles':
-                obabel_cmd = f"{self.obabel_bin} -:'{input_value}' -omol2 -O{temp_file.name} --gen3d"
-            else:
-                obabel_cmd = f"{self.obabel_bin} -i {input_type} {input_value} -omol2 -O{temp_file.name} --gen3d"
-            ret = subprocess.run(obabel_cmd, shell=True, capture_output=True, text=True)
-            mol = Chem.MolFromMol2File(temp_file.name, sanitize=False)
-            if '3D coordinate generation failed' in ret.stderr:
-                mol = generate_ETKDGv3_conformer(mol)
-            optimal_mol_wo_H = Chem.RemoveAllHs(mol, sanitize=False)
-            logging.debug(f'Converted `{input_type}`: {input_value}')
-            return optimal_mol_wo_H
+            with utils.timing(f'converting {input_type} to mol2: {input_value}'):
+                if input_type == 'smiles':
+                    obabel_cmd = f"{self.obabel_bin} -:'{input_value}' -omol2 -O{temp_file.name} --gen3d"
+                    if len(input_value)>60:
+                        logging.warning(f'This takes a while ...')
+                else:
+                    obabel_cmd = f"{self.obabel_bin} -i {input_type} {input_value} -omol2 -O{temp_file.name} --gen3d"
+                ret = subprocess.run(obabel_cmd, shell=True, capture_output=True, text=True)
+                mol = Chem.MolFromMol2File(temp_file.name, sanitize=False)
+                if '3D coordinate generation failed' in ret.stderr:
+                    mol = generate_ETKDGv3_conformer(mol)
+                optimal_mol_wo_H = Chem.RemoveAllHs(mol, sanitize=False)
+                return optimal_mol_wo_H
 
     def _convert_to_mol(self, input_type: str, input_value: str) -> Chem.Mol:
         if input_type not in self.supported_formats: