Merge pull request #138 from ChEB-AI/fix/read_data

Raise error for invalid smiles and return None

Author: sfluegel05

chebai/preprocessing/reader.py

@@ -199,23 +199,26 @@ def _read_data(self, raw_data: str) -> List[int]:
         Returns:
             List[int]: A list of integers representing the indices of the SMILES tokens.
         """
-        if self.canonicalize_smiles:
-            try:
-                mol = Chem.MolFromSmiles(raw_data.strip())
-                if mol is not None:
-                    raw_data = Chem.MolToSmiles(mol, canonical=True)
-            except Exception as e:
-                print(f"RDKit failed to process {raw_data}")
-                print(f"\t{e}")
         try:
-            return [self._get_token_index(v[1]) for v in _tokenize(raw_data)]
+            mol = Chem.MolFromSmiles(raw_data.strip())
+            if mol is None:
+                raise ValueError(f"Invalid SMILES: {raw_data}")
         except ValueError as e:
             print(f"could not process {raw_data}")
-            print(f"\t{e}")
+            print(f"\tError: {e}")
             return None
 
-    def _back_to_smiles(self, smiles_encoded):
+        if self.canonicalize_smiles:
+            try:
+                raw_data = Chem.MolToSmiles(mol, canonical=True)
+            except Exception as e:
+                print(f"RDKit failed to canonicalize the SMILES: {raw_data}")
+                print(f"\t{e}")
+                return None
+
+        return [self._get_token_index(v[1]) for v in _tokenize(raw_data)]
 
+    def _back_to_smiles(self, smiles_encoded):
         token_file = self.reader.token_path
         token_coding = {}
         counter = 0

pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "torch",
     "transformers",
     "pysmiles==1.1.2",
-    "rdkit",
+    "rdkit==2024.3.6",
     "lightning==2.5.1",
 ]
 

tests/unit/readers/testChemDataReader.py

@@ -42,19 +42,22 @@ def test_read_data(self) -> None:
         """
         Test the _read_data method with a SMILES string to ensure it correctly tokenizes the string.
         """
-        raw_data = "CC(=O)NC1[Mg-2]"
+        raw_data = "CC(=O)NC1CC1[Mg-2]"
         # Expected output as per the tokens already in the cache, and ")" getting added to it.
         expected_output: List[int] = [
             EMBEDDING_OFFSET + 0,  # C
             EMBEDDING_OFFSET + 0,  # C
-            EMBEDDING_OFFSET + 5,  # =
-            EMBEDDING_OFFSET + 3,  # O
-            EMBEDDING_OFFSET + 1,  # N
-            EMBEDDING_OFFSET + len(self.reader.cache),  # (
-            EMBEDDING_OFFSET + 2,  # C
+            EMBEDDING_OFFSET + 5,  # (
+            EMBEDDING_OFFSET + 3,  # =
+            EMBEDDING_OFFSET + 1,  # O
+            EMBEDDING_OFFSET + len(self.reader.cache),  # ) - new token
+            EMBEDDING_OFFSET + 2,  # N
             EMBEDDING_OFFSET + 0,  # C
             EMBEDDING_OFFSET + 4,  # 1
-            EMBEDDING_OFFSET + len(self.reader.cache) + 1,  # [Mg-2]
+            EMBEDDING_OFFSET + 0,  # C
+            EMBEDDING_OFFSET + 0,  # C
+            EMBEDDING_OFFSET + 4,  # 1
+            EMBEDDING_OFFSET + len(self.reader.cache) + 1,  # [Mg-2] - new token
         ]
         result = self.reader._read_data(raw_data)
         self.assertEqual(
@@ -99,13 +102,29 @@ def test_read_data_with_invalid_input(self) -> None:
         Test the _read_data method with an invalid input.
         The invalid token should prompt a return value None
         """
-        raw_data = "%INVALID%"
-
-        result = self.reader._read_data(raw_data)
-        self.assertIsNone(
-            result,
-            "The output for invalid token '%INVALID%' should be None.",
-        )
+        # see https://github.com/ChEB-AI/python-chebai/issues/137
+        raw_datas = ["%INVALID%", "ADADAD", "ADASDAD", "CC(=O)NC1[Mg-2]"]
+        for raw_data in raw_datas:
+            result = self.reader._read_data(raw_data)
+            self.assertIsNone(
+                result,
+                f"The output for invalid token '{raw_data}' should be None.",
+            )
+
+    def test_read_data_with_invalid_input_with_no_canonicalize(self) -> None:
+        """
+        Test the _read_data method with an invalid input.
+        The invalid token should prompt a return value None
+        """
+        self.reader.canonicalize_smiles = False
+        raw_datas = ["%INVALID%", "ADADAD", "ADASDAD", "CC(=O)NC1[Mg-2]"]
+        for raw_data in raw_datas:
+            result = self.reader._read_data(raw_data)
+            self.assertIsNone(
+                result,
+                f"The output for invalid token '{raw_data}' should be None.",
+            )
+        self.reader.canonicalize_smiles = True  # Reset to original state
 
     @patch("builtins.open", new_callable=mock_open)
     def test_finish_method_for_new_tokens(self, mock_file: mock_open) -> None: