From 3257aa18d7888b73618e04f24eddbba63b1033ab Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 11:56:12 +0200 Subject: [PATCH 1/3] accept InChI or mol objects for prediction, use chebi_utils SMILES parsing --- chebai/preprocessing/datasets/base.py | 21 +++++++------- chebai/preprocessing/reader.py | 24 ++++++++++------ chebai/result/prediction.py | 41 ++++++++++++++------------- 3 files changed, 48 insertions(+), 38 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index d5cd5653..b3110783 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/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,7 +408,7 @@ 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]]: @@ -415,7 +416,7 @@ def predict_dataloader( 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, } ) diff --git a/chebai/preprocessing/reader.py b/chebai/preprocessing/reader.py index 0dae39cd..664a8d8f 100644 --- a/chebai/preprocessing/reader.py +++ b/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}") diff --git a/chebai/result/prediction.py b/chebai/result/prediction.py index c80cdc34..543eeb2b 100644 --- a/chebai/result/prediction.py +++ b/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) From e064abaa4fded0c1381b418bcbc6bb8d12e12245 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 15:15:00 +0200 Subject: [PATCH 2/3] update chebi utils dependency --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6c00552f..f0594fc8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ dev = [ "omegaconf", "deepsmiles", "torchmetrics", - "chebi-utils>=0.3", + "chebi-utils>=0.4", ] linters = [ From 0063b1e841c552dcbdc40874c11fc974151d4304 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:21:38 +0000 Subject: [PATCH 3/3] fix(ci): install chebi-utils in verify constants workflow --- .github/workflows/verify_constants.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/verify_constants.yml b/.github/workflows/verify_constants.yml index 3246f64d..35e46ff5 100644 --- a/.github/workflows/verify_constants.yml +++ b/.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