fix molecule sanitization in reader.py

chebai_graph/preprocessing/fg_detection/fg_aware_rule_based.py

@@ -8,7 +8,7 @@
 from rdkit.Chem import AllChem
 from rdkit.Chem import MolToSmiles as m2s
 
-from chebi_utils.sdf_extractor import _sanitize_molecule
+from chebi_utils.read_molecule import smiles_or_inchi_to_mol
 
 from .fg_constants import ELEMENTS, FLAG_NO_FG
 
@@ -1913,11 +1913,7 @@ def get_structure(mol):
         structure[frag] = {"atom": atom_idx, "is_ring_fg": False}
 
         # Convert fragment SMILES back to mol to match with fused ring atom indices
-        frag_mol = Chem.MolFromSmiles(frag, sanitize=False)
-        try:
-            frag_mol = _sanitize_molecule(frag_mol)
-        except Exception:
-            pass
+        frag_mol = smiles_or_inchi_to_mol(frag)
         frag_rings = frag_mol.GetRingInfo().AtomRings()
         if len(frag_rings) >= 1:
             structure[frag]["is_ring_fg"] = True

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -4,7 +4,7 @@
 
 import torch
 from chebai.preprocessing.reader import DataReader
-from chebi_utils.sdf_extractor import _sanitize_molecule
+from chebi_utils.read_molecule import smiles_or_inchi_to_mol
 from rdkit import Chem
 from torch_geometric.data import Data as GeomData
 
@@ -75,7 +75,7 @@ def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, dict] | None:
             RuntimeError: If an unexpected error occurs during graph augmentation.
         """
         if isinstance(raw_data, str):
-            mol = self._smiles_to_mol(raw_data)
+            mol = smiles_or_inchi_to_mol(raw_data)
             smiles = raw_data
         else:
             mol = raw_data
@@ -139,29 +139,6 @@ def _read_data(self, raw_data: str | Chem.Mol) -> tuple[GeomData, dict] | None:
             augmented_molecule,
         )
 
-    def _smiles_to_mol(self, smiles: str) -> Chem.Mol | None:
-        """
-        Converts a SMILES string to an RDKit molecule object. Sanitizes the molecule.
-
-        Args:
-            smiles (str): SMILES string representing the molecule.
-
-        Returns:
-            Chem.Mol | None: RDKit molecule object if successful, else None.
-        """
-        mol = Chem.MolFromSmiles(smiles, sanitize=False)
-        if mol is None:
-            print(f"RDKit failed to parse {smiles} (returned None)")
-            self.f_cnt_for_smiles += 1
-        else:
-            try:
-                mol = _sanitize_molecule(mol)
-            except Exception as e:
-                print(f"RDKit failed at sanitizing {smiles}, Error {e}")
-                self.f_cnt_for_smiles += 1
-                mol = None
-        return mol
-
     def _create_augmented_graph(
         self, mol: Chem.Mol
     ) -> tuple[torch.Tensor, dict] | None:
@@ -315,7 +292,7 @@ def read_property(
                 smiles = raw_data
                 if smiles in self.mol_object_buffer:
                     return property.get_property_value(self.mol_object_buffer[smiles])
-                mol = self._smiles_to_mol(smiles)
+                mol = smiles_or_inchi_to_mol(smiles)
             if mol is None:
                 return None
             try:

chebai_graph/preprocessing/reader/reader.py

@@ -1,7 +1,7 @@
 import os
 
 import chebai.preprocessing.reader as dr
-from chebi_utils.sdf_extractor import _sanitize_molecule
+from chebi_utils.read_molecule import smiles_or_inchi_to_mol
 import networkx as nx
 import rdkit.Chem as Chem
 import torch
@@ -29,7 +29,7 @@ def __init__(
         """
         super().__init__(*args, **kwargs)
         self.failed_counter = 0
-        self.mol_object_buffer: dict[str, Chem.rdchem.Mol | None] = {}
+        self.mol_object_buffer: dict[str, Chem.Mol | None] = {}
 
     @classmethod
     def name(cls) -> str:
@@ -41,29 +41,20 @@ def name(cls) -> str:
         """
         return "graph_properties"
 
-    def _smiles_to_mol(self, smiles: str) -> Chem.rdchem.Mol | None:
+    def _smiles_to_mol(self, smiles: str) -> Chem.Mol | None:
         """
         Load SMILES string into an RDKit molecule object and cache it.
 
         Args:
             smiles (str): The SMILES string to parse.
 
         Returns:
-            Chem.rdchem.Mol | None: Parsed molecule object or None if parsing failed.
+            Chem.Mol | None: Parsed molecule object or None if parsing failed.
         """
         if smiles in self.mol_object_buffer:
             return self.mol_object_buffer[smiles]
 
-        mol = Chem.MolFromSmiles(smiles, sanitize=False)
-        if mol is None:
-            print(f"RDKit failed to at parsing {smiles} (returned None)")
-            self.failed_counter += 1
-        else:
-            try:
-                _sanitize_molecule(mol)
-            except Exception as e:
-                print(f"Rdkit failed at sanitizing {smiles}, \n Error: {e}")
-                self.failed_counter += 1
+        mol = smiles_or_inchi_to_mol(smiles)
         self.mol_object_buffer[smiles] = mol
         return mol
 
@@ -162,11 +153,12 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         # raw_data is a SMILES string
         try:
             mol = (
-                self._smiles_to_mol(raw_data) if isinstance(raw_data, str) else raw_data
+                smiles_or_inchi_to_mol(raw_data)
+                if isinstance(raw_data, str)
+                else raw_data
             )
         except ValueError:
             return None
-        assert isinstance(mol, nx.Graph)
         d: dict[int, int] = {}
         de: dict[tuple[int, int], int] = {}
         for node in mol.nodes:
@@ -197,27 +189,6 @@ def _read_data(self, raw_data: str | Chem.Mol) -> GeomData | None:
         data = from_networkx(mol)
         return data
 
-    def _smiles_to_mol(self, smiles: str) -> Chem.rdchem.Mol | None:
-        """
-        Load SMILES string into an RDKit molecule object.
-
-        Args:
-            smiles (str): The SMILES string to parse.
-
-        Returns:
-            Chem.rdchem.Mol | None: Parsed molecule object or None if parsing failed.
-        """
-
-        mol = Chem.MolFromSmiles(smiles, sanitize=False)
-        if mol is None:
-            print(f"RDKit failed to at parsing {smiles} (returned None)")
-        else:
-            try:
-                _sanitize_molecule(mol)
-            except Exception as e:
-                print(f"Rdkit failed at sanitizing {smiles}, \n Error: {e}")
-        return mol
-
     def collate(self, list_of_tuples: list) -> any:
         """
         Collate a list of samples into a batch.

pyproject.toml

@@ -8,6 +8,7 @@ authors = [
 ]
 dependencies = [
     "chebai",
+    "chebi_utils>=0.4",
     "descriptastorus",
     # below packages need to manually installed as mentioned in readme
     # torch-geometric