Use chebi_utils SMILES / InCHI parsing

.github/workflows/verify_constants.yml

@@ -63,6 +63,7 @@ jobs:
           python -m pip install --upgrade pip setuptools wheel
           python -m pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu
           python -m pip install -e .
+          python -m pip install "chebi-utils>=0.4"
 
       - name: Export constants
         run: python .github/workflows/export_constants.py

chebai/preprocessing/datasets/base.py

@@ -7,6 +7,7 @@
 import lightning as pl
 import numpy as np
 import pandas as pd
+from rdkit import Chem
 import torch
 import tqdm
 from lightning.pytorch.core.datamodule import LightningDataModule
@@ -407,15 +408,15 @@ def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]
 
     def predict_dataloader(
         self,
-        smiles_list: List[str],
+        molecule_list: List[str | Chem.Mol],
         model_hparams: dict,
         **kwargs,
     ) -> tuple[DataLoader, list[int]]:
         """
         Returns the predict DataLoader.
 
         Args:
-            smiles_list (List[str]): List of SMILES strings to predict.
+            molecule_list (List[str|Chem.Mol]): List of molecules (SMILES / InChI strings or RDKit molecule objects) to predict.
             model_hparams (Optional[dict]): Model hyperparameters.
                 Some prediction pre-processing pipelines may require these.
             **kwargs: Additional keyword arguments, passed to dataloader().
@@ -425,7 +426,7 @@ def predict_dataloader(
         """
 
         data, valid_indices = self._process_input_for_prediction(
-            smiles_list, model_hparams
+            molecule_list, model_hparams
         )
         return (
             DataLoader(
@@ -438,13 +439,13 @@ def predict_dataloader(
         )
 
     def _process_input_for_prediction(
-        self, smiles_list: list[str], model_hparams: dict
+        self, molecule_list: list[str | Chem.Mol], model_hparams: dict
     ) -> tuple[list, list]:
         """
         Process input data for prediction.
 
         Args:
-            smiles_list (List[str]): List of SMILES strings.
+            molecule_list (List[str|Chem.Mol]): List of molecules (SMILES / InChI strings or RDKit molecule objects) to predict.
             model_hparams (dict): Model hyperparameters.
                 Some prediction pre-processing pipelines may require these.
 
@@ -455,8 +456,8 @@ def _process_input_for_prediction(
         num_of_labels = int(model_hparams["out_dim"])
         self._dummy_labels: list = list(range(1, num_of_labels + 1))
 
-        for idx, smiles in enumerate(smiles_list):
-            result = self._preprocess_smiles_for_pred(idx, smiles, model_hparams)
+        for idx, molecule in enumerate(molecule_list):
+            result = self._preprocess_molecule_for_pred(idx, molecule, model_hparams)
             if result is None or result["features"] is None:
                 continue
             if not self._filter_to_token_limit(result):
@@ -466,8 +467,8 @@ def _process_input_for_prediction(
 
         return data, valid_indices
 
-    def _preprocess_smiles_for_pred(
-        self, idx: int, smiles: str, model_hparams: Optional[dict] = None
+    def _preprocess_molecule_for_pred(
+        self, idx: int, molecule: str | Chem.Mol, model_hparams: Optional[dict] = None
     ) -> dict:
         """Preprocess prediction data."""
         # Add dummy labels because the collate function requires them.
@@ -476,7 +477,7 @@ def _preprocess_smiles_for_pred(
         return self.reader.to_data(
             {
                 "id": f"smiles_{idx}",
-                "features": smiles,
+                "features": molecule,
                 "labels": self._dummy_labels,
             }
         )

chebai/preprocessing/reader.py

@@ -5,6 +5,7 @@
 from itertools import islice
 from typing import Any, Dict, List, Optional
 
+from chebi_utils.read_molecule import smiles_or_inchi_to_mol
 from pysmiles.read_smiles import _tokenize
 from rdkit import Chem
 
@@ -194,17 +195,18 @@ def name(cls) -> str:
 
     def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]:
         """
-        Reads and tokenizes SMILES strings (or SMILES strings generated from Chem.Mol objects) into a list of token indices. Optionally canonicalizes the SMILES string using RDKit.
+        Reads and tokenizes SMILES strings (or SMILES strings generated from Chem.Mol objects / InChI strings) into a list of token indices.
+        Optionally canonicalizes the SMILES string using RDKit (if the input is a mol object or InChI, the SMILES will always be canonicalized).
 
         Args:
-            raw_data (str|Chem.Mol): The raw SMILES string or Chem.Mol object to be tokenized.
+            raw_data (str|Chem.Mol): The raw SMILES / InChI string or Chem.Mol object to be tokenized.
 
         Returns:
             List[int]: A list of integers representing the indices of the SMILES tokens.
         """
         try:
             if isinstance(raw_data, str):
-                mol = Chem.MolFromSmiles(raw_data.strip())
+                mol = smiles_or_inchi_to_mol(raw_data.strip())
             else:
                 mol = raw_data
             if mol is None:
@@ -221,15 +223,17 @@ def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]:
                 print(f"RDKit failed to canonicalize the SMILES: {raw_data}")
                 print(f"\t{e}")
                 return None
-        elif not isinstance(raw_data, str):
+        elif isinstance(raw_data, str) and not raw_data.startswith("InChI="):
+            # only a raw SMILES string can be tokenized as-is
+            smiles = raw_data.strip()
+        else:
+            # Chem.Mol input, or an InChI string that has to be serialized first
             try:
                 smiles = Chem.MolToSmiles(mol)
             except Exception as e:
-                print(f"RDKit failed to convert Mol object to SMILES: {raw_data}")
+                print(f"RDKit failed to convert input to SMILES: {raw_data}")
                 print(f"\t{e}")
                 return None
-        else:
-            smiles = raw_data
 
         try:
             tokenized = [self._get_token_index(v[1]) for v in _tokenize(smiles)]
@@ -277,12 +281,14 @@ def name(cls) -> str:
         return "static_smiles"
 
     def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]:
-        """Tokenize raw SMILES data using BasicSmilesTokenizer with static vocabulary."""
+        """Tokenize SMILES / InChI / Mol object using BasicSmilesTokenizer with static vocabulary."""
         try:
             if isinstance(raw_data, str):
-                mol = Chem.MolFromSmiles(raw_data.strip())
+                mol = smiles_or_inchi_to_mol(raw_data.strip())
             else:
                 mol = raw_data
+            if mol is None:
+                raise ValueError(f"Invalid input: {raw_data}")
         except ValueError as e:
             print(f"could not process {raw_data}")
             print(f"\tError: {e}")

chebai/result/prediction.py

@@ -1,6 +1,7 @@
 from typing import List, Optional
 
 import pandas as pd
+from rdkit import Chem
 import torch
 from jsonargparse import CLI
 from lightning.fabric.utilities.types import _PATH
@@ -104,20 +105,22 @@ def __init__(
 
     def predict_from_file(
         self,
-        smiles_file_path: _PATH,
+        file_path: _PATH,
         save_to: _PATH = "predictions.csv",
     ) -> None:
         """
         Loads a model from a checkpoint and makes predictions on input data from a file.
 
         Args:
-            smiles_file_path: Path to the input file containing SMILES strings.
+            file_path: Path to the input file containing SMILES / InChI strings.
             save_to: Path to save the predictions CSV file.
         """
-        with open(smiles_file_path, "r") as input:
-            smiles_strings = [inp.strip() for inp in input.readlines()]
+        with open(file_path, "r") as input:
+            input_strings = [inp.strip() for inp in input.readlines()]
 
-        preds: list[torch.Tensor | None] = self.predict_smiles(smiles=smiles_strings)
+        preds: list[torch.Tensor | None] = self.predict_molecules(
+            molecules=input_strings
+        )
         if all(pred is None for pred in preds):
             print("No valid predictions were made. (All predictions are None.)")
             return
@@ -128,32 +131,32 @@ def predict_from_file(
             for pred in preds
         ]
         predictions_df = pd.DataFrame(
-            rows, columns=self._classification_labels, index=smiles_strings
+            rows, columns=self._classification_labels, index=input_strings
         )
 
         predictions_df.to_csv(save_to)
         print(f"Predictions saved to: {save_to}")
 
     @torch.inference_mode()
-    def predict_smiles(
+    def predict_molecules(
         self,
-        smiles: List[str],
+        molecules: List[str | Chem.Mol],
     ) -> list[torch.Tensor | None]:
         """
-        Predicts the output for a list of SMILES strings using the model.
+        Predicts the output for a list of molecules using the model.
 
         Args:
-            smiles: A list of SMILES strings.
+            molecules: A list of SMILES / InChI strings or RDKit molecule objects.
 
         Returns:
             A tensor containing the predictions.
         """
         # For certain data prediction pipelines, we may need model hyperparameters
         pred_dl, valid_indices = self._dm.predict_dataloader(
-            smiles_list=smiles, model_hparams=self._model_hparams
+            molecule_list=molecules, model_hparams=self._model_hparams
         )
         if valid_indices is None or len(valid_indices) == 0:
-            return [None] * len(smiles)
+            return [None] * len(molecules)
 
         preds = []
         for batch_idx, batch in enumerate(pred_dl):
@@ -165,7 +168,7 @@ def predict_smiles(
         preds = torch.cat(preds)
 
         # Initialize output with None
-        output: list[torch.Tensor | None] = [None] * len(smiles)
+        output: list[torch.Tensor | None] = [None] * len(molecules)
 
         # Scatter predictions back
         for pred, idx in zip(preds, valid_indices):
@@ -178,27 +181,27 @@ class MainPredictor:
     @staticmethod
     def predict_from_file(
         checkpoint_path: _PATH,
-        smiles_file_path: _PATH,
+        file_path: _PATH,
         save_to: _PATH = "predictions.csv",
         batch_size: Optional[int] = None,
     ) -> None:
         predictor = Predictor(checkpoint_path, batch_size)
         predictor.predict_from_file(
-            smiles_file_path,
+            file_path,
             save_to,
         )
 
     @staticmethod
-    def predict_smiles(
+    def predict(
         checkpoint_path: _PATH,
-        smiles: List[str],
+        molecules: List[str | Chem.Mol],
         batch_size: Optional[int] = None,
     ) -> list[torch.Tensor | None]:
         predictor = Predictor(checkpoint_path, batch_size)
-        return predictor.predict_smiles(smiles)
+        return predictor.predict_molecules(molecules=molecules)
 
 
 if __name__ == "__main__":
     # python chebai/result/prediction.py  predict_from_file --help
-    # python chebai/result/prediction.py  predict_smiles --help
+    # python chebai/result/prediction.py  predict --help
     CLI(MainPredictor, as_positional=False)

pyproject.toml

@@ -40,7 +40,7 @@ dev = [
     "omegaconf",
     "deepsmiles",
     "torchmetrics",
-    "chebi-utils>=0.3",
+    "chebi-utils>=0.4",
 ]
 
 linters = [