From ce757d5a0f4a5c2381b91792d28e5a75ab5040d0 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 28 Jul 2026 18:09:49 +0200 Subject: [PATCH 01/10] separate ensemble from base learner predictions, generalise both ensemble and base learners for inclusion of new ensemble strategies that are not voting-based --- chebifier/__init__.py | 4 +- chebifier/build_ensemble.py | 61 ++++ chebifier/ensemble/base_ensemble.py | 315 ++---------------- chebifier/ensemble/voting_ensemble.py | 304 +++++++++++++++++ .../ensemble/weighted_majority_ensemble.py | 41 ++- chebifier/model_registry.py | 4 +- chebifier/predict.py | 1 + chebifier/prediction_models/base_predictor.py | 16 +- chebifier/prediction_models/c3p_predictor.py | 2 +- chebifier/prediction_models/chebi_lookup.py | 8 +- .../prediction_models/chemlog_predictor.py | 10 +- chebifier/prediction_models/nn_predictor.py | 2 +- pyproject.toml | 1 + 13 files changed, 447 insertions(+), 322 deletions(-) create mode 100644 chebifier/build_ensemble.py create mode 100644 chebifier/ensemble/voting_ensemble.py create mode 100644 chebifier/predict.py diff --git a/chebifier/__init__.py b/chebifier/__init__.py index a4f770c..34967ee 100644 --- a/chebifier/__init__.py +++ b/chebifier/__init__.py @@ -2,10 +2,10 @@ # even if multiple subpackages are imported later. from ._custom_cache import PerSmilesPerModelLRUCache, modelwise_smiles_lru_cache -from .ensemble.base_ensemble import BaseEnsemble +from .ensemble.voting_ensemble import VotingEnsemble __all__ = [ - "BaseEnsemble", + "VotingEnsemble", "PerSmilesPerModelLRUCache", "modelwise_smiles_lru_cache", ] diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py new file mode 100644 index 0000000..eaf2cba --- /dev/null +++ b/chebifier/build_ensemble.py @@ -0,0 +1,61 @@ +import os + +import torch + + +class EnsembleBuilder: + """ + A class to build an ensemble model from base learners and validation data. + + Attributes: + base_learners (dict[str, BasePredictor]): A dictionary of base learner models. + ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model. + validation_data (list[Chem.Mol]): Validation data for calibration. + validation_labels (torch.Tensor): Validation labels for calibration. + prediction_cache_dir (str): Directory to cache predictions. + """ + + def __init__( + self, + base_learners, + ensemble_model, + validation_data, + validation_labels, + prediction_cache_dir, + ): + self.base_learners = base_learners + self.ensemble_model = ensemble_model + self.validation_data = validation_data + self.validation_labels = validation_labels + self.prediction_cache_dir = prediction_cache_dir + + def build_ensemble(self): + """ + Build an ensemble model from base learners and validation data. + + Base learner predictions are cached to avoid recomputation. + """ + + # Step 1: Get predictions from base learners on validation data + validation_predictions = {} + # get cached predictions if available, otherwise compute and cache them + for model_name, model in self.base_learners.items(): + cache_path = os.path.join( + self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" + ) + if os.path.exists(cache_path): + validation_predictions[model_name] = torch.load( + cache_path, weights_only=False + ) + else: + validation_predictions[model_name] = model.predict_list( + self.validation_data + ) + torch.save(validation_predictions[model_name], cache_path) + + # Step 2: Calibrate the ensemble model using validation predictions + self.ensemble_model.calibrate( + validation_predictions, self.validation_data, self.validation_labels + ) + + return self.ensemble_model diff --git a/chebifier/ensemble/base_ensemble.py b/chebifier/ensemble/base_ensemble.py index 59601f4..67f9fa3 100644 --- a/chebifier/ensemble/base_ensemble.py +++ b/chebifier/ensemble/base_ensemble.py @@ -1,303 +1,40 @@ -import importlib -import time -from pathlib import Path -from typing import Union +import os import torch -import tqdm -import yaml - -from chebifier.check_env import check_package_installed -from chebifier.hugging_face import download_model_files -from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother -from chebifier.prediction_models.base_predictor import BasePredictor -from chebifier.utils import ( - get_default_configs, - get_disjoint_files, - load_chebi_graph, - process_config, -) +from rdkit import Chem class BaseEnsemble: - def __init__( - self, - model_configs: Union[str, Path, dict, None] = None, - resolve_inconsistencies: bool = True, - verbose_output: bool = False, - use_confidence: bool = True, - ): - # Deferred Import: To avoid circular import error - from chebifier.model_registry import MODEL_TYPES - - # Load configuration from YAML file - if not model_configs: - config = get_default_configs() - elif isinstance(model_configs, dict): - config = model_configs - else: - print(f"Loading ensemble configuration from {model_configs}") - with open(model_configs, "r") as f: - config = yaml.safe_load(f) - - with ( - importlib.resources.files("chebifier") - .joinpath("model_registry.yml") - .open("r") as f - ): - model_registry = yaml.safe_load(f) - - processed_configs = process_config(config, model_registry) - self.verbose_output = verbose_output - self.use_confidence = use_confidence - - self.chebi_graph = load_chebi_graph() - self.disjoint_files = get_disjoint_files() + """Base class for ensemble models. + Each ensemble has to perform the following tasks: + 1. Calibration (e.g. calculating weights for WMV or fitting a meta-model) on validation data + 2. Prediction on test data (i.e., turning base learner predictions into aggregated predictions) - self.models = [] - self.positive_prediction_threshold = 0.5 - for model_name, model_config in processed_configs.items(): - model_cls = MODEL_TYPES[model_config["type"]] - if "hugging_face" in model_config: - hugging_face_kwargs = download_model_files(model_config["hugging_face"]) - else: - hugging_face_kwargs = {} - if "package_name" in model_config: - check_package_installed(model_config["package_name"]) + Each ensemble gets a directory where it can store its calibration results (e.g. weights for WMV or meta-model parameters). - model_instance = model_cls( - model_name, - **model_config, - **hugging_face_kwargs, - chebi_graph=self.chebi_graph, - ) - assert isinstance(model_instance, BasePredictor) - self.models.append(model_instance) + Not part of the ensemble are + - getting predictions from base learners + - resolving inconsistencies in the aggregated predictions + """ - if resolve_inconsistencies: - self.smoother = ScoreBasedPredictionSmoother( - self.chebi_graph, - label_names=None, - disjoint_files=self.disjoint_files, - verbose=self.verbose_output, - ) - else: - self.smoother = None + def __init__(self, ensemble_dir: str): + os.makedirs(ensemble_dir, exist_ok=True) + self.ensemble_dir = ensemble_dir - def gather_predictions(self, smiles_list): - # get predictions from all models for the SMILES list - # order them alphabetically by label class - model_predictions = [] - predicted_classes = set() - for model in self.models: - model_predictions.append(model.predict_smiles_list(smiles_list)) - for logits_for_smiles in model_predictions[-1]: - if logits_for_smiles is not None: - for cls in logits_for_smiles: - predicted_classes.add(cls) - if self.verbose_output: - print(f"Sorting predictions from {len(model_predictions)} models...") - predicted_classes = sorted(list(predicted_classes)) - predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)} - ordered_logits = ( - torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) - * torch.nan - ) - for i, model_prediction in enumerate(model_predictions): - for j, logits_for_smiles in tqdm.tqdm( - enumerate(model_prediction), - total=len(model_prediction), - desc=f"Sorting predictions for {self.models[i].model_name}", - ): - if logits_for_smiles is not None: - for cls in logits_for_smiles: - ordered_logits[j, predicted_classes_dict[cls], i] = ( - logits_for_smiles[cls] - ) + @property + def ensemble_name(self): + return self.__class__.__name__ - return ordered_logits, predicted_classes - - def consolidate_predictions( + def calibrate( self, - predictions, - classwise_weights, - return_intermediate_results=False, - **kwargs, + validation_predictions: dict[str, torch.Tensor], + validation_data: list[Chem.Mol], + validation_labels: torch.Tensor, ): + """Calibrate the ensemble model using validation predictions and labels. + At the end, save the calibration results (e.g. weights for WMV or meta-model parameters) to self.ensemble_dir. """ - Aggregates predictions from multiple models using weighted majority voting. - Optimized version using tensor operations instead of for loops. - """ - num_smiles, num_classes, num_models = predictions.shape - - # Get predictions for all classes - valid_predictions = ~torch.isnan(predictions) - valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension - - # Skip classes with no valid predictions - has_valid_predictions = valid_counts > 0 - - # Calculate positive and negative predictions for all classes at once - positive_mask = ( - predictions > self.positive_prediction_threshold - ) & valid_predictions - negative_mask = ( - predictions < self.positive_prediction_threshold - ) & valid_predictions - - # if use_confidence is passed in kwargs, it overrides the ensemble setting - use_confidence = kwargs.get("use_confidence", self.use_confidence) - if use_confidence: - confidence = 2 * torch.abs( - predictions.nan_to_num() - self.positive_prediction_threshold - ) - else: - confidence = torch.ones_like(predictions) - - # Extract positive and negative weights - pos_weights = classwise_weights[0] # Shape: (num_classes, num_models) - neg_weights = classwise_weights[1] # Shape: (num_classes, num_models) - - # Calculate weighted predictions using broadcasting - # predictions shape: (num_smiles, num_classes, num_models) - # weights shape: (num_classes, num_models) - positive_weighted = ( - positive_mask.float() * confidence * pos_weights.unsqueeze(0) - ) - negative_weighted = ( - negative_mask.float() * confidence * neg_weights.unsqueeze(0) - ) - - # Sum over models dimension - positive_sum = positive_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - negative_sum = negative_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - - # Determine which classes to include for each SMILES - net_score = positive_sum - negative_sum # Shape: (num_smiles, num_classes) - if return_intermediate_results: - return ( - net_score, - has_valid_predictions, - { - "positive_mask": positive_mask, - "negative_mask": negative_mask, - "confidence": confidence, - "positive_sum": positive_sum, - "negative_sum": negative_sum, - }, - ) - - return net_score, has_valid_predictions - - def apply_inconsistency_resolution( - self, net_score, class_names, has_valid_predictions - ): - # Smooth predictions - start_time = time.perf_counter() - if self.smoother is not None: - self.smoother.set_label_names(class_names) - smooth_net_score = self.smoother(net_score) - class_decisions = ( - smooth_net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - else: - class_decisions = ( - net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - end_time = time.perf_counter() - if self.verbose_output: - print(f"Prediction smoothing took {end_time - start_time:.2f} seconds") - - complete_failure = torch.all(~has_valid_predictions, dim=1) - return class_decisions, complete_failure - - def calculate_classwise_weights(self, predicted_classes): - """No weights, simple majority voting""" - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - - return positive_weights, negative_weights - - def predict_smiles_list( - self, smiles_list, return_intermediate_results=False, **kwargs - ) -> list: - ordered_predictions, predicted_classes = self.gather_predictions(smiles_list) - if len(predicted_classes) == 0: - print("Warning: No classes have been predicted for the given SMILES list.") - predicted_classes = {cls: i for i, cls in enumerate(predicted_classes)} - - classwise_weights = self.calculate_classwise_weights(predicted_classes) - if return_intermediate_results: - net_score, has_valid_predictions, intermediate_results_dict = ( - self.consolidate_predictions( - ordered_predictions, - classwise_weights, - return_intermediate_results=return_intermediate_results, - ) - ) - else: - net_score, has_valid_predictions = self.consolidate_predictions( - ordered_predictions, classwise_weights - ) - class_decisions, is_failure = self.apply_inconsistency_resolution( - net_score, list(predicted_classes.keys()), has_valid_predictions - ) - - class_names = list(predicted_classes.keys()) - class_indices = {predicted_classes[cls]: cls for cls in class_names} - result = [ - ( - [ - class_indices[idx.item()] - for idx in torch.nonzero(i, as_tuple=True)[0] - ] - if not failure - else None - ) - for i, failure in zip(class_decisions, is_failure) - ] - if return_intermediate_results: - intermediate_results_dict["predicted_classes"] = predicted_classes - intermediate_results_dict["classwise_weights"] = classwise_weights - intermediate_results_dict["net_score"] = net_score - return result, intermediate_results_dict - - return result - + pass -if __name__ == "__main__": - ensemble = BaseEnsemble( - { - "resgated_0ps1g189": { - "type": "resgated", - "ckpt_path": "data/0ps1g189/epoch=122.ckpt", - "molecular_properties": [ - "chebai_graph.preprocessing.properties.AtomType", - "chebai_graph.preprocessing.properties.NumAtomBonds", - "chebai_graph.preprocessing.properties.AtomCharge", - "chebai_graph.preprocessing.properties.AtomAromaticity", - "chebai_graph.preprocessing.properties.AtomHybridization", - "chebai_graph.preprocessing.properties.AtomNumHs", - "chebai_graph.preprocessing.properties.BondType", - "chebai_graph.preprocessing.properties.BondInRing", - "chebai_graph.preprocessing.properties.BondAromaticity", - "chebai_graph.preprocessing.properties.RDKit2DNormalized", - ], - # "classwise_weights_path" : "../python-chebai/metrics_0ps1g189_80-10-10.json" - }, - "electra_14ko0zcf": { - "type": "electra", - "ckpt_path": "data/14ko0zcf/epoch=193.ckpt", - # "classwise_weights_path": "../python-chebai/metrics_electra_14ko0zcf_80-10-10.json", - }, - } - ) - r = ensemble.predict_smiles_list( - [ - "[NH3+]CCCC[C@H](NC(=O)[C@@H]([NH3+])CC([O-])=O)C([O-])=O", - "C[C@H](N)C(=O)NCC(O)=O#", - "", - ], - load_preds_if_possible=False, - ) - print(len(r), r[0]) + def predict(self, test_predictions: dict[str, torch.Tensor]): + raise NotImplementedError() diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py new file mode 100644 index 0000000..38aa749 --- /dev/null +++ b/chebifier/ensemble/voting_ensemble.py @@ -0,0 +1,304 @@ +import importlib +import time +from pathlib import Path +from typing import Union + +import torch +import tqdm +import yaml + +from chebifier.check_env import check_package_installed +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.hugging_face import download_model_files +from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother +from chebifier.prediction_models.base_predictor import BasePredictor +from chebifier.utils import ( + get_default_configs, + get_disjoint_files, + load_chebi_graph, + process_config, +) + + +class VotingEnsemble(BaseEnsemble): + def __init__( + self, + model_configs: Union[str, Path, dict, None] = None, + resolve_inconsistencies: bool = True, + verbose_output: bool = False, + use_confidence: bool = True, + ): + # Deferred Import: To avoid circular import error + from chebifier.model_registry import MODEL_TYPES + + # Load configuration from YAML file + if not model_configs: + config = get_default_configs() + elif isinstance(model_configs, dict): + config = model_configs + else: + print(f"Loading ensemble configuration from {model_configs}") + with open(model_configs, "r") as f: + config = yaml.safe_load(f) + + with ( + importlib.resources.files("chebifier") + .joinpath("model_registry.yml") + .open("r") as f + ): + model_registry = yaml.safe_load(f) + + processed_configs = process_config(config, model_registry) + self.verbose_output = verbose_output + self.use_confidence = use_confidence + + self.chebi_graph = load_chebi_graph() + self.disjoint_files = get_disjoint_files() + + self.models = [] + self.positive_prediction_threshold = 0.5 + for model_name, model_config in processed_configs.items(): + model_cls = MODEL_TYPES[model_config["type"]] + if "hugging_face" in model_config: + hugging_face_kwargs = download_model_files(model_config["hugging_face"]) + else: + hugging_face_kwargs = {} + if "package_name" in model_config: + check_package_installed(model_config["package_name"]) + + model_instance = model_cls( + model_name, + **model_config, + **hugging_face_kwargs, + chebi_graph=self.chebi_graph, + ) + assert isinstance(model_instance, BasePredictor) + self.models.append(model_instance) + + if resolve_inconsistencies: + self.smoother = ScoreBasedPredictionSmoother( + self.chebi_graph, + label_names=None, + disjoint_files=self.disjoint_files, + verbose=self.verbose_output, + ) + else: + self.smoother = None + + def gather_predictions(self, smiles_list): + # get predictions from all models for the SMILES list + # order them alphabetically by label class + model_predictions = [] + predicted_classes = set() + for model in self.models: + model_predictions.append(model.predict_smiles_list(smiles_list)) + for logits_for_smiles in model_predictions[-1]: + if logits_for_smiles is not None: + for cls in logits_for_smiles: + predicted_classes.add(cls) + if self.verbose_output: + print(f"Sorting predictions from {len(model_predictions)} models...") + predicted_classes = sorted(list(predicted_classes)) + predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)} + ordered_logits = ( + torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) + * torch.nan + ) + for i, model_prediction in enumerate(model_predictions): + for j, logits_for_smiles in tqdm.tqdm( + enumerate(model_prediction), + total=len(model_prediction), + desc=f"Sorting predictions for {self.models[i].model_name}", + ): + if logits_for_smiles is not None: + for cls in logits_for_smiles: + ordered_logits[j, predicted_classes_dict[cls], i] = ( + logits_for_smiles[cls] + ) + + return ordered_logits, predicted_classes + + def consolidate_predictions( + self, + predictions, + classwise_weights, + return_intermediate_results=False, + **kwargs, + ): + """ + Aggregates predictions from multiple models using weighted majority voting. + Optimized version using tensor operations instead of for loops. + """ + num_smiles, num_classes, num_models = predictions.shape + + # Get predictions for all classes + valid_predictions = ~torch.isnan(predictions) + valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension + + # Skip classes with no valid predictions + has_valid_predictions = valid_counts > 0 + + # Calculate positive and negative predictions for all classes at once + positive_mask = ( + predictions > self.positive_prediction_threshold + ) & valid_predictions + negative_mask = ( + predictions < self.positive_prediction_threshold + ) & valid_predictions + + # if use_confidence is passed in kwargs, it overrides the ensemble setting + use_confidence = kwargs.get("use_confidence", self.use_confidence) + if use_confidence: + confidence = 2 * torch.abs( + predictions.nan_to_num() - self.positive_prediction_threshold + ) + else: + confidence = torch.ones_like(predictions) + + # Extract positive and negative weights + pos_weights = classwise_weights[0] # Shape: (num_classes, num_models) + neg_weights = classwise_weights[1] # Shape: (num_classes, num_models) + + # Calculate weighted predictions using broadcasting + # predictions shape: (num_smiles, num_classes, num_models) + # weights shape: (num_classes, num_models) + positive_weighted = ( + positive_mask.float() * confidence * pos_weights.unsqueeze(0) + ) + negative_weighted = ( + negative_mask.float() * confidence * neg_weights.unsqueeze(0) + ) + + # Sum over models dimension + positive_sum = positive_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) + negative_sum = negative_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) + + # Determine which classes to include for each SMILES + net_score = positive_sum - negative_sum # Shape: (num_smiles, num_classes) + if return_intermediate_results: + return ( + net_score, + has_valid_predictions, + { + "positive_mask": positive_mask, + "negative_mask": negative_mask, + "confidence": confidence, + "positive_sum": positive_sum, + "negative_sum": negative_sum, + }, + ) + + return net_score, has_valid_predictions + + def apply_inconsistency_resolution( + self, net_score, class_names, has_valid_predictions + ): + # Smooth predictions + start_time = time.perf_counter() + if self.smoother is not None: + self.smoother.set_label_names(class_names) + smooth_net_score = self.smoother(net_score) + class_decisions = ( + smooth_net_score > 0 + ) & has_valid_predictions # Shape: (num_smiles, num_classes) + else: + class_decisions = ( + net_score > 0 + ) & has_valid_predictions # Shape: (num_smiles, num_classes) + end_time = time.perf_counter() + if self.verbose_output: + print(f"Prediction smoothing took {end_time - start_time:.2f} seconds") + + complete_failure = torch.all(~has_valid_predictions, dim=1) + return class_decisions, complete_failure + + def calculate_classwise_weights(self, predicted_classes): + """No weights, simple majority voting""" + positive_weights = torch.ones(len(predicted_classes), len(self.models)) + negative_weights = torch.ones(len(predicted_classes), len(self.models)) + + return positive_weights, negative_weights + + def predict_smiles_list( + self, smiles_list, return_intermediate_results=False, **kwargs + ) -> list: + ordered_predictions, predicted_classes = self.gather_predictions(smiles_list) + if len(predicted_classes) == 0: + print("Warning: No classes have been predicted for the given SMILES list.") + predicted_classes = {cls: i for i, cls in enumerate(predicted_classes)} + + classwise_weights = self.calculate_classwise_weights(predicted_classes) + if return_intermediate_results: + net_score, has_valid_predictions, intermediate_results_dict = ( + self.consolidate_predictions( + ordered_predictions, + classwise_weights, + return_intermediate_results=return_intermediate_results, + ) + ) + else: + net_score, has_valid_predictions = self.consolidate_predictions( + ordered_predictions, classwise_weights + ) + class_decisions, is_failure = self.apply_inconsistency_resolution( + net_score, list(predicted_classes.keys()), has_valid_predictions + ) + + class_names = list(predicted_classes.keys()) + class_indices = {predicted_classes[cls]: cls for cls in class_names} + result = [ + ( + [ + class_indices[idx.item()] + for idx in torch.nonzero(i, as_tuple=True)[0] + ] + if not failure + else None + ) + for i, failure in zip(class_decisions, is_failure) + ] + if return_intermediate_results: + intermediate_results_dict["predicted_classes"] = predicted_classes + intermediate_results_dict["classwise_weights"] = classwise_weights + intermediate_results_dict["net_score"] = net_score + return result, intermediate_results_dict + + return result + + +if __name__ == "__main__": + ensemble = VotingEnsemble( + { + "resgated_0ps1g189": { + "type": "resgated", + "ckpt_path": "data/0ps1g189/epoch=122.ckpt", + "molecular_properties": [ + "chebai_graph.preprocessing.properties.AtomType", + "chebai_graph.preprocessing.properties.NumAtomBonds", + "chebai_graph.preprocessing.properties.AtomCharge", + "chebai_graph.preprocessing.properties.AtomAromaticity", + "chebai_graph.preprocessing.properties.AtomHybridization", + "chebai_graph.preprocessing.properties.AtomNumHs", + "chebai_graph.preprocessing.properties.BondType", + "chebai_graph.preprocessing.properties.BondInRing", + "chebai_graph.preprocessing.properties.BondAromaticity", + "chebai_graph.preprocessing.properties.RDKit2DNormalized", + ], + # "classwise_weights_path" : "../python-chebai/metrics_0ps1g189_80-10-10.json" + }, + "electra_14ko0zcf": { + "type": "electra", + "ckpt_path": "data/14ko0zcf/epoch=193.ckpt", + # "classwise_weights_path": "../python-chebai/metrics_electra_14ko0zcf_80-10-10.json", + }, + } + ) + r = ensemble.predict_smiles_list( + [ + "[NH3+]CCCC[C@H](NC(=O)[C@@H]([NH3+])CC([O-])=O)C([O-])=O", + "C[C@H](N)C(=O)NCC(O)=O#", + "", + ], + load_preds_if_possible=False, + ) + print(len(r), r[0]) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 1353325..59d0c32 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,9 +1,12 @@ +import json +import os + import torch -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import VotingEnsemble -class WMVwithPPVNPVEnsemble(BaseEnsemble): +class WMVwithPPVNPVEnsemble(VotingEnsemble): def __init__( self, config_path=None, weighting_strength=1, weighting_exponent=1, **kwargs @@ -19,6 +22,18 @@ def __init__( self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_classwise_weights = dict() + for model in self.models: + classwise_weights_path = os.path.join( + self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + ) + if os.path.exists(classwise_weights_path): + self.model_classwise_weights[model.model_name] = json.load( + open(classwise_weights_path, encoding="utf-8") + ) + else: + self.model_classwise_weights[model.model_name] = None + def calculate_classwise_weights(self, predicted_classes): """ Given the positions of predicted classes in the predictions tensor, assign weights to each class. The @@ -30,9 +45,9 @@ def calculate_classwise_weights(self, predicted_classes): for j, model in enumerate(self.models): positive_weights[:, j] *= model.model_weight negative_weights[:, j] *= model.model_weight - if model.classwise_weights is None: + if self.model_classwise_weights[model.model_name] is None: continue - for cls, weights in model.classwise_weights.items(): + for cls, weights in self.model_classwise_weights[model.model_name].items(): if cls not in predicted_classes: continue ppv = ( @@ -64,7 +79,7 @@ def calculate_classwise_weights(self, predicted_classes): return positive_weights, negative_weights -class WMVwithF1Ensemble(BaseEnsemble): +class WMVwithF1Ensemble(VotingEnsemble): def __init__( self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs @@ -77,6 +92,18 @@ def __init__( self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_classwise_weights = dict() + for model in self.models: + classwise_weights_path = os.path.join( + self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + ) + if os.path.exists(classwise_weights_path): + self.model_classwise_weights[model.model_name] = json.load( + open(classwise_weights_path, encoding="utf-8") + ) + else: + self.model_classwise_weights[model.model_name] = None + def calculate_classwise_weights(self, predicted_classes): """ Given the positions of predicted classes in the predictions tensor, assign weights to each class. The @@ -86,9 +113,9 @@ def calculate_classwise_weights(self, predicted_classes): weights_by_cls = torch.ones(len(predicted_classes), len(self.models)) for j, model in enumerate(self.models): weights_by_cls[:, j] *= model.model_weight - if model.classwise_weights is None: + if self.model_classwise_weights[model.model_name] is None: continue - for cls, weights in model.classwise_weights.items(): + for cls, weights in self.model_classwise_weights[model.model_name].items(): if cls in predicted_classes: if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0: f1 = ( diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index a8287bb..ef13e7e 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,4 +1,4 @@ -from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.ensemble.voting_ensemble import VotingEnsemble from chebifier.ensemble.weighted_majority_ensemble import ( WMVwithF1Ensemble, WMVwithPPVNPVEnsemble, @@ -19,7 +19,7 @@ ) ENSEMBLES = { - "mv": BaseEnsemble, + "mv": VotingEnsemble, "wmv-ppvnpv": WMVwithPPVNPVEnsemble, "wmv-f1": WMVwithF1Ensemble, } diff --git a/chebifier/predict.py b/chebifier/predict.py new file mode 100644 index 0000000..9ff2530 --- /dev/null +++ b/chebifier/predict.py @@ -0,0 +1 @@ +# Get end-to-end predictions (from SMILES via base learners + ensemble + inconsistency resolution to ChEBI classes) diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 37851b2..1082b8e 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,6 +1,7 @@ -import json from abc import ABC +from rdkit import Chem + from .._custom_cache import modelwise_smiles_lru_cache @@ -9,27 +10,20 @@ def __init__( self, model_name: str, model_weight: int = 1, - classwise_weights_path: str = None, **kwargs, ): self.model_name = model_name self.model_weight = model_weight - if classwise_weights_path is not None: - self.classwise_weights = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.classwise_weights = None self._description = kwargs.get("description", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> dict: + def predict_list(self, molecule_list: list[str | Chem.Mol]) -> dict: raise NotImplementedError() - def predict_smiles(self, smiles: str) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict: # by default, use list-based prediction - return self.predict_smiles_list([smiles])[0] + return self.predict_list([molecule])[0] @property def info_text(self): diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index bf6c39b..f0a1cfd 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/chebifier/prediction_models/c3p_predictor.py @@ -25,7 +25,7 @@ def __init__( self.chebi_graph = kwargs.get("chebi_graph", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: from c3p import classifier as c3p_classifier result_list = [] diff --git a/chebifier/prediction_models/chebi_lookup.py b/chebifier/prediction_models/chebi_lookup.py index d68c9cd..006af48 100644 --- a/chebifier/prediction_models/chebi_lookup.py +++ b/chebifier/prediction_models/chebi_lookup.py @@ -69,7 +69,7 @@ def build_smiles_lookup(self): ) return smiles_lookup - def predict_smiles(self, smiles: str) -> Optional[dict]: + def predict(self, smiles: str) -> Optional[dict]: if not smiles: return None mol = _smiles_to_mol(smiles) @@ -96,10 +96,10 @@ def predict_smiles(self, smiles: str) -> Optional[dict]: return None @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: predictions = [] for smiles in smiles_list: - predictions.append(self.predict_smiles(smiles)) + predictions.append(self.predict(smiles)) return predictions @@ -150,5 +150,5 @@ def explain_smiles(self, smiles: str) -> dict: "C1=CC=CC=C1", "*C(=O)OC[C@H](COP(=O)([O-])OCC[N+](C)(C)C)OC(*)=O", ] # SMILES with 251 matches in ChEBI - predictions = predictor.predict_smiles_list(smiles_list) + predictions = predictor.predict_list(smiles_list) print(predictions) diff --git a/chebifier/prediction_models/chemlog_predictor.py b/chebifier/prediction_models/chemlog_predictor.py index 51fccd0..f637e98 100644 --- a/chebifier/prediction_models/chemlog_predictor.py +++ b/chebifier/prediction_models/chemlog_predictor.py @@ -43,7 +43,7 @@ def __init__(self, model_name: str, **kwargs): ] @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: results = [] for predictor in self.predictors: predictor_results = predictor._predict_smiles_list(smiles_list) @@ -66,7 +66,7 @@ def __init__(self, model_name: str, **kwargs): self.classifier = None @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: return self._predict_smiles_list(smiles_list) def _predict_smiles_list(self, smiles_list: list[str]) -> list: @@ -141,7 +141,7 @@ def __init__(self, model_name: str, **kwargs): # fmt: on print(f"Initialised ChemLog model {self.model_name}") - def predict_smiles(self, smiles: str) -> Optional[dict]: + def predict(self, smiles: str) -> Optional[dict]: from chemlog.cli import _smiles_to_mol, strategy_call mol = _smiles_to_mol(smiles) @@ -168,13 +168,13 @@ def predict_smiles(self, smiles: str) -> Optional[dict]: } @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: return self._predict_smiles_list(smiles_list) def _predict_smiles_list(self, smiles_list: list[str]) -> list: results = [] for i, smiles in tqdm.tqdm(enumerate(smiles_list)): - results.append(self.predict_smiles(smiles)) + results.append(self.predict(smiles)) for classifier in self.classifier_instances.values(): classifier.on_finish() diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 971a42d..2bfd389 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -29,7 +29,7 @@ def __init__( ) @modelwise_smiles_lru_cache.batch_decorator - def predict_smiles_list(self, smiles_list: list[str]) -> list: + def predict_list(self, smiles_list: list[str]) -> list: """ Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary of classes and predicted values. diff --git a/pyproject.toml b/pyproject.toml index 60cd4fb..d73db76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "pyyaml", "tqdm", "rdkit", + "chebi-utils>=0.3", # Package to install manually if required #"chebai>=1.0.1", #"chemlog>=1.0.4", From 4a167f96ff5aa77e7cfbf6aa7926813adfc86362 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Wed, 29 Jul 2026 10:43:28 +0200 Subject: [PATCH 02/10] remove chebi_graph builder (now done via chebi_utils library) --- chebifier/utils.py | 88 ---------------------------------------------- 1 file changed, 88 deletions(-) diff --git a/chebifier/utils.py b/chebifier/utils.py index 96f4a80..8bfa6f3 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -3,9 +3,6 @@ import os import pickle -import fastobo -import networkx as nx -import requests import yaml from rdkit import Chem @@ -29,82 +26,6 @@ def load_chebi_graph(filename=None): return pickle.load(open(file, "rb")) -def term_callback(doc): - """Similar to the chebai function, but reduced to the necessary fields. Also, ChEBI IDs are strings""" - parents = [] - name = None - smiles = None - subset = None - for clause in doc: - if isinstance(clause, fastobo.term.PropertyValueClause): - t = clause.property_value - if str(t.relation) == "http://purl.obolibrary.org/obo/chebi/smiles": - assert smiles is None - smiles = t.value - # in older chebi versions, smiles strings are synonyms - # e.g. synonym: "[F-].[Na+]" RELATED SMILES [ChEBI] - elif isinstance(clause, fastobo.term.SynonymClause): - if "SMILES" in clause.raw_value(): - assert smiles is None - smiles = clause.raw_value().split('"')[1] - elif isinstance(clause, fastobo.term.IsAClause): - chebi_id = str(clause.term) - chebi_id = chebi_id[chebi_id.index(":") + 1 :] - parents.append(chebi_id) - elif isinstance(clause, fastobo.term.NameClause): - name = str(clause.name) - elif isinstance(clause, fastobo.term.SubsetClause): - subset = str(clause.subset) - if isinstance(clause, fastobo.term.IsObsoleteClause): - if clause.obsolete: - # if the term document contains clause as obsolete as true, skips this document. - return False - chebi_id = str(doc.id) - chebi_id = chebi_id[chebi_id.index(":") + 1 :] - return { - "id": chebi_id, - "parents": parents, - "name": name, - "smiles": smiles, - "subset": subset, - } - - -def build_chebi_graph(chebi_version=241): - """Creates a networkx graph for the ChEBI hierarchy. Usually, you don't want to call this function directly, but rather use the `load_chebi_graph` function.""" - chebi_path = os.path.join("data", f"chebi_v{chebi_version}", "chebi.obo") - os.makedirs(os.path.join("data", f"chebi_v{chebi_version}"), exist_ok=True) - if not os.path.exists(chebi_path): - url = f"http://purl.obolibrary.org/obo/chebi/{chebi_version}/chebi.obo" - r = requests.get(url, allow_redirects=True) - open(chebi_path, "wb").write(r.content) - with open(chebi_path, encoding="utf-8") as chebi: - chebi = "\n".join(line for line in chebi if not line.startswith("xref:")) - - elements = [] - for term_doc in fastobo.loads(chebi): - if ( - term_doc - and isinstance(term_doc.id, fastobo.id.PrefixedIdent) - and term_doc.id.prefix == "CHEBI" - ): - term_dict = term_callback(term_doc) - if term_dict: - elements.append(term_dict) - - g = nx.DiGraph() - for n in elements: - g.add_node(n["id"], **n) - - # Only take the edges which connect the existing nodes, to avoid internal creation of obsolete nodes - # https://github.com/ChEB-AI/python-chebai/pull/55#issuecomment-2386654142 - g.add_edges_from( - [(p, q["id"]) for q in elements for p in q["parents"] if g.has_node(p)], - label="direct_child", - ) - return nx.transitive_closure_dag(g) - - def get_disjoint_files(): """Gets local disjointness files if they are present in the right location, otherwise downloads them from Hugging Face.""" local_disjoint_files = [ @@ -168,12 +89,3 @@ def _smiles_to_mol(smiles: str): except Chem.KekulizeException as e: print(f"Failed to Kekulize {smiles}: {e}") return mol - - -if __name__ == "__main__": - chebi_graph = build_chebi_graph(chebi_version=244) - os.makedirs(os.path.join("data", "chebi_v244"), exist_ok=True) - pickle.dump( - chebi_graph, - open(os.path.join("data", "chebi_v244", "chebi_graph.pkl"), "wb"), - ) From b26ba7c96c1afb74ccd10d1fe2395afed41f2694 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Wed, 29 Jul 2026 12:56:55 +0200 Subject: [PATCH 03/10] update CLI and integrate MV ensemble into new workflow --- chebifier/build_ensemble.py | 8 + chebifier/cli.py | 341 +++++++++++++---- chebifier/ensemble/voting_ensemble.py | 347 +++++------------- chebifier/predict.py | 146 +++++++- chebifier/prediction_models/base_predictor.py | 4 +- chebifier/prediction_models/nn_predictor.py | 2 +- 6 files changed, 506 insertions(+), 342 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index eaf2cba..5a9df68 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -2,6 +2,8 @@ import torch +from chebifier.predict import collect_base_learner_predictions + class EnsembleBuilder: """ @@ -28,6 +30,7 @@ def __init__( self.validation_data = validation_data self.validation_labels = validation_labels self.prediction_cache_dir = prediction_cache_dir + os.makedirs(self.prediction_cache_dir, exist_ok=True) def build_ensemble(self): """ @@ -38,6 +41,7 @@ def build_ensemble(self): # Step 1: Get predictions from base learners on validation data validation_predictions = {} + classes = {} # get cached predictions if available, otherwise compute and cache them for model_name, model in self.base_learners.items(): cache_path = os.path.join( @@ -53,6 +57,10 @@ def build_ensemble(self): ) torch.save(validation_predictions[model_name], cache_path) + validation_predictions, classes = collect_base_learner_predictions( + validation_predictions + ) + # Step 2: Calibrate the ensemble model using validation predictions self.ensemble_model.calibrate( validation_predictions, self.validation_data, self.validation_labels diff --git a/chebifier/cli.py b/chebifier/cli.py index 2f0468c..be9bbd7 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -1,6 +1,160 @@ +import importlib.resources +import os +from typing import Literal + import click +import pandas as pd +import yaml +from rdkit import Chem + +from chebifier.build_ensemble import EnsembleBuilder +from chebifier.check_env import check_package_installed +from chebifier.hugging_face import download_model_files +from chebifier.model_registry import ENSEMBLES, MODEL_TYPES +from chebifier.predict import predict as predict_molecules +from chebifier.utils import get_default_configs, load_chebi_graph, process_config + + +def read_molecules(molecules, molecule_file): + """Collect molecules from CLI arguments and/or a file (one molecule per line) and convert them to RDKit mol objects.""" + raw_inputs = list(molecules) + if molecule_file: + with open(molecule_file, "r", encoding="utf-8") as f: + raw_inputs.extend([line.strip() for line in f if line.strip()]) + + mol_list = [] + for raw_input in raw_inputs: + try: + if raw_input.startswith("InChI="): + mol = Chem.MolFromInchi(raw_input, sanitize=False) + if mol is None: + click.echo(f"Failed to parse InChI: {raw_input}") + mol_list.append(None) + mol_list.append(mol) + elif Chem.MolFromSmiles(raw_input, sanitize=False) is None: + click.echo(f"Failed to parse SMILES: {raw_input}") + mol_list.append(None) + else: + mol_list.append(Chem.MolFromSmiles(raw_input, sanitize=False)) + except Exception as e: + click.echo(f"Error parsing molecule '{raw_input}': {e}.") + mol_list.append(None) + return mol_list + + +def build_base_learners(ensemble_config): + """Instantiate the base learners described by an ensemble configuration file.""" + if ensemble_config is None: + config = get_default_configs() + else: + print(f"Loading ensemble configuration from {ensemble_config}") + with open(ensemble_config, "r") as f: + config = yaml.safe_load(f) + + with ( + importlib.resources.files("chebifier") + .joinpath("model_registry.yml") + .open("r") as f + ): + model_registry = yaml.safe_load(f) + + chebi_graph = load_chebi_graph() + base_learners = {} + for model_name, model_config in process_config(config, model_registry).items(): + if "hugging_face" in model_config: + hugging_face_kwargs = download_model_files(model_config["hugging_face"]) + else: + hugging_face_kwargs = {} + if "package_name" in model_config: + check_package_installed(model_config["package_name"]) + base_learners[model_name] = MODEL_TYPES[model_config["type"]]( + model_name, + **model_config, + **hugging_face_kwargs, + chebi_graph=chebi_graph, + ) + return base_learners + + +def load_dataset(data_path, split: Literal["train", "validation", "test"]): + data_file = os.path.join(data_path, "data.pkl") + splits_file = os.path.join(data_path, "splits.csv") + if not os.path.exists(data_file) or not os.path.exists(splits_file): + raise FileNotFoundError( + f"Required dataset files not found. Expected to find 'data.pkl' and 'splits.csv' in the provided data path ({data_path})." + ) + data_df = pd.read_pickle(data_file) + splits_df = pd.read_csv(splits_file) + # merge dataframe on id column and filter by the specified split + splits_df["id"] = splits_df["id"].astype(str) + merged_df = data_df.merge(splits_df, left_on="chebi_id", right_on="id", how="inner") + merged_df = merged_df[merged_df["split"] == split].reset_index(drop=True) + + mol_list = merged_df["mol"].tolist() + + # extract labels from data_df: every column other than chebi_id/mol/id/split is a ChEBI class label + label_columns = [c for c in data_df.columns if c not in ("chebi_id", "mol")] + labels_df = merged_df[label_columns].astype(bool) + labels_df.columns = [str(c) for c in label_columns] + + print( + f"Loaded {len(mol_list)} molecules and {len(labels_df.columns)} labels for split '{split}' from {data_path}." + ) + + return mol_list, labels_df -from chebifier.model_registry import ENSEMBLES + +def ensemble_options(command): + """Options shared by all commands that use an ensemble.""" + for option in reversed( + [ + click.option( + "--ensemble-config", + "-e", + type=click.Path(exists=True), + default=None, + help="Configuration file listing the base learners of the ensemble", + ), + click.option( + "--ensemble-type", + "-t", + type=click.Choice(ENSEMBLES.keys()), + default="wmv-f1", + help="Type of ensemble to use (default: Weighted Majority Voting with F1 weights)", + ), + click.option( + "--ensemble-dir", + "-d", + type=click.Path(), + required=True, + help="Directory where the calibration results of the ensemble are stored", + ), + click.option( + "--prediction-cache-dir", + type=click.Path(), + default=None, + help="Directory for caching base learner predictions", + ), + ] + ): + command = option(command) + return command + + +def data_options(command): + """Options shared by the commands that work on a ChEBI dataset split.""" + for option in reversed( + [ + click.option( + "--data-path", + type=str, + required=True, + help="Data source: local dataset directory or Hugging Face repo id", + ), + ] + ): + command = option(command) + return command @click.group() @@ -10,107 +164,138 @@ def cli(): @cli.command() +@ensemble_options +@data_options +def build( + ensemble_config, ensemble_type, ensemble_dir, prediction_cache_dir, data_path +): + """Build (calibrate) an ensemble on the ChEBI validation set.""" + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + + # TODO: Hugging Face support + validation_data, validation_labels = load_dataset(data_path, split="validation") + + builder = EnsembleBuilder( + base_learners, + ensemble_model, + validation_data, + validation_labels, + prediction_cache_dir, + ) + builder.build_ensemble() + + +@cli.command() +@ensemble_options +@data_options @click.option( - "--ensemble-config", - "-e", - type=click.Path(exists=True), - default=None, - help="Configuration file for ensemble models", -) -@click.option("--smiles", "-s", multiple=True, help="SMILES strings to predict") -@click.option( - "--smiles-file", - "-f", - type=click.Path(exists=True), - help="File containing SMILES strings (one per line)", + "--resolve-inconsistencies/--no-resolve-inconsistencies", + default=True, + help="Resolve inconsistencies in the aggregated predictions (default: True)", ) @click.option( "--output", "-o", type=click.Path(), - help="Output file to save predictions (optional)", + default=None, + help="Output file to save the evaluation results (optional)", ) +def evaluate( + ensemble_config, + ensemble_type, + ensemble_dir, + prediction_cache_dir, + data_path, + resolve_inconsistencies, + output, +): + """Evaluate an ensemble on the ChEBI test set.""" + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) + + # TODO: Hugging Face support + test_data, test_labels = load_dataset(data_path, split="test") + + predictions = predict_molecules( + base_learners, + ensemble_model, + test_data, + prediction_cache_dir=prediction_cache_dir, + resolve_inconsistencies=resolve_inconsistencies, + ) + + print(f"Predictions: {predictions}") + + # TODO: compare predictions to test_labels, report metrics and save them to output + + +@cli.command() +@ensemble_options @click.option( - "--ensemble-type", - "-t", - type=click.Choice(ENSEMBLES.keys()), - default="wmv-f1", - help="Type of ensemble to use (default: Weighted Majority Voting)", + "--molecules", "-m", multiple=True, help="SMILES or InChI strings to predict" ) @click.option( - "--use-confidence", - "-c", - is_flag=True, - default=True, - help="Weight predictions based on how 'confident' a model is in its prediction (default: True)", + "--molecule-file", + "-f", + type=click.Path(exists=True), + default=None, + help="File containing SMILES or InChI strings (one per line)", ) @click.option( - "--resolve-inconsistencies", - "-r", - is_flag=True, + "--resolve-inconsistencies/--no-resolve-inconsistencies", default=True, - help="Resolve inconsistencies in predictions automatically (default: True)", + help="Resolve inconsistencies in the aggregated predictions (default: True)", +) +@click.option( + "--output", + "-o", + type=click.Path(), + default=None, + help="Output file to save the predictions (optional)", ) @click.option( - "--verbose", - "-v", - is_flag=True, - default=False, - help="Enable verbose output", + "--decision-threshold", + "-dt", + type=float, + default=0, + help="Threshold for classifying predictions (default: 0)", ) def predict( ensemble_config, - smiles, - smiles_file, - output, ensemble_type, - use_confidence, - resolve_inconsistencies=True, - verbose=False, + ensemble_dir, + prediction_cache_dir, + molecules, + molecule_file, + resolve_inconsistencies, + decision_threshold, + output, ): - """Predict ChEBI classes for SMILES strings using an ensemble model.""" - - # Instantiate ensemble model - ensemble = ENSEMBLES[ensemble_type]( - ensemble_config, - resolve_inconsistencies=resolve_inconsistencies, - verbose_output=verbose, - use_confidence=use_confidence, - ) - - # Collect SMILES strings from arguments and/or file - smiles_list = list(smiles) - if smiles_file: - with open(smiles_file, "r") as f: - smiles_list.extend([line.strip() for line in f if line.strip()]) - - if not smiles_list: - click.echo("No SMILES strings provided. Use --smiles or --smiles-file options.") + """Predict ChEBI classes for a list of SMILES / InChI strings.""" + molecules_list = read_molecules(molecules, molecule_file) + if not molecules_list: + click.echo("No molecules provided. Use --molecules or --molecule-file.") return - # Make predictions - predictions = ensemble.predict_smiles_list(smiles_list) + base_learners = build_base_learners(ensemble_config) + ensemble_model = ENSEMBLES[ensemble_type](ensemble_dir) - if output: - # save as json - import json - - with open(output, "w") as f: - json.dump( - {smiles: pred for smiles, pred in zip(smiles_list, predictions)}, - f, - indent=2, - ) + predictions = predict_molecules( + base_learners, + ensemble_model, + molecules_list, + prediction_cache_dir=prediction_cache_dir, + resolve_inconsistencies=resolve_inconsistencies, + decision_threshold=decision_threshold, + ) - else: - # Print results - for i, (smiles, prediction) in enumerate(zip(smiles_list, predictions)): - click.echo(f"Result for: {smiles}") - if prediction: - click.echo(f" Predicted classes: {', '.join(map(str, prediction))}") - else: - click.echo(" No predictions") + print(f"Predictions: {predictions}") + # TODO: turn the aggregated predictions into ChEBI classes per molecule, print them / save to output if __name__ == "__main__": - cli() + # cli() + load_dataset( + os.path.join("data", "chebi_v252", "ChEBI25", "processed"), split="validation" + ) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 38aa749..1424d71 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -1,304 +1,131 @@ -import importlib -import time from pathlib import Path -from typing import Union import torch -import tqdm import yaml +from torchmetrics import F1Score -from chebifier.check_env import check_package_installed from chebifier.ensemble.base_ensemble import BaseEnsemble -from chebifier.hugging_face import download_model_files -from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother -from chebifier.prediction_models.base_predictor import BasePredictor -from chebifier.utils import ( - get_default_configs, - get_disjoint_files, - load_chebi_graph, - process_config, -) class VotingEnsemble(BaseEnsemble): def __init__( self, - model_configs: Union[str, Path, dict, None] = None, - resolve_inconsistencies: bool = True, - verbose_output: bool = False, + ensemble_dir: str, use_confidence: bool = True, ): - # Deferred Import: To avoid circular import error - from chebifier.model_registry import MODEL_TYPES - - # Load configuration from YAML file - if not model_configs: - config = get_default_configs() - elif isinstance(model_configs, dict): - config = model_configs - else: - print(f"Loading ensemble configuration from {model_configs}") - with open(model_configs, "r") as f: - config = yaml.safe_load(f) - - with ( - importlib.resources.files("chebifier") - .joinpath("model_registry.yml") - .open("r") as f - ): - model_registry = yaml.safe_load(f) - - processed_configs = process_config(config, model_registry) - self.verbose_output = verbose_output + super().__init__(ensemble_dir) self.use_confidence = use_confidence + self.macro_f1 = None - self.chebi_graph = load_chebi_graph() - self.disjoint_files = get_disjoint_files() - - self.models = [] - self.positive_prediction_threshold = 0.5 - for model_name, model_config in processed_configs.items(): - model_cls = MODEL_TYPES[model_config["type"]] - if "hugging_face" in model_config: - hugging_face_kwargs = download_model_files(model_config["hugging_face"]) - else: - hugging_face_kwargs = {} - if "package_name" in model_config: - check_package_installed(model_config["package_name"]) + def find_best_threshold(self, predictions, val_labels_tensor): - model_instance = model_cls( - model_name, - **model_config, - **hugging_face_kwargs, - chebi_graph=self.chebi_graph, + best_threshold = 0.5 + best_f1 = 0.0 + for threshold in range(0, 100): + threshold_value = threshold / 100 + macro_f1_score = self.macro_f1( + predictions > threshold_value, val_labels_tensor ) - assert isinstance(model_instance, BasePredictor) - self.models.append(model_instance) - - if resolve_inconsistencies: - self.smoother = ScoreBasedPredictionSmoother( - self.chebi_graph, - label_names=None, - disjoint_files=self.disjoint_files, - verbose=self.verbose_output, - ) - else: - self.smoother = None - - def gather_predictions(self, smiles_list): - # get predictions from all models for the SMILES list - # order them alphabetically by label class - model_predictions = [] - predicted_classes = set() - for model in self.models: - model_predictions.append(model.predict_smiles_list(smiles_list)) - for logits_for_smiles in model_predictions[-1]: - if logits_for_smiles is not None: - for cls in logits_for_smiles: - predicted_classes.add(cls) - if self.verbose_output: - print(f"Sorting predictions from {len(model_predictions)} models...") - predicted_classes = sorted(list(predicted_classes)) - predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)} - ordered_logits = ( - torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) - * torch.nan + if macro_f1_score > best_f1: + best_f1 = macro_f1_score.item() + best_threshold = threshold_value + return best_threshold + + def calibrate(self, validation_predictions, validation_data, validation_labels): + self.macro_f1 = F1Score( + task="multilabel", num_labels=len(validation_labels), average="macro" ) - for i, model_prediction in enumerate(model_predictions): - for j, logits_for_smiles in tqdm.tqdm( - enumerate(model_prediction), - total=len(model_prediction), - desc=f"Sorting predictions for {self.models[i].model_name}", - ): - if logits_for_smiles is not None: - for cls in logits_for_smiles: - ordered_logits[j, predicted_classes_dict[cls], i] = ( - logits_for_smiles[cls] - ) - return ordered_logits, predicted_classes + prediction_thresholds = { + model_name: self.find_best_threshold(predictions, validation_labels) + for model_name, predictions in validation_predictions.items() + } + self._save_prediction_thresholds(prediction_thresholds) + + def _save_prediction_thresholds(self, thresholds: dict[str, float]): + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + with open(thresholds_path, "w+", encoding="utf-8") as f: + yaml.dump(thresholds, f) + + def _load_prediction_thresholds(self) -> dict[str, float]: + thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" + if thresholds_path.exists(): + with open(thresholds_path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + else: + raise FileNotFoundError( + f"Prediction thresholds file not found in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) - def consolidate_predictions( - self, - predictions, - classwise_weights, - return_intermediate_results=False, - **kwargs, - ): + def predict(self, test_predictions: dict[str, torch.Tensor]): """ Aggregates predictions from multiple models using weighted majority voting. - Optimized version using tensor operations instead of for loops. + weights are only the self-reported confidence (=difference between prediction and threshold). If set to false, all models are weighted equally. """ - num_smiles, num_classes, num_models = predictions.shape - + predictions_tensor = torch.stack( + list(test_predictions.values()), dim=2 + ) # Shape: (num_molecules, num_classes, num_models) # Get predictions for all classes - valid_predictions = ~torch.isnan(predictions) + valid_predictions = ~torch.isnan(predictions_tensor) valid_counts = valid_predictions.sum(dim=2) # Sum over models dimension + thresholds = self._load_prediction_thresholds() + if any(model_name not in thresholds for model_name in test_predictions.keys()): + raise ValueError( + "Prediction thresholds not found for all models. Please calibrate the ensemble first. Models missing thresholds: " + + ", ".join( + model_name + for model_name in test_predictions.keys() + if model_name not in thresholds + ) + ) + threshold_mask = torch.tensor( + [thresholds[model_name] for model_name in test_predictions.keys()], + dtype=predictions_tensor.dtype, + device=predictions_tensor.device, + ) + # Skip classes with no valid predictions has_valid_predictions = valid_counts > 0 # Calculate positive and negative predictions for all classes at once positive_mask = ( - predictions > self.positive_prediction_threshold + predictions_tensor > threshold_mask.unsqueeze(0).unsqueeze(0) ) & valid_predictions negative_mask = ( - predictions < self.positive_prediction_threshold + predictions_tensor < threshold_mask.unsqueeze(0).unsqueeze(0) ) & valid_predictions - # if use_confidence is passed in kwargs, it overrides the ensemble setting - use_confidence = kwargs.get("use_confidence", self.use_confidence) - if use_confidence: + if self.use_confidence: confidence = 2 * torch.abs( - predictions.nan_to_num() - self.positive_prediction_threshold + predictions_tensor.nan_to_num() + - threshold_mask.unsqueeze(0).unsqueeze(0) ) else: - confidence = torch.ones_like(predictions) - - # Extract positive and negative weights - pos_weights = classwise_weights[0] # Shape: (num_classes, num_models) - neg_weights = classwise_weights[1] # Shape: (num_classes, num_models) + confidence = torch.ones_like(predictions_tensor) # Calculate weighted predictions using broadcasting - # predictions shape: (num_smiles, num_classes, num_models) + # predictions shape: (num_molecules, num_classes, num_models) # weights shape: (num_classes, num_models) - positive_weighted = ( - positive_mask.float() * confidence * pos_weights.unsqueeze(0) - ) - negative_weighted = ( - negative_mask.float() * confidence * neg_weights.unsqueeze(0) - ) + positive_weighted = positive_mask.float() * confidence + negative_weighted = negative_mask.float() * confidence # Sum over models dimension - positive_sum = positive_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - negative_sum = negative_weighted.sum(dim=2) # Shape: (num_smiles, num_classes) - - # Determine which classes to include for each SMILES - net_score = positive_sum - negative_sum # Shape: (num_smiles, num_classes) - if return_intermediate_results: - return ( - net_score, - has_valid_predictions, - { - "positive_mask": positive_mask, - "negative_mask": negative_mask, - "confidence": confidence, - "positive_sum": positive_sum, - "negative_sum": negative_sum, - }, - ) - - return net_score, has_valid_predictions - - def apply_inconsistency_resolution( - self, net_score, class_names, has_valid_predictions - ): - # Smooth predictions - start_time = time.perf_counter() - if self.smoother is not None: - self.smoother.set_label_names(class_names) - smooth_net_score = self.smoother(net_score) - class_decisions = ( - smooth_net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - else: - class_decisions = ( - net_score > 0 - ) & has_valid_predictions # Shape: (num_smiles, num_classes) - end_time = time.perf_counter() - if self.verbose_output: - print(f"Prediction smoothing took {end_time - start_time:.2f} seconds") - - complete_failure = torch.all(~has_valid_predictions, dim=1) - return class_decisions, complete_failure - - def calculate_classwise_weights(self, predicted_classes): - """No weights, simple majority voting""" - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - - return positive_weights, negative_weights - - def predict_smiles_list( - self, smiles_list, return_intermediate_results=False, **kwargs - ) -> list: - ordered_predictions, predicted_classes = self.gather_predictions(smiles_list) - if len(predicted_classes) == 0: - print("Warning: No classes have been predicted for the given SMILES list.") - predicted_classes = {cls: i for i, cls in enumerate(predicted_classes)} - - classwise_weights = self.calculate_classwise_weights(predicted_classes) - if return_intermediate_results: - net_score, has_valid_predictions, intermediate_results_dict = ( - self.consolidate_predictions( - ordered_predictions, - classwise_weights, - return_intermediate_results=return_intermediate_results, - ) - ) - else: - net_score, has_valid_predictions = self.consolidate_predictions( - ordered_predictions, classwise_weights - ) - class_decisions, is_failure = self.apply_inconsistency_resolution( - net_score, list(predicted_classes.keys()), has_valid_predictions - ) - - class_names = list(predicted_classes.keys()) - class_indices = {predicted_classes[cls]: cls for cls in class_names} - result = [ - ( - [ - class_indices[idx.item()] - for idx in torch.nonzero(i, as_tuple=True)[0] - ] - if not failure - else None - ) - for i, failure in zip(class_decisions, is_failure) - ] - if return_intermediate_results: - intermediate_results_dict["predicted_classes"] = predicted_classes - intermediate_results_dict["classwise_weights"] = classwise_weights - intermediate_results_dict["net_score"] = net_score - return result, intermediate_results_dict - - return result - - -if __name__ == "__main__": - ensemble = VotingEnsemble( - { - "resgated_0ps1g189": { - "type": "resgated", - "ckpt_path": "data/0ps1g189/epoch=122.ckpt", - "molecular_properties": [ - "chebai_graph.preprocessing.properties.AtomType", - "chebai_graph.preprocessing.properties.NumAtomBonds", - "chebai_graph.preprocessing.properties.AtomCharge", - "chebai_graph.preprocessing.properties.AtomAromaticity", - "chebai_graph.preprocessing.properties.AtomHybridization", - "chebai_graph.preprocessing.properties.AtomNumHs", - "chebai_graph.preprocessing.properties.BondType", - "chebai_graph.preprocessing.properties.BondInRing", - "chebai_graph.preprocessing.properties.BondAromaticity", - "chebai_graph.preprocessing.properties.RDKit2DNormalized", - ], - # "classwise_weights_path" : "../python-chebai/metrics_0ps1g189_80-10-10.json" - }, - "electra_14ko0zcf": { - "type": "electra", - "ckpt_path": "data/14ko0zcf/epoch=193.ckpt", - # "classwise_weights_path": "../python-chebai/metrics_electra_14ko0zcf_80-10-10.json", - }, + positive_sum = positive_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + negative_sum = negative_weighted.sum( + dim=2 + ) # Shape: (num_molecules, num_classes) + + # Determine which classes to include for each molecule + net_score = positive_sum - negative_sum # Shape: (num_molecules, num_classes) + return { + "net_score": net_score, + "has_valid_predictions": has_valid_predictions, + "positive_sum": positive_sum, + "negative_sum": negative_sum, + "confidence": confidence, + "positive_mask": positive_mask, + "negative_mask": negative_mask, } - ) - r = ensemble.predict_smiles_list( - [ - "[NH3+]CCCC[C@H](NC(=O)[C@@H]([NH3+])CC([O-])=O)C([O-])=O", - "C[C@H](N)C(=O)NCC(O)=O#", - "", - ], - load_preds_if_possible=False, - ) - print(len(r), r[0]) diff --git a/chebifier/predict.py b/chebifier/predict.py index 9ff2530..c4c6811 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -1 +1,145 @@ -# Get end-to-end predictions (from SMILES via base learners + ensemble + inconsistency resolution to ChEBI classes) +# Get end-to-end predictions (from SMILES / molecule list via base learners + ensemble + inconsistency resolution to ChEBI classes) + + +import os +from typing import Optional + +import torch +from rdkit import Chem + +from chebifier.ensemble.base_ensemble import BaseEnsemble +from chebifier.inconsistency_resolution import ScoreBasedPredictionSmoother +from chebifier.prediction_models.base_predictor import BasePredictor +from chebifier.utils import get_disjoint_files, load_chebi_graph + + +def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions): + smoother.set_label_names(class_names) + smooth_net_score = smoother(aggregated_predictions["net_score"]) + aggregated_predictions["net_score"] = smooth_net_score + return aggregated_predictions + + +def collect_base_learner_predictions( + predictions: dict[str, list[dict | None]], +) -> (dict[str, torch.Tensor], list[str]): + """ + Collect predictions from base learners into a single dictionary. + + Args: + predictions (dict): A dictionary where keys are model names and values are lists of predictions. + Assumes those lists have the same length and each entry in the list is either None or a dict + mapping class labels to predicted values. + + Returns: + dict: A dictionary where keys are model names and values are tensors of predictions with shape (num_samples, num_classes). + If a prediction is None, it will be replaced with NaN. + ensemble_classes (list): A list of class labels that are present in the predictions. + """ + collected_predictions = {} + ensemble_classes = set() + n_samples = -1 + # step 1: collect classes + for model_name, model_predictions in predictions.items(): + for pred in model_predictions: + if pred is not None: + ensemble_classes.update(pred.keys()) + if n_samples == -1: + n_samples = len(model_predictions) + else: + assert n_samples == len( + model_predictions + ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." + ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} + # step 2: map predictions to tensors + for model_name, model_predictions in predictions.items(): + # Replace None values with NaN + model_predictions = torch.zeros((n_samples, len(ensemble_classes))) * float( + "nan" + ) + for i, pred in enumerate(model_predictions): + if pred is not None: + for cls, value in pred.items(): + model_predictions[i, cls_to_idx[cls]] = value + collected_predictions[model_name] = model_predictions + + return collected_predictions, list(ensemble_classes) + + +def predict( + base_learners: dict[str, BasePredictor], + ensemble_model: BaseEnsemble, + molecules: list[str | Chem.Mol], + prediction_cache_dir: Optional[str] = None, + resolve_inconsistencies: bool = True, + decision_threshold: float = 0, +) -> dict: + """ + Get end-to-end predictions from base learners and an ensemble model. + + Args: + base_learners (dict[str, BasePredictor]): A dictionary of base learner models. + ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model. + molecules (list[str | Chem.Mol]): List of molecules for prediction (either SMILES strings or molecule objects). + prediction_cache_dir (Optional[str]): Directory to cache predictions. If None, no caching is performed. If provided, + predictions from base learners will be cached to avoid recomputation (warning: not checked against the molecules provided + -> if the molecules change, you have to empty the cache or provide a new cache directory). + resolve_inconsistencies (bool): Whether to resolve inconsistencies in the aggregated predictions. + decision_threshold (float): Threshold for class decisions based on net score. Default is 0. + + Returns: + dict: A dictionary containing the final predictions and optionally the smoothed predictions. + """ + + # Step 1: Get predictions from base learners on test data + test_predictions = {} + for model_name, model in base_learners.items(): + if prediction_cache_dir is None: + test_predictions[model_name] = model.predict_list(molecules) + else: + cache_path = os.path.join( + prediction_cache_dir, f"{model_name}_test_predictions.pt" + ) + if os.path.exists(cache_path): + test_predictions[model_name] = torch.load( + cache_path, weights_only=False + ) + else: + test_predictions[model_name] = model.predict_list(molecules) + torch.save(test_predictions[model_name], cache_path) + + test_predictions, predicted_classes = collect_base_learner_predictions( + test_predictions + ) + + # Step 2: Get aggregated predictions from the ensemble model + aggregated_predictions = ensemble_model.predict(test_predictions) + # net_score, has_valid_predictions, intermediate_results_dict + + # Step 3: Optionally resolve inconsistencies in the aggregated predictions + if resolve_inconsistencies: + chebi_graph = load_chebi_graph() + disjoint_files = get_disjoint_files() + smoother = ScoreBasedPredictionSmoother( + chebi_graph=chebi_graph, label_names=None, disjoint_files=disjoint_files + ) + aggregated_predictions = apply_inconsistency_resolution( + smoother, predicted_classes, aggregated_predictions + ) + + class_decisions = ( + aggregated_predictions["net_score"] > decision_threshold + ) & aggregated_predictions[ + "has_valid_predictions" + ] # Shape: (num_smiles, num_classes) + + complete_failure = torch.all( + ~aggregated_predictions["has_valid_predictions"], dim=1 + ) + aggregated_predictions["class_decisions"] = class_decisions + aggregated_predictions["complete_failure"] = complete_failure + + aggregated_predictions["predicted_classes"] = predicted_classes + + return aggregated_predictions diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 1082b8e..806cd6f 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -18,10 +18,10 @@ def __init__( self._description = kwargs.get("description", None) @modelwise_smiles_lru_cache.batch_decorator - def predict_list(self, molecule_list: list[str | Chem.Mol]) -> dict: + def predict_list(self, molecule_list: list[str | Chem.Mol]) -> list[dict | None]: raise NotImplementedError() - def predict(self, molecule: str | Chem.Mol) -> dict: + def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction return self.predict_list([molecule])[0] diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 2bfd389..cf48658 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -21,7 +21,7 @@ def __init__( super().__init__(model_name, **kwargs) self.batch_size = kwargs.get("batch_size", None) # compile_model will run the model in eager mode, which gives better performance, but does not return intermediate states - # such as attention weights. Therfore, ELECTRA attention graphs will only work with compile_model=False. + # such as attention weights. Therefore, ELECTRA attention graphs will only work with compile_model=False. compile_model = kwargs.get("compile_model", True) # If batch_size is not provided, it will be set to default batch size used during training in Predictor self.predictor: Predictor = Predictor( From 8b48598e304e65e2c10f1d398521001de2f3cc32 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 15:09:52 +0200 Subject: [PATCH 04/10] bug fixes and optimized performance for aggregation --- chebifier/build_ensemble.py | 19 +++++- chebifier/cli.py | 5 +- chebifier/ensemble/voting_ensemble.py | 14 +++- chebifier/predict.py | 70 +++++++++++++++++--- chebifier/prediction_models/c3p_predictor.py | 26 ++++++-- 5 files changed, 106 insertions(+), 28 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index 5a9df68..23de772 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -13,7 +13,8 @@ class EnsembleBuilder: base_learners (dict[str, BasePredictor]): A dictionary of base learner models. ensemble_model (BaseEnsemble): An instance of a BaseEnsemble model. validation_data (list[Chem.Mol]): Validation data for calibration. - validation_labels (torch.Tensor): Validation labels for calibration. + validation_labels (pd.DataFrame): Validation labels for calibration, one column per class. + The column names define the label set the base learner predictions are mapped onto. prediction_cache_dir (str): Directory to cache predictions. """ @@ -48,22 +49,34 @@ def build_ensemble(self): self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" ) if os.path.exists(cache_path): + print(f"{model_name} validation predictions found in cache, loading...") validation_predictions[model_name] = torch.load( cache_path, weights_only=False ) else: + print(f"Computing {model_name} validation predictions...") validation_predictions[model_name] = model.predict_list( self.validation_data ) torch.save(validation_predictions[model_name], cache_path) + # Base learners may be trained on different label sets (e.g. ChEBI25 vs. ChEBI25_3_STAR), + # so their union does not match the labels we calibrate against. Map every base learner + # onto the label set of the validation data instead. + label_classes = [str(cls) for cls in self.validation_labels.columns] validation_predictions, classes = collect_base_learner_predictions( - validation_predictions + validation_predictions, classes=label_classes + ) + validation_labels = torch.from_numpy( + self.validation_labels.to_numpy(dtype=bool) ) + print( + f"Collected validation predictions from {len(validation_predictions)} base learners with {len(classes)} unique classes. Calibrating ensemble model..." + ) # Step 2: Calibrate the ensemble model using validation predictions self.ensemble_model.calibrate( - validation_predictions, self.validation_data, self.validation_labels + validation_predictions, self.validation_data, validation_labels ) return self.ensemble_model diff --git a/chebifier/cli.py b/chebifier/cli.py index be9bbd7..49d345b 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -295,7 +295,4 @@ def predict( if __name__ == "__main__": - # cli() - load_dataset( - os.path.join("data", "chebi_v252", "ChEBI25", "processed"), split="validation" - ) + cli() diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 1424d71..8a224ee 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -18,7 +18,9 @@ def __init__( self.macro_f1 = None def find_best_threshold(self, predictions, val_labels_tensor): - + print( + f"Finding best threshold for predictions with shape {predictions.shape} and validation labels with shape {val_labels_tensor.shape}" + ) best_threshold = 0.5 best_f1 = 0.0 for threshold in range(0, 100): @@ -32,10 +34,15 @@ def find_best_threshold(self, predictions, val_labels_tensor): return best_threshold def calibrate(self, validation_predictions, validation_data, validation_labels): + print( + f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." + ) self.macro_f1 = F1Score( - task="multilabel", num_labels=len(validation_labels), average="macro" + task="multilabel", num_labels=validation_labels.shape[1], average="macro" + ) + print( + f"Validation labels: {validation_labels.shape}, Validation predictions: {validation_predictions[list(validation_predictions.keys())[0]].shape}" ) - prediction_thresholds = { model_name: self.find_best_threshold(predictions, validation_labels) for model_name, predictions in validation_predictions.items() @@ -46,6 +53,7 @@ def _save_prediction_thresholds(self, thresholds: dict[str, float]): thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" with open(thresholds_path, "w+", encoding="utf-8") as f: yaml.dump(thresholds, f) + print(f"Saved prediction thresholds to {thresholds_path}: {thresholds}") def _load_prediction_thresholds(self) -> dict[str, float]: thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" diff --git a/chebifier/predict.py b/chebifier/predict.py index c4c6811..fb0726f 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -4,6 +4,7 @@ import os from typing import Optional +import numpy as np import torch from rdkit import Chem @@ -22,6 +23,7 @@ def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions def collect_base_learner_predictions( predictions: dict[str, list[dict | None]], + classes: Optional[list[str]] = None, ) -> (dict[str, torch.Tensor], list[str]): """ Collect predictions from base learners into a single dictionary. @@ -30,6 +32,12 @@ def collect_base_learner_predictions( predictions (dict): A dictionary where keys are model names and values are lists of predictions. Assumes those lists have the same length and each entry in the list is either None or a dict mapping class labels to predicted values. + classes (Optional[list]): Column space to map the predictions onto. If None (the default), the + union of all classes reported by the base learners is used. Pass an explicit list to align + the predictions with a fixed label set (e.g. the labels of an evaluation dataset) - base + learners may be trained on different label sets, so the union does not necessarily match the + labels you want to compare against. Classes predicted by a base learner but missing from + `classes` are dropped, classes that no base learner covers stay NaN. Returns: dict: A dictionary where keys are model names and values are tensors of predictions with shape (num_samples, num_classes). @@ -39,30 +47,70 @@ def collect_base_learner_predictions( collected_predictions = {} ensemble_classes = set() n_samples = -1 + print(f"Collecting base learner predictions from {len(predictions)} models...") # step 1: collect classes + # Base learners typically return the same label set for every sample, so we only + # touch the class set when the label set actually changes from one sample to the next. for model_name, model_predictions in predictions.items(): - for pred in model_predictions: - if pred is not None: - ensemble_classes.update(pred.keys()) + if classes is None: + previous_keys = None + for pred in model_predictions: + if pred: + keys = tuple(pred) + if keys != previous_keys: + ensemble_classes.update(keys) + previous_keys = keys if n_samples == -1: n_samples = len(model_predictions) else: assert n_samples == len( model_predictions ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." - ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + if classes is None: + ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + else: + ensemble_classes = list(classes) cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} # step 2: map predictions to tensors + # Filled row-wise via numpy fancy indexing: one vectorised write per sample instead + # of one Python-level tensor assignment per predicted class. The column indices are + # reused as long as consecutive samples share the same label set (see step 1). for model_name, model_predictions in predictions.items(): - # Replace None values with NaN - model_predictions = torch.zeros((n_samples, len(ensemble_classes))) * float( - "nan" + # Samples without a prediction (and classes the model does not cover) stay NaN + predictions_array = np.full( + (n_samples, len(ensemble_classes)), np.nan, dtype=np.float32 ) + previous_keys = None + columns = None + # positions of the kept values within pred.values(), None if nothing is dropped + kept_positions = None for i, pred in enumerate(model_predictions): - if pred is not None: - for cls, value in pred.items(): - model_predictions[i, cls_to_idx[cls]] = value - collected_predictions[model_name] = model_predictions + if pred: + keys = tuple(pred) + if keys != previous_keys: + kept = [ + (position, cls_to_idx[cls]) + for position, cls in enumerate(keys) + if cls in cls_to_idx + ] + columns = np.fromiter( + (column for _, column in kept), dtype=np.intp, count=len(kept) + ) + kept_positions = ( + None + if len(kept) == len(keys) + else np.fromiter( + (position for position, _ in kept), + dtype=np.intp, + count=len(kept), + ) + ) + previous_keys = keys + values = np.fromiter(pred.values(), dtype=np.float32, count=len(pred)) + predictions_array[i, columns] = ( + values if kept_positions is None else values[kept_positions] + ) + collected_predictions[model_name] = torch.from_numpy(predictions_array) return collected_predictions, list(ensemble_classes) diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index f0a1cfd..ec59748 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/chebifier/prediction_models/c3p_predictor.py @@ -42,17 +42,29 @@ def predict_list(self, smiles_list: list[str]) -> list: ) ) + # Look up the position of each SMILES via a dict instead of scanning smiles_list + # for every result (C3P returns one result per class and molecule, so the scan + # made reformatting quadratic in the number of molecules). Repeated SMILES map to + # all of their positions, which list.index could not do (it always returned the + # first one, leaving the later rows without any predictions). + indices_by_smiles: dict[str, list[int]] = {} + for idx, smiles in enumerate(smiles_list): + indices_by_smiles.setdefault(smiles, []).append(idx) + result_reformatted = [dict() for _ in range(len(smiles_list))] for result in tqdm.tqdm(result_list, desc="Reformatting C3P results"): chebi_id = result.class_id.split(":")[1] - result_reformatted[smiles_list.index(result.input_smiles)][ - chebi_id - ] = result.is_match if result.is_match and self.chebi_graph is not None: - for parent in list(self.chebi_graph.predecessors(chebi_id)): - result_reformatted[smiles_list.index(result.input_smiles)][ - str(parent) - ] = 1 + parents = [ + str(parent) for parent in self.chebi_graph.predecessors(chebi_id) + ] + else: + parents = [] + for idx in indices_by_smiles[result.input_smiles]: + preds_i = result_reformatted[idx] + preds_i[chebi_id] = result.is_match + for parent in parents: + preds_i[parent] = 1 return result_reformatted def explain_smiles(self, smiles): From deb681ac5b0dba1f38a5e7d7d6a8005757f448db Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 15:54:00 +0200 Subject: [PATCH 05/10] integrate wmv-f1 ensemble into new workflow --- chebifier/ensemble/voting_ensemble.py | 25 ++- .../ensemble/weighted_majority_ensemble.py | 168 ++++++------------ chebifier/model_registry.py | 6 +- 3 files changed, 67 insertions(+), 132 deletions(-) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 8a224ee..6b23d5f 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -15,19 +15,16 @@ def __init__( ): super().__init__(ensemble_dir) self.use_confidence = use_confidence - self.macro_f1 = None + self.classwise_f1 = None def find_best_threshold(self, predictions, val_labels_tensor): - print( - f"Finding best threshold for predictions with shape {predictions.shape} and validation labels with shape {val_labels_tensor.shape}" - ) best_threshold = 0.5 best_f1 = 0.0 for threshold in range(0, 100): threshold_value = threshold / 100 - macro_f1_score = self.macro_f1( + macro_f1_score = self.classwise_f1( predictions > threshold_value, val_labels_tensor - ) + ).mean() if macro_f1_score > best_f1: best_f1 = macro_f1_score.item() best_threshold = threshold_value @@ -37,11 +34,8 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): print( f"Calibrating {self.ensemble_name} with {len(validation_predictions)} base learners..." ) - self.macro_f1 = F1Score( - task="multilabel", num_labels=validation_labels.shape[1], average="macro" - ) - print( - f"Validation labels: {validation_labels.shape}, Validation predictions: {validation_predictions[list(validation_predictions.keys())[0]].shape}" + self.classwise_f1 = F1Score( + task="multilabel", num_labels=validation_labels.shape[1], average=None ) prediction_thresholds = { model_name: self.find_best_threshold(predictions, validation_labels) @@ -65,6 +59,10 @@ def _load_prediction_thresholds(self) -> dict[str, float]: f"Prediction thresholds file not found in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." ) + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # No trust for MV, only used in WMV + return 1 + def predict(self, test_predictions: dict[str, torch.Tensor]): """ Aggregates predictions from multiple models using weighted majority voting. @@ -112,11 +110,12 @@ def predict(self, test_predictions: dict[str, torch.Tensor]): else: confidence = torch.ones_like(predictions_tensor) + trust = self.calculate_trust(test_predictions) # Calculate weighted predictions using broadcasting # predictions shape: (num_molecules, num_classes, num_models) # weights shape: (num_classes, num_models) - positive_weighted = positive_mask.float() * confidence - negative_weighted = negative_mask.float() * confidence + positive_weighted = positive_mask.float() * confidence * trust + negative_weighted = negative_mask.float() * confidence * trust # Sum over models dimension positive_sum = positive_weighted.sum( diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 59d0c32..495fa18 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,134 +1,74 @@ -import json -import os +from pathlib import Path import torch from chebifier.ensemble.voting_ensemble import VotingEnsemble -class WMVwithPPVNPVEnsemble(VotingEnsemble): - - def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=1, **kwargs - ): - """WMV ensemble that weights models based on their class-wise positive / negative predictive values. For each class, the weight is calculated as: - weight = (weighting_strength * PPV + (1 - weighting_strength)) ** weighting_exponent - where PPV is the class-specific positive predictive value of the model on the validation set - or (if the prediction is negative): - weight = (weighting_strength * NPV + (1 - weighting_strength)) ** weighting_exponent - where NPV is the class-specific negative predictive value of the model on the validation set. - """ - super().__init__(config_path, **kwargs) - self.weighting_strength = weighting_strength - self.weighting_exponent = weighting_exponent - - self.model_classwise_weights = dict() - for model in self.models: - classwise_weights_path = os.path.join( - self.ensemble_dir, f"{model.model_name}_classwise_weights.json" - ) - if os.path.exists(classwise_weights_path): - self.model_classwise_weights[model.model_name] = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.model_classwise_weights[model.model_name] = None - - def calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by the class-specific positive / negative weight (default 1). - """ - positive_weights = torch.ones(len(predicted_classes), len(self.models)) - negative_weights = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - positive_weights[:, j] *= model.model_weight - negative_weights[:, j] *= model.model_weight - if self.model_classwise_weights[model.model_name] is None: - continue - for cls, weights in self.model_classwise_weights[model.model_name].items(): - if cls not in predicted_classes: - continue - ppv = ( - weights["TP"] / (weights["TP"] + weights["FP"]) - if (weights["TP"] + weights["FP"]) > 0 - else 1.0 - ) - npv = ( - weights["TN"] / (weights["TN"] + weights["FN"]) - if (weights["TN"] + weights["FN"]) > 0 - else 1.0 - ) - positive_weights[predicted_classes[cls], j] *= ( - ppv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - negative_weights[predicted_classes[cls], j] *= ( - npv * self.weighting_strength + (1 - self.weighting_strength) - ) ** self.weighting_exponent - - if self.verbose_output: - print( - "Calculated model weightings. The averages for positive / negative weights are:" - ) - for i, model in enumerate(self.models): - print( - f"{model.model_name}: {positive_weights[:, i].mean().item():.3f} / {negative_weights[:, i].mean().item():.3f}" - ) - - return positive_weights, negative_weights - - class WMVwithF1Ensemble(VotingEnsemble): def __init__( - self, config_path=None, weighting_strength=1, weighting_exponent=6.25, **kwargs + self, + ensemble_dir: str, + use_confidence: bool = True, + weighting_strength=1, + weighting_exponent=1, + **kwargs, ): """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent where F1 is the class-specific F1 score ("trust") of the model on the validation set. """ - super().__init__(config_path, **kwargs) + super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent - self.model_classwise_weights = dict() - for model in self.models: - classwise_weights_path = os.path.join( - self.ensemble_dir, f"{model.model_name}_classwise_weights.json" + def calibrate(self, validation_predictions, validation_data, validation_labels): + super().calibrate(validation_predictions, validation_data, validation_labels) + self._save_classwise_f1(validation_predictions, validation_labels) + + def _save_classwise_f1(self, validation_predictions, validation_labels): + thresholds = self._load_prediction_thresholds() + for model_name, predictions in validation_predictions.items(): + f1 = self.classwise_f1( + predictions > thresholds[model_name], validation_labels + ) + f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + with open(f1_path, "w+") as f: + f.write("\n".join(f1.tolist())) + print( + f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) - if os.path.exists(classwise_weights_path): - self.model_classwise_weights[model.model_name] = json.load( - open(classwise_weights_path, encoding="utf-8") - ) - else: - self.model_classwise_weights[model.model_name] = None - def calculate_classwise_weights(self, predicted_classes): - """ - Given the positions of predicted classes in the predictions tensor, assign weights to each class. The - result is two tensors of shape (num_predicted_classes, num_models). The weight for each class is the model_weight - (default: 1) multiplied by (1 + the class-specific validation-f1 (default 1)). - """ - weights_by_cls = torch.ones(len(predicted_classes), len(self.models)) - for j, model in enumerate(self.models): - weights_by_cls[:, j] *= model.model_weight - if self.model_classwise_weights[model.model_name] is None: - continue - for cls, weights in self.model_classwise_weights[model.model_name].items(): - if cls in predicted_classes: - if (2 * weights["TP"] + weights["FP"] + weights["FN"]) > 0: - f1 = ( - 2 - * weights["TP"] - / (2 * weights["TP"] + weights["FP"] + weights["FN"]) - ) - weights_by_cls[predicted_classes[cls], j] *= ( - self.weighting_strength * f1 + 1 - self.weighting_strength - ) ** self.weighting_exponent - if self.verbose_output: - print("Calculated model weightings. The average weights are:") - for i, model in enumerate(self.models): - print(f"{model.model_name}: {weights_by_cls[:, i].mean().item():.3f}") + def _load_classwise_f1(self, model_name: str) -> torch.Tensor: + classwise_f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" + if classwise_f1_path.exists(): + with open(classwise_f1_path, "r", encoding="utf-8") as f: + return torch.tensor([float(x) for x in f.read().splitlines()]) + else: + raise FileNotFoundError( + f"Class-wise F1 scores file not found for model {model_name} in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." + ) - return weights_by_cls, weights_by_cls + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: + # Calculate trust based on class-wise F1 scores for each model + # target shape: (num_molecules, num_classes, num_models) + num_models = len(predictions) + num_molecules = list(predictions.values())[0].shape[0] + num_classes = list(predictions.values())[0].shape[1] + trust_tensor = torch.ones( + (num_molecules, num_classes, num_models), dtype=torch.float32 + ) + for model_idx, (model_name, prediction_tensor) in enumerate( + predictions.items() + ): + classwise_f1 = self._load_classwise_f1(model_name) + assert ( + classwise_f1.shape[0] == num_classes + ), f"Class-wise F1 scores for model {model_name} do not match number of classes in predictions." + # Expand classwise_f1 to match the shape of trust_tensor for broadcasting + classwise_f1 = classwise_f1.unsqueeze(0).expand(num_molecules, -1) + trust_tensor[:, :, model_idx] = ( + self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) + ) ** self.weighting_exponent + return trust_tensor diff --git a/chebifier/model_registry.py b/chebifier/model_registry.py index ef13e7e..8d737de 100644 --- a/chebifier/model_registry.py +++ b/chebifier/model_registry.py @@ -1,8 +1,5 @@ from chebifier.ensemble.voting_ensemble import VotingEnsemble -from chebifier.ensemble.weighted_majority_ensemble import ( - WMVwithF1Ensemble, - WMVwithPPVNPVEnsemble, -) +from chebifier.ensemble.weighted_majority_ensemble import WMVwithF1Ensemble from chebifier.prediction_models import ( ChEBILookupPredictor, ChemlogPeptidesPredictor, @@ -20,7 +17,6 @@ ENSEMBLES = { "mv": VotingEnsemble, - "wmv-ppvnpv": WMVwithPPVNPVEnsemble, "wmv-f1": WMVwithF1Ensemble, } From fbafc8e8ae46a6965511ab408fd6c0524c9dd514 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 16:42:48 +0200 Subject: [PATCH 06/10] optimize prediction storage --- chebifier/build_ensemble.py | 16 ++- .../ensemble/weighted_majority_ensemble.py | 2 +- chebifier/predict.py | 128 ++++++++---------- chebifier/prediction_models/base_predictor.py | 37 +++++ chebifier/prediction_models/nn_predictor.py | 14 +- 5 files changed, 114 insertions(+), 83 deletions(-) diff --git a/chebifier/build_ensemble.py b/chebifier/build_ensemble.py index 23de772..64c6b0c 100644 --- a/chebifier/build_ensemble.py +++ b/chebifier/build_ensemble.py @@ -2,7 +2,11 @@ import torch -from chebifier.predict import collect_base_learner_predictions +from chebifier.predict import ( + collect_base_learner_predictions, + load_dense_predictions, + save_dense_predictions, +) class EnsembleBuilder: @@ -46,19 +50,17 @@ def build_ensemble(self): # get cached predictions if available, otherwise compute and cache them for model_name, model in self.base_learners.items(): cache_path = os.path.join( - self.prediction_cache_dir, f"{model_name}_validation_predictions.pt" + self.prediction_cache_dir, f"{model_name}_validation_predictions.npz" ) if os.path.exists(cache_path): print(f"{model_name} validation predictions found in cache, loading...") - validation_predictions[model_name] = torch.load( - cache_path, weights_only=False - ) + validation_predictions[model_name] = load_dense_predictions(cache_path) else: print(f"Computing {model_name} validation predictions...") - validation_predictions[model_name] = model.predict_list( + validation_predictions[model_name] = model.predict_dense( self.validation_data ) - torch.save(validation_predictions[model_name], cache_path) + save_dense_predictions(cache_path, *validation_predictions[model_name]) # Base learners may be trained on different label sets (e.g. ChEBI25 vs. ChEBI25_3_STAR), # so their union does not match the labels we calibrate against. Map every base learner diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 495fa18..ff51020 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -35,7 +35,7 @@ def _save_classwise_f1(self, validation_predictions, validation_labels): ) f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" with open(f1_path, "w+") as f: - f.write("\n".join(f1.tolist())) + f.writelines(f"{x}\n" for x in f1.tolist()) print( f"Saved class-wise F1 scores to {f1_path}: {len(f1.tolist())} classes (macro-f1: {f1.mean().item():.4f})." ) diff --git a/chebifier/predict.py b/chebifier/predict.py index fb0726f..18a6fa8 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -21,17 +21,26 @@ def apply_inconsistency_resolution(smoother, class_names, aggregated_predictions return aggregated_predictions +def save_dense_predictions(path: str, classes: list[str], scores: np.ndarray) -> None: + np.savez_compressed(path, classes=np.array(classes), scores=scores) + + +def load_dense_predictions(path: str) -> tuple[list[str], np.ndarray]: + with np.load(path) as data: + return [str(cls) for cls in data["classes"]], data["scores"] + + def collect_base_learner_predictions( - predictions: dict[str, list[dict | None]], + predictions: dict[str, tuple[list[str], np.ndarray]], classes: Optional[list[str]] = None, ) -> (dict[str, torch.Tensor], list[str]): """ Collect predictions from base learners into a single dictionary. Args: - predictions (dict): A dictionary where keys are model names and values are lists of predictions. - Assumes those lists have the same length and each entry in the list is either None or a dict - mapping class labels to predicted values. + predictions (dict): A dictionary where keys are model names and values are + (class labels, score matrix) pairs as returned by BasePredictor.predict_dense. + Assumes all score matrices have the same number of rows. classes (Optional[list]): Column space to map the predictions onto. If None (the default), the union of all classes reported by the base learners is used. Pass an explicit list to align the predictions with a fixed label set (e.g. the labels of an evaluation dataset) - base @@ -41,78 +50,51 @@ def collect_base_learner_predictions( Returns: dict: A dictionary where keys are model names and values are tensors of predictions with shape (num_samples, num_classes). - If a prediction is None, it will be replaced with NaN. + If a prediction is missing, it will be NaN. ensemble_classes (list): A list of class labels that are present in the predictions. """ - collected_predictions = {} - ensemble_classes = set() - n_samples = -1 print(f"Collecting base learner predictions from {len(predictions)} models...") - # step 1: collect classes - # Base learners typically return the same label set for every sample, so we only - # touch the class set when the label set actually changes from one sample to the next. - for model_name, model_predictions in predictions.items(): - if classes is None: - previous_keys = None - for pred in model_predictions: - if pred: - keys = tuple(pred) - if keys != previous_keys: - ensemble_classes.update(keys) - previous_keys = keys + n_samples = -1 + for model_name, (_, scores) in predictions.items(): if n_samples == -1: - n_samples = len(model_predictions) + n_samples = scores.shape[0] else: - assert n_samples == len( - model_predictions - ), f"All prediction lists must have the same length. Model {model_name} has {len(model_predictions)} predictions, expected {n_samples}." + assert ( + n_samples == scores.shape[0] + ), f"All prediction matrices must have the same length. Model {model_name} has {scores.shape[0]} predictions, expected {n_samples}." + if classes is None: - ensemble_classes = sorted(ensemble_classes) # Sort for consistent ordering + ensemble_classes = sorted( + {cls for model_classes, _ in predictions.values() for cls in model_classes} + ) else: ensemble_classes = list(classes) cls_to_idx = {cls: idx for idx, cls in enumerate(ensemble_classes)} - # step 2: map predictions to tensors - # Filled row-wise via numpy fancy indexing: one vectorised write per sample instead - # of one Python-level tensor assignment per predicted class. The column indices are - # reused as long as consecutive samples share the same label set (see step 1). - for model_name, model_predictions in predictions.items(): - # Samples without a prediction (and classes the model does not cover) stay NaN - predictions_array = np.full( - (n_samples, len(ensemble_classes)), np.nan, dtype=np.float32 - ) - previous_keys = None - columns = None - # positions of the kept values within pred.values(), None if nothing is dropped - kept_positions = None - for i, pred in enumerate(model_predictions): - if pred: - keys = tuple(pred) - if keys != previous_keys: - kept = [ - (position, cls_to_idx[cls]) - for position, cls in enumerate(keys) - if cls in cls_to_idx - ] - columns = np.fromiter( - (column for _, column in kept), dtype=np.intp, count=len(kept) - ) - kept_positions = ( - None - if len(kept) == len(keys) - else np.fromiter( - (position for position, _ in kept), - dtype=np.intp, - count=len(kept), - ) - ) - previous_keys = keys - values = np.fromiter(pred.values(), dtype=np.float32, count=len(pred)) - predictions_array[i, columns] = ( - values if kept_positions is None else values[kept_positions] - ) - collected_predictions[model_name] = torch.from_numpy(predictions_array) - - return collected_predictions, list(ensemble_classes) + + collected_predictions = {} + for model_name, (model_classes, scores) in predictions.items(): + if model_classes == ensemble_classes: + collected_predictions[model_name] = torch.from_numpy( + scores.astype(np.float32) + ) + continue + shared = [ + (source, cls_to_idx[cls]) + for source, cls in enumerate(model_classes) + if cls in cls_to_idx + ] + mapped = np.full((n_samples, len(ensemble_classes)), np.nan, dtype=np.float32) + if shared: + source_idx = np.fromiter( + (source for source, _ in shared), dtype=np.intp, count=len(shared) + ) + target_idx = np.fromiter( + (target for _, target in shared), dtype=np.intp, count=len(shared) + ) + mapped[:, target_idx] = scores[:, source_idx] + collected_predictions[model_name] = torch.from_numpy(mapped) + + return collected_predictions, ensemble_classes def predict( @@ -144,18 +126,16 @@ def predict( test_predictions = {} for model_name, model in base_learners.items(): if prediction_cache_dir is None: - test_predictions[model_name] = model.predict_list(molecules) + test_predictions[model_name] = model.predict_dense(molecules) else: cache_path = os.path.join( - prediction_cache_dir, f"{model_name}_test_predictions.pt" + prediction_cache_dir, f"{model_name}_test_predictions.npz" ) if os.path.exists(cache_path): - test_predictions[model_name] = torch.load( - cache_path, weights_only=False - ) + test_predictions[model_name] = load_dense_predictions(cache_path) else: - test_predictions[model_name] = model.predict_list(molecules) - torch.save(test_predictions[model_name], cache_path) + test_predictions[model_name] = model.predict_dense(molecules) + save_dense_predictions(cache_path, *test_predictions[model_name]) test_predictions, predicted_classes = collect_base_learner_predictions( test_predictions diff --git a/chebifier/prediction_models/base_predictor.py b/chebifier/prediction_models/base_predictor.py index 806cd6f..6caf8c7 100644 --- a/chebifier/prediction_models/base_predictor.py +++ b/chebifier/prediction_models/base_predictor.py @@ -1,9 +1,41 @@ from abc import ABC +import numpy as np from rdkit import Chem from .._custom_cache import modelwise_smiles_lru_cache +SCORE_DTYPE = np.float16 + + +def dicts_to_dense(predictions: list[dict | None]) -> tuple[list[str], np.ndarray]: + class_set = set() + previous_keys = None + for pred in predictions: + if pred: + keys = tuple(pred) + if keys != previous_keys: + class_set.update(keys) + previous_keys = keys + classes = sorted(class_set) + cls_to_idx = {cls: idx for idx, cls in enumerate(classes)} + + scores = np.full((len(predictions), len(classes)), np.nan, dtype=SCORE_DTYPE) + previous_keys = None + columns = None + for i, pred in enumerate(predictions): + if pred: + keys = tuple(pred) + if keys != previous_keys: + columns = np.fromiter( + (cls_to_idx[cls] for cls in keys), dtype=np.intp, count=len(keys) + ) + previous_keys = keys + scores[i, columns] = np.fromiter( + pred.values(), dtype=SCORE_DTYPE, count=len(pred) + ) + return classes, scores + class BasePredictor(ABC): def __init__( @@ -25,6 +57,11 @@ def predict(self, molecule: str | Chem.Mol) -> dict | None: # by default, use list-based prediction return self.predict_list([molecule])[0] + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + return dicts_to_dense(self.predict_list(molecule_list)) + @property def info_text(self): if self._description is None: diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index cf48658..7e4e1b4 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -1,11 +1,13 @@ from abc import ABC from typing import TYPE_CHECKING +import numpy as np from chebai.result.prediction import Predictor +from rdkit import Chem from chebifier import modelwise_smiles_lru_cache -from .base_predictor import BasePredictor +from .base_predictor import SCORE_DTYPE, BasePredictor if TYPE_CHECKING: from torch import Tensor @@ -51,6 +53,16 @@ def predict_list(self, smiles_list: list[str]) -> list: else: return [None for _ in smiles_list] + def predict_dense( + self, molecule_list: list[str | Chem.Mol] + ) -> tuple[list[str], np.ndarray]: + raw_preds: Tensor = self.predictor.predict_smiles(molecule_list) + if raw_preds is None: + return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) + classes = [str(label) for label in self.predictor._classification_labels] + scores = raw_preds.detach().cpu().numpy().astype(SCORE_DTYPE) + return classes, scores + def calculate_results(self, batch): collator = self.predictor._dm.reader.COLLATOR() dat = self.predictor._model._process_batch( From ea8177b1f07366b5695970ab43766ccdbe465906 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 17:18:19 +0200 Subject: [PATCH 07/10] add hyperparameter optimization for WMV-F1 --- chebifier/ensemble/voting_ensemble.py | 15 +- .../ensemble/weighted_majority_ensemble.py | 155 +++++++++++++++++- chebifier/prediction_models/nn_predictor.py | 26 +-- 3 files changed, 175 insertions(+), 21 deletions(-) diff --git a/chebifier/ensemble/voting_ensemble.py b/chebifier/ensemble/voting_ensemble.py index 6b23d5f..28e4e9c 100644 --- a/chebifier/ensemble/voting_ensemble.py +++ b/chebifier/ensemble/voting_ensemble.py @@ -16,6 +16,7 @@ def __init__( super().__init__(ensemble_dir) self.use_confidence = use_confidence self.classwise_f1 = None + self.prediction_thresholds = None def find_best_threshold(self, predictions, val_labels_tensor): best_threshold = 0.5 @@ -37,11 +38,15 @@ def calibrate(self, validation_predictions, validation_data, validation_labels): self.classwise_f1 = F1Score( task="multilabel", num_labels=validation_labels.shape[1], average=None ) - prediction_thresholds = { - model_name: self.find_best_threshold(predictions, validation_labels) - for model_name, predictions in validation_predictions.items() + self._save_prediction_thresholds( + self._fit_prediction_thresholds(validation_predictions, validation_labels) + ) + + def _fit_prediction_thresholds(self, predictions, labels) -> dict[str, float]: + return { + model_name: self.find_best_threshold(model_predictions, labels) + for model_name, model_predictions in predictions.items() } - self._save_prediction_thresholds(prediction_thresholds) def _save_prediction_thresholds(self, thresholds: dict[str, float]): thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" @@ -50,6 +55,8 @@ def _save_prediction_thresholds(self, thresholds: dict[str, float]): print(f"Saved prediction thresholds to {thresholds_path}: {thresholds}") def _load_prediction_thresholds(self) -> dict[str, float]: + if self.prediction_thresholds is not None: + return self.prediction_thresholds thresholds_path = Path(self.ensemble_dir) / "prediction_thresholds.yaml" if thresholds_path.exists(): with open(thresholds_path, "r", encoding="utf-8") as f: diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index ff51020..eca7b41 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -1,9 +1,13 @@ from pathlib import Path +import pandas as pd import torch from chebifier.ensemble.voting_ensemble import VotingEnsemble +N_FOLDS = 5 +WEIGHTING_STRENGTH_GRID = [0, 0.25, 0.5, 0.75, 1] + class WMVwithF1Ensemble(VotingEnsemble): @@ -22,17 +26,27 @@ def __init__( super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent + self.model_f1_scores = None def calibrate(self, validation_predictions, validation_data, validation_labels): super().calibrate(validation_predictions, validation_data, validation_labels) self._save_classwise_f1(validation_predictions, validation_labels) + self._optimize_hyperparameters(validation_predictions, validation_labels) + + def _fit_classwise_f1(self, predictions, labels, thresholds): + return { + model_name: self.classwise_f1( + model_predictions > thresholds[model_name], labels + ) + for model_name, model_predictions in predictions.items() + } def _save_classwise_f1(self, validation_predictions, validation_labels): thresholds = self._load_prediction_thresholds() - for model_name, predictions in validation_predictions.items(): - f1 = self.classwise_f1( - predictions > thresholds[model_name], validation_labels - ) + classwise_f1 = self._fit_classwise_f1( + validation_predictions, validation_labels, thresholds + ) + for model_name, f1 in classwise_f1.items(): f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" with open(f1_path, "w+") as f: f.writelines(f"{x}\n" for x in f1.tolist()) @@ -41,6 +55,8 @@ def _save_classwise_f1(self, validation_predictions, validation_labels): ) def _load_classwise_f1(self, model_name: str) -> torch.Tensor: + if self.model_f1_scores is not None: + return self.model_f1_scores[model_name] classwise_f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" if classwise_f1_path.exists(): with open(classwise_f1_path, "r", encoding="utf-8") as f: @@ -72,3 +88,134 @@ def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) ) ** self.weighting_exponent return trust_tensor + + def _build_folds(self, validation_predictions, validation_labels): + """Split the validation set into N_FOLDS folds and calibrate prediction thresholds and + class-wise F1 scores on the training part of each fold. Returns one + (thresholds, class-wise F1 scores, held-out indices) tuple per fold.""" + permutation = torch.randperm( + validation_labels.shape[0], generator=torch.Generator().manual_seed(0) + ) + fold_indices = [permutation[i::N_FOLDS] for i in range(N_FOLDS)] + folds = [] + for fold, test_idx in enumerate(fold_indices): + print(f"Calibrating fold {fold + 1}/{N_FOLDS}...") + train_idx = torch.cat( + [idx for other, idx in enumerate(fold_indices) if other != fold] + ) + train_predictions = { + model_name: model_predictions[train_idx] + for model_name, model_predictions in validation_predictions.items() + } + train_labels = validation_labels[train_idx] + thresholds = self._fit_prediction_thresholds( + train_predictions, train_labels + ) + classwise_f1 = self._fit_classwise_f1( + train_predictions, train_labels, thresholds + ) + folds.append((thresholds, classwise_f1, test_idx)) + return folds + + def _score_hyperparameters( + self, + folds, + validation_predictions, + validation_labels, + weighting_strength, + weighting_exponent, + ): + """Macro F1 of the aggregated predictions on each held-out fold.""" + self.weighting_strength = weighting_strength + self.weighting_exponent = weighting_exponent + scores = [] + for thresholds, classwise_f1, test_idx in folds: + self.prediction_thresholds = thresholds + self.model_f1_scores = classwise_f1 + aggregated = self.predict( + { + model_name: model_predictions[test_idx] + for model_name, model_predictions in validation_predictions.items() + } + ) + decisions = (aggregated["net_score"] > 0) & aggregated[ + "has_valid_predictions" + ] + scores.append( + self.classwise_f1(decisions, validation_labels[test_idx]).mean().item() + ) + self.prediction_thresholds = None + self.model_f1_scores = None + return scores + + def _optimize_hyperparameters(self, validation_predictions, validation_labels): + print( + f"Optimizing hyperparameters with {N_FOLDS}-fold cross-validation on the validation set..." + ) + weighting_strength = self.weighting_strength + weighting_exponent = self.weighting_exponent + folds = self._build_folds(validation_predictions, validation_labels) + results = [] + + def score(stage, strength, exponent): + scores = self._score_hyperparameters( + folds, validation_predictions, validation_labels, strength, exponent + ) + mean_score = sum(scores) / len(scores) + results.append( + { + "stage": stage, + "weighting_strength": strength, + "weighting_exponent": exponent, + "mean_macro_f1": mean_score, + "std_macro_f1": torch.tensor(scores).std().item(), + **{f"fold_{i}_macro_f1": s for i, s in enumerate(scores)}, + } + ) + print( + f"weighting_strength={strength}, weighting_exponent={exponent}: macro-f1 {mean_score:.4f}" + ) + return mean_score + + strength_scores = { + strength: score("weighting_strength", strength, 1) + for strength in WEIGHTING_STRENGTH_GRID + } + best_strength = max(strength_scores, key=strength_scores.get) + + best_exponent = 1 + best_score = strength_scores[best_strength] + exponent = 2 + while True: + exponent_score = score("weighting_exponent", best_strength, exponent) + if exponent_score <= best_score: + break + best_score = exponent_score + best_exponent = exponent + exponent += 1 + + self.weighting_strength = weighting_strength + self.weighting_exponent = weighting_exponent + self._save_hyperparameter_results( + results, best_strength, best_exponent, best_score + ) + + def _save_hyperparameter_results( + self, results, best_strength, best_exponent, best_score + ): + results_path = Path(self.ensemble_dir) / "hyperparameter_search.csv" + pd.DataFrame(results).to_csv(results_path, index=False) + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + pd.DataFrame( + [ + { + "weighting_strength": best_strength, + "weighting_exponent": best_exponent, + "mean_macro_f1": best_score, + } + ] + ).to_csv(best_path, index=False) + print( + f"Saved hyperparameter search results to {results_path}. Recommended parameters (saved to {best_path}): " + f"weighting_strength={best_strength}, weighting_exponent={best_exponent} (cross-validated macro-f1: {best_score:.4f})." + ) diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 7e4e1b4..53ab46b 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -10,7 +10,7 @@ from .base_predictor import SCORE_DTYPE, BasePredictor if TYPE_CHECKING: - from torch import Tensor + pass class NNPredictor(BasePredictor, ABC): @@ -36,18 +36,16 @@ def predict_list(self, smiles_list: list[str]) -> list: Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary of classes and predicted values. """ - raw_preds: Tensor = self.predictor.predict_smiles(smiles_list) + raw_preds = self.predictor.predict_smiles(smiles_list) if raw_preds is not None: preds = [ - ( - { - label: pred - for label, pred in zip( - self.predictor._classification_labels, raw_preds[i].tolist() - ) - } - ) - for i in range(len(smiles_list)) + { + label: pred + for label, pred in zip( + self.predictor._classification_labels, pred_tensor.tolist() + ) + } + for pred_tensor in raw_preds ] return preds else: @@ -56,11 +54,13 @@ def predict_list(self, smiles_list: list[str]) -> list: def predict_dense( self, molecule_list: list[str | Chem.Mol] ) -> tuple[list[str], np.ndarray]: - raw_preds: Tensor = self.predictor.predict_smiles(molecule_list) + raw_preds = self.predictor.predict_smiles(molecule_list) if raw_preds is None: return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) classes = [str(label) for label in self.predictor._classification_labels] - scores = raw_preds.detach().cpu().numpy().astype(SCORE_DTYPE) + scores = np.stack([pred.detach().cpu().numpy() for pred in raw_preds]).astype( + SCORE_DTYPE + ) return classes, scores def calculate_results(self, batch): From 844768a6ba2a18b75b4f1050c1daad0fab5f7239 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 30 Jul 2026 17:24:30 +0200 Subject: [PATCH 08/10] apply optimized hyperparameters if available --- .../ensemble/weighted_majority_ensemble.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index eca7b41..56d08da 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -15,13 +15,17 @@ def __init__( self, ensemble_dir: str, use_confidence: bool = True, - weighting_strength=1, - weighting_exponent=1, + weighting_strength=None, + weighting_exponent=None, **kwargs, ): """WMV ensemble that weights models based on their class-wise F1 scores. For each class, the weight is calculated as: weight = model_weight * (weighting_strength * F1 + (1 - weighting_strength)) ** weighting_exponent where F1 is the class-specific F1 score ("trust") of the model on the validation set. + + weighting_strength and weighting_exponent default to the optimal values determined during + calibration (best_hyperparameters.csv in the ensemble directory), falling back to 1 if the + ensemble has not been calibrated. Values passed here take precedence over both. """ super().__init__(ensemble_dir, use_confidence, **kwargs) self.weighting_strength = weighting_strength @@ -66,9 +70,26 @@ def _load_classwise_f1(self, model_name: str) -> torch.Tensor: f"Class-wise F1 scores file not found for model {model_name} in ensemble directory: {self.ensemble_dir}. Please calibrate the ensemble first." ) + def _load_hyperparameters(self) -> tuple[float, int]: + """Hyperparameters set explicitly take precedence, otherwise the optimal values found during + calibration are used (falling back to 1 if the ensemble has not been calibrated). + """ + best = {} + best_path = Path(self.ensemble_dir) / "best_hyperparameters.csv" + if best_path.exists(): + best = pd.read_csv(best_path).iloc[0].to_dict() + strength = self.weighting_strength + if strength is None: + strength = float(best.get("weighting_strength", 1)) + exponent = self.weighting_exponent + if exponent is None: + exponent = int(best.get("weighting_exponent", 1)) + return strength, exponent + def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: # Calculate trust based on class-wise F1 scores for each model # target shape: (num_molecules, num_classes, num_models) + weighting_strength, weighting_exponent = self._load_hyperparameters() num_models = len(predictions) num_molecules = list(predictions.values())[0].shape[0] num_classes = list(predictions.values())[0].shape[1] @@ -85,8 +106,8 @@ def calculate_trust(self, predictions: dict[str, torch.Tensor]) -> torch.Tensor: # Expand classwise_f1 to match the shape of trust_tensor for broadcasting classwise_f1 = classwise_f1.unsqueeze(0).expand(num_molecules, -1) trust_tensor[:, :, model_idx] = ( - self.weighting_strength * classwise_f1 + (1 - self.weighting_strength) - ) ** self.weighting_exponent + weighting_strength * classwise_f1 + (1 - weighting_strength) + ) ** weighting_exponent return trust_tensor def _build_folds(self, validation_predictions, validation_labels): From 4ea20f64ad0c13a821d979309fac9cfeb71e91f3 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 14:51:26 +0200 Subject: [PATCH 09/10] fix handling of missing predictions --- chebifier/prediction_models/nn_predictor.py | 35 ++++++++++++--------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/chebifier/prediction_models/nn_predictor.py b/chebifier/prediction_models/nn_predictor.py index 53ab46b..a39eff6 100644 --- a/chebifier/prediction_models/nn_predictor.py +++ b/chebifier/prediction_models/nn_predictor.py @@ -36,31 +36,36 @@ def predict_list(self, smiles_list: list[str]) -> list: Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary of classes and predicted values. """ - raw_preds = self.predictor.predict_smiles(smiles_list) - if raw_preds is not None: - preds = [ - { + raw_preds = self.predictor.predict_molecules(smiles_list) + if raw_preds is None: + return [None for _ in smiles_list] + return [ + ( + None + if pred_tensor is None + else { label: pred for label, pred in zip( self.predictor._classification_labels, pred_tensor.tolist() ) } - for pred_tensor in raw_preds - ] - return preds - else: - return [None for _ in smiles_list] + ) + for pred_tensor in raw_preds + ] def predict_dense( self, molecule_list: list[str | Chem.Mol] ) -> tuple[list[str], np.ndarray]: - raw_preds = self.predictor.predict_smiles(molecule_list) - if raw_preds is None: - return [], np.full((len(molecule_list), 0), np.nan, dtype=SCORE_DTYPE) + raw_preds = self.predictor.predict_molecules(molecule_list) classes = [str(label) for label in self.predictor._classification_labels] - scores = np.stack([pred.detach().cpu().numpy() for pred in raw_preds]).astype( - SCORE_DTYPE - ) + # molecules the model could not process stay NaN, so the ensemble skips this + # model for those rows only (see dicts_to_dense in base_predictor.py) + scores = np.full((len(molecule_list), len(classes)), np.nan, dtype=SCORE_DTYPE) + if raw_preds is None: + return classes, scores + for idx, pred in enumerate(raw_preds): + if pred is not None: + scores[idx] = pred.detach().cpu().numpy() return classes, scores def calculate_results(self, batch): From 32e14a154ae8306673b08ff81a129f414f2d3061 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 14:53:27 +0200 Subject: [PATCH 10/10] use chebi_utils SMILES / InChI parsing --- chebifier/cli.py | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/chebifier/cli.py b/chebifier/cli.py index 49d345b..a9f9901 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -5,7 +5,7 @@ import click import pandas as pd import yaml -from rdkit import Chem +from chebi_utils.read_molecule import smiles_or_inchi_to_mol from chebifier.build_ensemble import EnsembleBuilder from chebifier.check_env import check_package_installed @@ -22,24 +22,7 @@ def read_molecules(molecules, molecule_file): with open(molecule_file, "r", encoding="utf-8") as f: raw_inputs.extend([line.strip() for line in f if line.strip()]) - mol_list = [] - for raw_input in raw_inputs: - try: - if raw_input.startswith("InChI="): - mol = Chem.MolFromInchi(raw_input, sanitize=False) - if mol is None: - click.echo(f"Failed to parse InChI: {raw_input}") - mol_list.append(None) - mol_list.append(mol) - elif Chem.MolFromSmiles(raw_input, sanitize=False) is None: - click.echo(f"Failed to parse SMILES: {raw_input}") - mol_list.append(None) - else: - mol_list.append(Chem.MolFromSmiles(raw_input, sanitize=False)) - except Exception as e: - click.echo(f"Error parsing molecule '{raw_input}': {e}.") - mol_list.append(None) - return mol_list + return [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] def build_base_learners(ensemble_config):