From 010e30007fbf4c3122e00d12abc273a679645724 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 12:49:40 +0200 Subject: [PATCH 01/56] make dynamic dataset generic --- chebai/preprocessing/datasets/base.py | 18 ++++++++++++++---- chebai/preprocessing/datasets/chebi.py | 19 +++++++++++++++++-- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 24655b0d..3b53506f 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -909,10 +909,7 @@ def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None: print(f"Missing processed data file (`{processed_name}` file)") os.makedirs(self.processed_dir_main, exist_ok=True) data_path = self._download_required_data() - from chebi_utils import build_chebi_graph - - g = build_chebi_graph(data_path) - data_df = self._graph_to_raw_dataset(g) + data_df = self._preprocess_data_into_dataframe(data_path) self.save_processed(data_df, processed_name) @abstractmethod @@ -925,6 +922,19 @@ def _download_required_data(self) -> str: """ pass + @abstractmethod + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + """ + Preprocesses the raw data into a DataFrame. + + Args: + raw_data_path (str): Path to the raw data. + + Returns: + pd.DataFrame: The preprocessed data as a DataFrame. + """ + pass + @abstractmethod def _graph_to_raw_dataset(self, graph: "nx.DiGraph") -> pd.DataFrame: """ diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index 4d77495d..1a7e3d83 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -9,9 +9,9 @@ from itertools import cycle, permutations, product from typing import TYPE_CHECKING, Any, Generator, List, Literal, Optional -from networkx import DiGraph import numpy as np import pandas as pd +from networkx import DiGraph from rdkit import Chem from chebai.preprocessing import reader as dr @@ -150,6 +150,21 @@ def _download_required_data(self) -> str: self._load_sdf() return self._load_chebi() + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + """ + Preprocesses the raw data into a DataFrame. + + Args: + raw_data_path (str): Path to the raw data. + + Returns: + pd.DataFrame: The preprocessed data as a DataFrame. + """ + from chebi_utils import build_chebi_graph + + g = build_chebi_graph(raw_data_path) + return self._graph_to_raw_dataset(g) + def _load_chebi(self, version: Optional[int] = None) -> str: """ Load the ChEBI ontology file. @@ -746,12 +761,12 @@ def _graph_to_raw_dataset(self, g: "nx.DiGraph") -> pd.DataFrame: """ # Extract mol objects from SDF using chebi-utils + import networkx as nx from chebi_utils import ( build_labeled_dataset, extract_molecules, get_hierarchy_subgraph, ) - import networkx as nx sdf_path = os.path.join(self.raw_dir, self.raw_file_names_dict["sdf"]) mol_df = extract_molecules(sdf_path) From 8927a69a5173f84036de8cc789b6e124c0a492c0 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 13:50:56 +0200 Subject: [PATCH 02/56] add common class for molecule net --- chebai/preprocessing/datasets/base.py | 3 +- .../datasets/molecule_classification.py | 125 +++++++----------- 2 files changed, 53 insertions(+), 75 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 3b53506f..83161362 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -981,7 +981,7 @@ def setup_processed(self) -> None: Transforms `data.pkl` into a model input data format (`data.pt`), ensuring that the data is in a format compatible for input to the model. The transformed data contains the following keys: `ident`, `features`, `labels`, and `group`. - This method uses a subclass of Data Reader to perform the transformation. + This method uses assigned subclass of `DataReader` to perform the transformation. Returns: None @@ -1116,6 +1116,7 @@ def _retrieve_splits_from_csv(self) -> None: splits.csv to reconstruct the train, validation, and test splits. """ print(f"\nLoading splits from {self.splits_file_path}...") + assert self.splits_file_path is not None, "splits_file_path should not be None" splits_df = pd.read_csv(self.splits_file_path) filename = self.processed_file_names_dict["data"] diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index c2916675..03bec787 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -2,19 +2,54 @@ import gzip import os import shutil +from abc import ABC from tempfile import NamedTemporaryFile -from typing import Dict, List +from typing import Any, Dict, Generator, List from urllib import request import numpy as np +import pandas as pd import torch from sklearn.model_selection import GroupShuffleSplit, train_test_split from chebai.preprocessing import reader as dr -from chebai.preprocessing.datasets.base import XYBaseDataModule +from chebai.preprocessing.datasets.base import XYBaseDataModule, _DynamicDataset -class ClinTox(XYBaseDataModule): +class MoleculeNetDataExtractor(_DynamicDataset, ABC): + LABLES_COLUMNS = [] + FEATURE_COLUMN_NAME = "smiles" + ID_COLUMN_NAME = None + + @property + def _name(self) -> str: + """Returns the name of the dataset.""" + return str(self.__class__.__name__) + + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any]]: + """Loads data from a CSV file. + + Args: + input_file_path (str): Path to the CSV file. + + Returns: + List[Dict]: List of data dictionaries. + """ + with open(input_file_path, "rb") as input_file: + df = pd.read_pickle(input_file) + + features = df[self.FEATURE_COLUMN_NAME].to_numpy() + if self.ID_COLUMN_NAME is not None and self.ID_COLUMN_NAME in df.columns: + idents = df[self.ID_COLUMN_NAME].to_numpy() + else: + idents = np.arange(len(df)) + labels = df[self.LABLES_COLUMNS].to_numpy() + + for feat, labels, ident in zip(features, labels, idents): + yield dict(features=feat, labels=labels, ident=ident) + + +class ClinTox(MoleculeNetDataExtractor): """Data module for ClinTox MoleculeNet dataset.""" HEADERS = [ @@ -22,35 +57,12 @@ class ClinTox(XYBaseDataModule): "CT_TOX", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "ClinTox" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 2 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["clintox.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: + def _download_required_data(self) -> None: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: request.urlretrieve( @@ -61,6 +73,18 @@ def download(self) -> None: with open(os.path.join(self.raw_dir, "clintox.csv"), "wt") as fout: fout.write(gfile.read().decode()) + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + """ + Preprocesses the raw data into a DataFrame. + + Args: + raw_data_path (str): Path to the raw data. + + Returns: + pd.DataFrame: The preprocessed data as a DataFrame. + """ + return pd.read_csv(raw_data_path, header=0) + def setup_processed(self) -> None: """Processes and splits the dataset.""" print("Create splits") @@ -116,21 +140,6 @@ def setup_processed(self) -> None: os.path.join(self.processed_dir, f"{k}.pt"), ) - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - def _set_processed_data_props(self): """ Load processed data and extract metadata. @@ -147,38 +156,6 @@ def _set_processed_data_props(self): self._num_of_labels = len(data_pt[0]["labels"]) self._feature_vector_size = max(len(d["features"]) for d in data_pt) - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # group=group - ) - # yield dict(features=smiles, labels=labels, ident=i) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - class BBBP(XYBaseDataModule): """Data module for ClinTox MoleculeNet dataset.""" From 3209b56a2081047b0836039c75575b3833ef950d Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 13:57:02 +0200 Subject: [PATCH 03/56] remain code --- .../datasets/molecule_classification.py | 29 +++-- chebai/preprocessing/splitters/group.py | 111 ++++++++++++++++++ 2 files changed, 128 insertions(+), 12 deletions(-) create mode 100644 chebai/preprocessing/splitters/group.py diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 03bec787..d52e5995 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -140,21 +140,26 @@ def setup_processed(self) -> None: os.path.join(self.processed_dir, f"{k}.pt"), ) - def _set_processed_data_props(self): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. + Loads encoded/transformed data and generates training, validation, and test splits. """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + from chebi_utils import create_multilabel_splits + + splits = create_multilabel_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["val"], splits["test"] class BBBP(XYBaseDataModule): diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py new file mode 100644 index 00000000..1b6fa5f7 --- /dev/null +++ b/chebai/preprocessing/splitters/group.py @@ -0,0 +1,111 @@ +"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" + +from __future__ import annotations + +import pandas as pd + + +def create_group_splits( + df: pd.DataFrame, + label_start_col: int = 2, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, + seed: int | None = 42, +) -> dict[str, pd.DataFrame]: + """Create stratified train/validation/test splits for multilabel DataFrames. + + Columns from index *label_start_col* onwards are treated as binary label + columns (one boolean column per label). The stratification strategy is + chosen automatically based on the number of label columns: + + - More than one label column: ``MultilabelStratifiedShuffleSplit`` from + the ``iterative-stratification`` package. + - Single label column: ``StratifiedShuffleSplit`` from ``scikit-learn``. + + Parameters + ---------- + df : pd.DataFrame + Input data. Columns ``0`` to ``label_start_col - 1`` are treated as + feature/metadata columns; all remaining columns are boolean label + columns. A typical ChEBI DataFrame has columns + ``["chebi_id", "mol", "label1", "label2", ...]``. + label_start_col : int + Index of the first label column (default 2). + train_ratio : float + Fraction of data for training (default 0.8). + val_ratio : float + Fraction of data for validation (default 0.1). + test_ratio : float + Fraction of data for testing (default 0.1). + seed : int or None + Random seed for reproducibility. + + Returns + ------- + dict + Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each + containing a DataFrame. + + Raises + ------ + ValueError + If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, or + *label_start_col* is out of range. + """ + if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6: + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]): + raise ValueError("All ratios must be between 0 and 1") + if label_start_col >= len(df.columns): + raise ValueError( + f"label_start_col={label_start_col} is out of range for a DataFrame " + f"with {len(df.columns)} columns" + ) + + from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit + from sklearn.model_selection import StratifiedShuffleSplit + + labels_matrix = df.iloc[:, label_start_col:].values + is_multilabel = labels_matrix.shape[1] > 1 + # StratifiedShuffleSplit requires a 1-D label array + y = labels_matrix if is_multilabel else labels_matrix[:, 0] + + df_reset = df.reset_index(drop=True) + + # ── Step 1: carve out the test set ────────────────────────────────────── + if is_multilabel: + test_splitter = MultilabelStratifiedShuffleSplit( + n_splits=1, test_size=test_ratio, random_state=seed + ) + else: + test_splitter = StratifiedShuffleSplit( + n_splits=1, test_size=test_ratio, random_state=seed + ) + train_val_idx, test_idx = next(test_splitter.split(y, y)) + + df_test = df_reset.iloc[test_idx] + df_trainval = df_reset.iloc[train_val_idx] + + # ── Step 2: split train/val from the remaining data ───────────────────── + y_trainval = y[train_val_idx] + val_ratio_adjusted = val_ratio / (1.0 - test_ratio) + + if is_multilabel: + val_splitter = MultilabelStratifiedShuffleSplit( + n_splits=1, test_size=val_ratio_adjusted, random_state=seed + ) + else: + val_splitter = StratifiedShuffleSplit( + n_splits=1, test_size=val_ratio_adjusted, random_state=seed + ) + train_idx_inner, val_idx_inner = next(val_splitter.split(y_trainval, y_trainval)) + + df_train = df_trainval.iloc[train_idx_inner] + df_val = df_trainval.iloc[val_idx_inner] + + return { + "train": df_train.reset_index(drop=True), + "val": df_val.reset_index(drop=True), + "test": df_test.reset_index(drop=True), + } From 6c759e6ec1108c5ad911e9f8e972486230c92cd4 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 14:38:10 +0200 Subject: [PATCH 04/56] add group splitter class --- .../datasets/molecule_classification.py | 803 +----------------- chebai/preprocessing/splitters/__init__.py | 3 + chebai/preprocessing/splitters/group.py | 71 +- 3 files changed, 81 insertions(+), 796 deletions(-) create mode 100644 chebai/preprocessing/splitters/__init__.py diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index d52e5995..4b21441d 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -1,22 +1,24 @@ -import csv import gzip import os import shutil from abc import ABC from tempfile import NamedTemporaryFile -from typing import Any, Dict, Generator, List +from typing import Any, Generator, List from urllib import request import numpy as np import pandas as pd import torch -from sklearn.model_selection import GroupShuffleSplit, train_test_split +from sklearn.model_selection import train_test_split from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import XYBaseDataModule, _DynamicDataset +from chebai.preprocessing.splitters import GroupSplitter class MoleculeNetDataExtractor(_DynamicDataset, ABC): + READER = dr.ChemDataReader + LABLES_COLUMNS = [] FEATURE_COLUMN_NAME = "smiles" ID_COLUMN_NAME = None @@ -26,6 +28,18 @@ def _name(self) -> str: """Returns the name of the dataset.""" return str(self.__class__.__name__) + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + """ + Preprocesses the raw data into a DataFrame. + + Args: + raw_data_path (str): Path to the raw data. + + Returns: + pd.DataFrame: The preprocessed data as a DataFrame. + """ + return pd.read_csv(raw_data_path, header=0) + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any]]: """Loads data from a CSV file. @@ -49,10 +63,10 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any]]: yield dict(features=feat, labels=labels, ident=ident) -class ClinTox(MoleculeNetDataExtractor): +class ClinTox(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABLES_COLUMNS = [ "FDA_APPROVED", "CT_TOX", ] @@ -73,131 +87,20 @@ def _download_required_data(self) -> None: with open(os.path.join(self.raw_dir, "clintox.csv"), "wt") as fout: fout.write(gfile.read().decode()) - def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: - """ - Preprocesses the raw data into a DataFrame. - - Args: - raw_data_path (str): Path to the raw data. - - Returns: - pd.DataFrame: The preprocessed data as a DataFrame. - """ - return pd.read_csv(raw_data_path, header=0) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list( - self._load_data_from_file(os.path.join(self.raw_dir, "clintox.csv")) - ) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d for d in (data[temp_split_index[i]] for i in test_split_index) - ] - validation_split = [ - d for d in (data[temp_split_index[i]] for i in validation_split_index) - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: - """ - Loads encoded/transformed data and generates training, validation, and test splits. - """ - - filename = self.processed_file_names_dict["data"] - data = self.load_processed_data_from_file(filename) - df_data = pd.DataFrame(data) - - from chebi_utils import create_multilabel_splits - splits = create_multilabel_splits( - df_data, - self._LABELS_START_IDX, - 1 - self.validation_split - self.test_split, - self.validation_split, - self.test_split, - self.dynamic_data_split_seed, - ) - return splits["train"], splits["val"], splits["test"] - - -class BBBP(XYBaseDataModule): +class BBBP(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABLES_COLUMNS = [ "p_np", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "BBBP" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["bbbp.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: + def _download_required_data(self) -> None: """Downloads and extracts the dataset.""" with open(os.path.join(self.raw_dir, "bbbp.csv"), "ab") as dst: with request.urlopen( @@ -205,128 +108,11 @@ def download(self) -> None: ) as src: shutil.copyfileobj(src, dst) - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "bbbp.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - print("Group shuffled") - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [int(row["p_np"])] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - class Sider(XYBaseDataModule): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABLES_COLUMNS = [ "Hepatobiliary disorders", "Metabolism and nutrition disorders", "Product issues", @@ -356,35 +142,12 @@ class Sider(XYBaseDataModule): "Injury, poisoning and procedural complications", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Sider" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 27 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["sider.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: + def _download_required_data(self) -> None: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: request.urlretrieve( @@ -395,161 +158,19 @@ def download(self) -> None: with open(os.path.join(self.raw_dir, "sider.csv"), "wt") as fout: fout.write(gfile.read().decode()) - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "sider.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = row["group"] - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - class Bace(XYBaseDataModule): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABELS_COLUMNS = [ "class", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Bace" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["bace.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - def download(self) -> None: """Downloads and extracts the dataset.""" with open(os.path.join(self.raw_dir, "bace.csv"), "ab") as dst: @@ -608,97 +229,20 @@ def setup_processed(self) -> None: os.path.join(self.processed_dir, f"{k}.pt"), ) - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["mol"] - labels = [int(row["Class"])] - # group = row["group"] - yield dict(features=smiles, labels=labels, ident=i) # , group=group - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class HIV(XYBaseDataModule): +class HIV(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABELS_COLUMNS = [ "HIV_active", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "HIV" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["hiv.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: + def _download_required_data(self) -> None: """Downloads and extracts the dataset.""" with open(os.path.join(self.raw_dir, "hiv.csv"), "ab") as dst: with request.urlopen( @@ -706,125 +250,11 @@ def download(self) -> None: ) as src: shutil.copyfileobj(src, dst) - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "hiv.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - print("Group shuffled") - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d for d in (data[temp_split_index[i]] for i in test_split_index) - ] - validation_split = [ - d for d in (data[temp_split_index[i]] for i in validation_split_index) - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - if len(row) > 1: - i += 1 - smiles = row["smiles"] - labels = [int(row["HIV_active"])] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class MUV(XYBaseDataModule): +class MUV(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" - HEADERS = [ + LABELS_COLUMNS = [ "MUV-466", "MUV-548", "MUV-600", @@ -844,34 +274,11 @@ class MUV(XYBaseDataModule): "MUV-859", ] - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "MUV" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 17 - @property def raw_file_names(self) -> List[str]: """Returns a list of raw file names.""" return ["muv.csv"] - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - def download(self) -> None: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: @@ -882,153 +289,3 @@ def download(self) -> None: with gzip.open(gout.name) as gfile: with open(os.path.join(self.raw_dir, "muv.csv"), "wt") as fout: fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "muv.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = row["group"] - yield dict(features=smiles, labels=labels, ident=i) # , group=group) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class BaceChem(Bace): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class SiderChem(Sider): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class BBBPChem(BBBP): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class ClinToxChem(ClinTox): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class HIVChem(HIV): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class MUVChem(MUV): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader diff --git a/chebai/preprocessing/splitters/__init__.py b/chebai/preprocessing/splitters/__init__.py new file mode 100644 index 00000000..961c4115 --- /dev/null +++ b/chebai/preprocessing/splitters/__init__.py @@ -0,0 +1,3 @@ +from .group import GroupSplitter + +__all__ = ["GroupSplitter"] diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py index 1b6fa5f7..7f1338ec 100644 --- a/chebai/preprocessing/splitters/group.py +++ b/chebai/preprocessing/splitters/group.py @@ -2,7 +2,33 @@ from __future__ import annotations +from abc import ABC + import pandas as pd +from sklearn.model_selection import GroupShuffleSplit + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class GroupSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + splits = create_group_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["val"], splits["test"] def create_group_splits( @@ -63,26 +89,27 @@ def create_group_splits( f"with {len(df.columns)} columns" ) - from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit - from sklearn.model_selection import StratifiedShuffleSplit + if "group" not in df.columns: + raise ValueError( + "Input DataFrame must contain a 'group' column for group split" + ) + + if len(df["group"].unique()) < 2: + raise ValueError( + "Input DataFrame must contain at least 2 unique groups for group split" + ) - labels_matrix = df.iloc[:, label_start_col:].values - is_multilabel = labels_matrix.shape[1] > 1 + y = df.iloc[:, label_start_col:].values # StratifiedShuffleSplit requires a 1-D label array - y = labels_matrix if is_multilabel else labels_matrix[:, 0] df_reset = df.reset_index(drop=True) # ── Step 1: carve out the test set ────────────────────────────────────── - if is_multilabel: - test_splitter = MultilabelStratifiedShuffleSplit( - n_splits=1, test_size=test_ratio, random_state=seed - ) - else: - test_splitter = StratifiedShuffleSplit( - n_splits=1, test_size=test_ratio, random_state=seed - ) - train_val_idx, test_idx = next(test_splitter.split(y, y)) + test_splitter = GroupShuffleSplit( + n_splits=1, test_size=test_ratio, random_state=seed + ) + + train_val_idx, test_idx = next(test_splitter.split(y, y, groups=df_reset["group"])) df_test = df_reset.iloc[test_idx] df_trainval = df_reset.iloc[train_val_idx] @@ -91,15 +118,13 @@ def create_group_splits( y_trainval = y[train_val_idx] val_ratio_adjusted = val_ratio / (1.0 - test_ratio) - if is_multilabel: - val_splitter = MultilabelStratifiedShuffleSplit( - n_splits=1, test_size=val_ratio_adjusted, random_state=seed - ) - else: - val_splitter = StratifiedShuffleSplit( - n_splits=1, test_size=val_ratio_adjusted, random_state=seed - ) - train_idx_inner, val_idx_inner = next(val_splitter.split(y_trainval, y_trainval)) + val_splitter = GroupShuffleSplit( + n_splits=1, test_size=val_ratio_adjusted, random_state=seed + ) + + train_idx_inner, val_idx_inner = next( + val_splitter.split(y_trainval, y_trainval, groups=df_trainval["group"]) + ) df_train = df_trainval.iloc[train_idx_inner] df_val = df_trainval.iloc[val_idx_inner] From e492b241d0bf5322bf945815af43810392079ad6 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 14:44:10 +0200 Subject: [PATCH 05/56] add class for multilabel splitter --- chebai/preprocessing/datasets/chebi.py | 24 ++------------- chebai/preprocessing/splitters/__init__.py | 3 +- chebai/preprocessing/splitters/multilabel.py | 32 ++++++++++++++++++++ 3 files changed, 36 insertions(+), 23 deletions(-) create mode 100644 chebai/preprocessing/splitters/multilabel.py diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index 1a7e3d83..9f209bc2 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -16,12 +16,13 @@ from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import _DynamicDataset +from chebai.preprocessing.splitters import MultiLabelSplitter if TYPE_CHECKING: import networkx as nx -class _ChEBIDataExtractor(_DynamicDataset, ABC): +class _ChEBIDataExtractor(MultiLabelSplitter, _DynamicDataset, ABC): """ A class for extracting and processing data from the ChEBI dataset. @@ -389,27 +390,6 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No for feat, labels, ident in zip(features, all_labels, idents): yield dict(features=feat, labels=labels, ident=ident) - def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: - """ - Loads encoded/transformed data and generates training, validation, and test splits. - """ - - filename = self.processed_file_names_dict["data"] - data = self.load_processed_data_from_file(filename) - df_data = pd.DataFrame(data) - - from chebi_utils import create_multilabel_splits - - splits = create_multilabel_splits( - df_data, - self._LABELS_START_IDX, - 1 - self.validation_split - self.test_split, - self.validation_split, - self.test_split, - self.dynamic_data_split_seed, - ) - return splits["train"], splits["val"], splits["test"] - def _setup_pruned_test_set( self, df_test_chebi_version: pd.DataFrame ) -> pd.DataFrame: diff --git a/chebai/preprocessing/splitters/__init__.py b/chebai/preprocessing/splitters/__init__.py index 961c4115..f46b77fe 100644 --- a/chebai/preprocessing/splitters/__init__.py +++ b/chebai/preprocessing/splitters/__init__.py @@ -1,3 +1,4 @@ from .group import GroupSplitter +from .multilabel import MultiLabelSplitter -__all__ = ["GroupSplitter"] +__all__ = ["GroupSplitter", "MultiLabelSplitter"] diff --git a/chebai/preprocessing/splitters/multilabel.py b/chebai/preprocessing/splitters/multilabel.py new file mode 100644 index 00000000..77366ea8 --- /dev/null +++ b/chebai/preprocessing/splitters/multilabel.py @@ -0,0 +1,32 @@ +"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" + +from __future__ import annotations + +from abc import ABC + +import pandas as pd + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class MultiLabelSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + from chebi_utils import create_multilabel_splits + + splits = create_multilabel_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["val"], splits["test"] From 3f33b546c26e975acfbc69e4d66d138575409c04 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 14:47:04 +0200 Subject: [PATCH 06/56] update classification molecule net config --- configs/data/moleculenet/bace_moleculenet.yml | 2 +- configs/data/moleculenet/bbbp_moleculenet.yml | 2 +- configs/data/moleculenet/clintox_moleculenet.yml | 2 +- configs/data/moleculenet/hiv_moleculenet.yml | 2 +- configs/data/moleculenet/muv_moleculenet.yml | 2 +- configs/data/moleculenet/sider_moleculenet.yml | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index bd6c04a8..f801cc9f 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.BaceChem +class_path: chebai.preprocessing.datasets.molecule_classification.Bace init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index 01479443..2ef99142 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.BBBPChem +class_path: chebai.preprocessing.datasets.molecule_classification.BBBP init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index d7b7c3be..870a6445 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.ClinToxChem +class_path: chebai.preprocessing.datasets.molecule_classification.ClinTox init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index 3bef06b2..f3499e26 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.HIVChem +class_path: chebai.preprocessing.datasets.molecule_classification.HIV init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index d7498305..d276a76e 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.MUVChem +class_path: chebai.preprocessing.datasets.molecule_classification.MUV init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 1a1d81ee..3b774843 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.SiderChem +class_path: chebai.preprocessing.datasets.molecule_classification.Sider init_args: batch_size: 10 validation_split: 0.05 From 75d6301ce2a95112a923a1c29489e5b5f12fa14d Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 17:21:04 +0200 Subject: [PATCH 07/56] general splitter class --- .../datasets/molecule_classification.py | 60 +--------- chebai/preprocessing/splitters/__init__.py | 3 +- chebai/preprocessing/splitters/general.py | 111 ++++++++++++++++++ 3 files changed, 117 insertions(+), 57 deletions(-) create mode 100644 chebai/preprocessing/splitters/general.py diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 4b21441d..285e551f 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -8,12 +8,10 @@ import numpy as np import pandas as pd -import torch -from sklearn.model_selection import train_test_split from chebai.preprocessing import reader as dr -from chebai.preprocessing.datasets.base import XYBaseDataModule, _DynamicDataset -from chebai.preprocessing.splitters import GroupSplitter +from chebai.preprocessing.datasets.base import _DynamicDataset +from chebai.preprocessing.splitters import GeneralSplitter, GroupSplitter class MoleculeNetDataExtractor(_DynamicDataset, ABC): @@ -109,7 +107,7 @@ def _download_required_data(self) -> None: shutil.copyfileobj(src, dst) -class Sider(XYBaseDataModule): +class Sider(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" LABLES_COLUMNS = [ @@ -159,7 +157,7 @@ def _download_required_data(self) -> None: fout.write(gfile.read().decode()) -class Bace(XYBaseDataModule): +class Bace(MoleculeNetDataExtractor, GeneralSplitter): """Data module for ClinTox MoleculeNet dataset.""" LABELS_COLUMNS = [ @@ -179,56 +177,6 @@ def download(self) -> None: ) as src: shutil.copyfileobj(src, dst) - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "bace.csv"))) - # groups = np.array([d.get("group") for d in data]) - - # if not all(g is None for g in groups): - # split_size = int(len(set(groups)) * (1 - self.test_split - self.validation_split)) - # os.makedirs(self.processed_dir, exist_ok=True) - # splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - # train_split_index, temp_split_index = next( - # splitter.split(data, groups=groups) - # ) - - # split_groups = groups[temp_split_index] - - # splitter = GroupShuffleSplit( - # train_size=int(len(set(split_groups)) * (1 - self.test_split - self.validation_split)), n_splits=1 - # ) - # test_split_index, validation_split_index = next( - # splitter.split(temp_split_index, groups=split_groups) - # ) - # train_split = [data[i] for i in train_split_index] - # test_split = [ - # d - # for d in (data[temp_split_index[i]] for i in test_split_index) - # ] - # validation_split = [ - # d - # for d in (data[temp_split_index[i]] for i in validation_split_index) - # ] - # else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - class HIV(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" diff --git a/chebai/preprocessing/splitters/__init__.py b/chebai/preprocessing/splitters/__init__.py index f46b77fe..84c176bf 100644 --- a/chebai/preprocessing/splitters/__init__.py +++ b/chebai/preprocessing/splitters/__init__.py @@ -1,4 +1,5 @@ +from .general import GeneralSplitter from .group import GroupSplitter from .multilabel import MultiLabelSplitter -__all__ = ["GroupSplitter", "MultiLabelSplitter"] +__all__ = ["GroupSplitter", "MultiLabelSplitter", "GeneralSplitter"] diff --git a/chebai/preprocessing/splitters/general.py b/chebai/preprocessing/splitters/general.py new file mode 100644 index 00000000..64230130 --- /dev/null +++ b/chebai/preprocessing/splitters/general.py @@ -0,0 +1,111 @@ +"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" + +from __future__ import annotations + +from abc import ABC + +import pandas as pd +from sklearn.model_selection import train_test_split + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class GeneralSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + splits = create_general_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["val"], splits["test"] + + +def create_general_splits( + df: pd.DataFrame, + label_start_col: int = 2, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, +) -> dict[str, pd.DataFrame]: + """Create stratified train/validation/test splits for multilabel DataFrames. + + Columns from index *label_start_col* onwards are treated as binary label + columns (one boolean column per label). The stratification strategy is + chosen automatically based on the number of label columns: + + - More than one label column: ``MultilabelStratifiedShuffleSplit`` from + the ``iterative-stratification`` package. + - Single label column: ``StratifiedShuffleSplit`` from ``scikit-learn``. + + Parameters + ---------- + df : pd.DataFrame + Input data. Columns ``0`` to ``label_start_col - 1`` are treated as + feature/metadata columns; all remaining columns are boolean label + columns. A typical ChEBI DataFrame has columns + ``["chebi_id", "mol", "label1", "label2", ...]``. + label_start_col : int + Index of the first label column (default 2). + train_ratio : float + Fraction of data for training (default 0.8). + val_ratio : float + Fraction of data for validation (default 0.1). + test_ratio : float + Fraction of data for testing (default 0.1). + seed : int or None + Random seed for reproducibility. + + Returns + ------- + dict + Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each + containing a DataFrame. + + Raises + ------ + ValueError + If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, or + *label_start_col* is out of range. + """ + if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6: + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]): + raise ValueError("All ratios must be between 0 and 1") + if label_start_col >= len(df.columns): + raise ValueError( + f"label_start_col={label_start_col} is out of range for a DataFrame " + f"with {len(df.columns)} columns" + ) + + df_reset = df.reset_index(drop=True) + + # ── Step 1: carve out the test set ────────────────────────────────────── + df_trainval, df_test = train_test_split( + df_reset, test_size=test_ratio, shuffle=True + ) + + # ── Step 2: split train/val from the remaining data ───────────────────── + val_ratio_adjusted = val_ratio / (1.0 - test_ratio) + + df_train, df_val = train_test_split( + df_trainval, + test_size=val_ratio_adjusted, + shuffle=True, + ) + + return { + "train": df_train.reset_index(drop=True), + "val": df_val.reset_index(drop=True), + "test": df_test.reset_index(drop=True), + } From e87df7694725f548d94af48d691a7a9494567d91 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 17:21:23 +0200 Subject: [PATCH 08/56] certific dependency for ssl error --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index beda3a3b..7e59f0c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,9 @@ dev = [ "deepsmiles", "torchmetrics", "chebi-utils>=0.1.1", + # In case of urllib.error.URLError: Date: Fri, 17 Jul 2026 19:58:09 +0200 Subject: [PATCH 09/56] fix file path error --- chebai/preprocessing/datasets/base.py | 18 +--- .../datasets/molecule_classification.py | 99 +++++++++++++------ chebai/preprocessing/datasets/pubchem.py | 5 - 3 files changed, 70 insertions(+), 52 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 83161362..33caec15 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -617,6 +617,7 @@ def raw_file_names(self) -> List[str]: return list(self.raw_file_names_dict.values()) @property + @abstractmethod def raw_file_names_dict(self) -> dict: """ Returns the dictionary of raw file names (i.e., files that are directly obtained from an external source). @@ -815,7 +816,7 @@ class _DynamicDataset(XYBaseDataModule, ABC): apply_id_filter (Optional[str]): Path to a data.pt file for ID filtering. """ - # ---- Index for columns of processed `data.pkl` (should be derived from `_graph_to_raw_dataset` method) ------ + # ---- Index for columns of processed `data.pkl` (should be derived from `_preprocess_data_into_dataframe` method) ------ _ID_IDX: int = None _DATA_REPRESENTATION_IDX: int = None _LABELS_START_IDX: int = None @@ -935,21 +936,6 @@ def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: """ pass - @abstractmethod - def _graph_to_raw_dataset(self, graph: "nx.DiGraph") -> pd.DataFrame: - """ - Converts the graph to a raw dataset. - Uses the graph created by chebi_utils to extract the - raw data in Dataframe format with additional columns corresponding to each multi-label class. - - Args: - graph (nx.DiGraph): The class hierarchy graph. - - Returns: - pd.DataFrame: The raw dataset. - """ - pass - def save_processed(self, data: pd.DataFrame, filename: str) -> None: """ Save the processed dataset to a pickle file. diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 285e551f..829f7ef3 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -38,7 +38,7 @@ def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: """ return pd.read_csv(raw_data_path, header=0) - def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any]]: + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: """Loads data from a CSV file. Args: @@ -60,6 +60,16 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any]]: for feat, labels, ident in zip(features, labels, idents): yield dict(features=feat, labels=labels, ident=ident) + @property + def base_dir(self) -> str: + """ + Return the base directory path for data. + + Returns: + str: The base directory path for data. + """ + return os.path.join("data", "MoleculeNetClassification") + class ClinTox(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" @@ -70,11 +80,11 @@ class ClinTox(MoleculeNetDataExtractor, GroupSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["clintox.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"clintox": "clintox.csv"} - def _download_required_data(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: request.urlretrieve( @@ -82,8 +92,12 @@ def _download_required_data(self) -> None: gout.name, ) with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "clintox.csv"), "wt") as fout: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["clintox"]), + "wt", + ) as fout: fout.write(gfile.read().decode()) + return os.path.join(self.raw_dir, self.raw_file_names_dict["clintox"]) class BBBP(MoleculeNetDataExtractor, GroupSplitter): @@ -94,17 +108,20 @@ class BBBP(MoleculeNetDataExtractor, GroupSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["bbbp.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"bbbp": "bbbp.csv"} - def _download_required_data(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "bbbp.csv"), "ab") as dst: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["bbbp"]), "ab" + ) as dst: with request.urlopen( "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv", ) as src: shutil.copyfileobj(src, dst) + return os.path.join(self.raw_dir, self.raw_file_names_dict["bbbp"]) class Sider(MoleculeNetDataExtractor, GroupSplitter): @@ -141,11 +158,11 @@ class Sider(MoleculeNetDataExtractor, GroupSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["sider.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"sider": "sider.csv"} - def _download_required_data(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: request.urlretrieve( @@ -153,8 +170,11 @@ def _download_required_data(self) -> None: gout.name, ) with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "sider.csv"), "wt") as fout: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["sider"]), "wt" + ) as fout: fout.write(gfile.read().decode()) + return os.path.join(self.raw_dir, self.raw_file_names_dict["sider"]) class Bace(MoleculeNetDataExtractor, GeneralSplitter): @@ -165,17 +185,20 @@ class Bace(MoleculeNetDataExtractor, GeneralSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["bace.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"bace": "bace.csv"} - def download(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "bace.csv"), "ab") as dst: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["bace"]), "ab" + ) as dst: with request.urlopen( "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/bace.csv", ) as src: shutil.copyfileobj(src, dst) + return os.path.join(self.raw_dir, self.raw_file_names_dict["bace"]) class HIV(MoleculeNetDataExtractor, GroupSplitter): @@ -186,17 +209,20 @@ class HIV(MoleculeNetDataExtractor, GroupSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["hiv.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"hiv": "hiv.csv"} - def _download_required_data(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "hiv.csv"), "ab") as dst: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["hiv"]), "ab" + ) as dst: with request.urlopen( "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/HIV.csv", ) as src: shutil.copyfileobj(src, dst) + return os.path.join(self.raw_dir, self.raw_file_names_dict["hiv"]) class MUV(MoleculeNetDataExtractor, GroupSplitter): @@ -223,11 +249,11 @@ class MUV(MoleculeNetDataExtractor, GroupSplitter): ] @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["muv.csv"] + def raw_file_names_dict(self) -> dict: + """Returns a dictionary of raw file names.""" + return {"muv": "muv.csv"} - def download(self) -> None: + def _download_required_data(self) -> str: """Downloads and extracts the dataset.""" with NamedTemporaryFile("rb") as gout: request.urlretrieve( @@ -235,5 +261,16 @@ def download(self) -> None: gout.name, ) with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "muv.csv"), "wt") as fout: + with open( + os.path.join(self.raw_dir, self.raw_file_names_dict["muv"]), "wt" + ) as fout: fout.write(gfile.read().decode()) + + return os.path.join(self.raw_dir, self.raw_file_names_dict["muv"]) + + +if __name__ == "__main__": + # Example usage + dataset = ClinTox() + dataset.prepare_data() + dataset.setup() diff --git a/chebai/preprocessing/datasets/pubchem.py b/chebai/preprocessing/datasets/pubchem.py index ea5e8978..5ac33439 100644 --- a/chebai/preprocessing/datasets/pubchem.py +++ b/chebai/preprocessing/datasets/pubchem.py @@ -138,11 +138,6 @@ def _download_required_data(self) -> str: self.download() return self._raw_data_source_path - def _graph_to_raw_dataset(self, graph): - raise NotImplementedError( - "PubChem does not use a graph-based data preparation pipeline." - ) - def download(self): """ Downloads PubChem data based on `_k` parameter. From c9841ad74e822d62cb679aded0222fe5abab8867 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 21:12:22 +0200 Subject: [PATCH 10/56] right column names for each data --- .../datasets/molecule_classification.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 829f7ef3..424fa846 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -20,6 +20,7 @@ class MoleculeNetDataExtractor(_DynamicDataset, ABC): LABLES_COLUMNS = [] FEATURE_COLUMN_NAME = "smiles" ID_COLUMN_NAME = None + GROUP_COLUMN_NAME = "group" @property def _name(self) -> str: @@ -57,8 +58,13 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No idents = np.arange(len(df)) labels = df[self.LABLES_COLUMNS].to_numpy() - for feat, labels, ident in zip(features, labels, idents): - yield dict(features=feat, labels=labels, ident=ident) + if self.GROUP_COLUMN_NAME in df.columns: + groups = df[self.GROUP_COLUMN_NAME].to_numpy() + for feat, labels, ident, group in zip(features, labels, idents, groups): + yield dict(features=feat, labels=labels, ident=ident, group=group) + else: + for feat, labels, ident in zip(features, labels, idents): + yield dict(features=feat, labels=labels, ident=ident) @property def base_dir(self) -> str: @@ -68,7 +74,7 @@ def base_dir(self) -> str: Returns: str: The base directory path for data. """ - return os.path.join("data", "MoleculeNetClassification") + return os.path.join("data", f"{self._name}:MNClassification") class ClinTox(MoleculeNetDataExtractor, GroupSplitter): @@ -101,8 +107,9 @@ def _download_required_data(self) -> str: class BBBP(MoleculeNetDataExtractor, GroupSplitter): - """Data module for ClinTox MoleculeNet dataset.""" + """Data module for BBBP MoleculeNet dataset.""" + ID_COLUMN_NAME = "num" LABLES_COLUMNS = [ "p_np", ] @@ -125,7 +132,7 @@ def _download_required_data(self) -> str: class Sider(MoleculeNetDataExtractor, GroupSplitter): - """Data module for ClinTox MoleculeNet dataset.""" + """Data module for Sider MoleculeNet dataset.""" LABLES_COLUMNS = [ "Hepatobiliary disorders", @@ -178,10 +185,12 @@ def _download_required_data(self) -> str: class Bace(MoleculeNetDataExtractor, GeneralSplitter): - """Data module for ClinTox MoleculeNet dataset.""" + """Data module for Bace MoleculeNet dataset.""" + ID_COLUMN_NAME = "CID" + FEATURE_COLUMN_NAME = "mol" LABELS_COLUMNS = [ - "class", + "Class", ] @property @@ -202,7 +211,7 @@ def _download_required_data(self) -> str: class HIV(MoleculeNetDataExtractor, GroupSplitter): - """Data module for ClinTox MoleculeNet dataset.""" + """Data module for HIV MoleculeNet dataset.""" LABELS_COLUMNS = [ "HIV_active", @@ -226,8 +235,9 @@ def _download_required_data(self) -> str: class MUV(MoleculeNetDataExtractor, GroupSplitter): - """Data module for ClinTox MoleculeNet dataset.""" + """Data module for MUV MoleculeNet dataset.""" + ID_COLUMN_NAME = "mol_id" LABELS_COLUMNS = [ "MUV-466", "MUV-548", @@ -271,6 +281,6 @@ def _download_required_data(self) -> str: if __name__ == "__main__": # Example usage - dataset = ClinTox() + dataset = Sider() dataset.prepare_data() - dataset.setup() + # dataset.setup() From f6c9a846d20931c6fc87b6cdd769d3cac1737297 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 17 Jul 2026 21:15:23 +0200 Subject: [PATCH 11/56] addd certifi --- chebai/preprocessing/datasets/base.py | 5 +---- chebai/preprocessing/datasets/molecule_classification.py | 2 +- pyproject.toml | 2 +- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 33caec15..94e4e5aa 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -2,7 +2,7 @@ import random from abc import ABC, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union import lightning as pl import numpy as np @@ -15,9 +15,6 @@ from chebai.preprocessing import reader as dr -if TYPE_CHECKING: - import networkx as nx - class XYBaseDataModule(LightningDataModule): """ diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 424fa846..e38a0429 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -3,7 +3,7 @@ import shutil from abc import ABC from tempfile import NamedTemporaryFile -from typing import Any, Generator, List +from typing import Any, Generator from urllib import request import numpy as np diff --git a/pyproject.toml b/pyproject.toml index 7e59f0c5..89f3856d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dev = [ "chebi-utils>=0.1.1", # In case of urllib.error.URLError: Date: Sat, 18 Jul 2026 16:10:38 +0200 Subject: [PATCH 12/56] raiser error if data is empty before loading data into dataloader --- chebai/preprocessing/datasets/base.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 94e4e5aa..b5f00c1a 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -280,6 +280,13 @@ def dataloader(self, kind: str, **kwargs) -> DataLoader: random.shuffle(dataset) if self.data_limit is not None: dataset = dataset[: self.data_limit] + + if len(dataset) == 0: + raise ValueError( + f"Dataset is empty for {kind} data.", + "Please check the data preparation and filtering steps.", + ) + return DataLoader( dataset, collate_fn=self.reader.collator, From 6877db0958ec4fdfda1515be9e195407a34b5315 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 00:27:39 +0200 Subject: [PATCH 13/56] use deepchem to access the data --- chebai/preprocessing/datasets/base.py | 2 +- .../datasets/molecule_classification.py | 191 ++++++++++++------ 2 files changed, 125 insertions(+), 68 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index b5f00c1a..008047d9 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -948,7 +948,7 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: data (pd.DataFrame): The processed dataset to be saved. filename (str): The filename for the pickle file. """ - pd.to_pickle(data, open(os.path.join(self.processed_dir_main, filename), "wb")) + data.to_pickle(open(os.path.join(self.processed_dir_main, filename), "wb")) def get_processed_pickled_df_file(self, filename: str) -> Optional[pd.DataFrame]: """ diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index e38a0429..585a9a12 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -1,12 +1,12 @@ import gzip import os import shutil -from abc import ABC +from abc import ABC, abstractmethod from tempfile import NamedTemporaryFile from typing import Any, Generator from urllib import request -import numpy as np +import deepchem as dc import pandas as pd from chebai.preprocessing import reader as dr @@ -17,29 +17,36 @@ class MoleculeNetDataExtractor(_DynamicDataset, ABC): READER = dr.ChemDataReader - LABLES_COLUMNS = [] - FEATURE_COLUMN_NAME = "smiles" - ID_COLUMN_NAME = None - GROUP_COLUMN_NAME = "group" - @property def _name(self) -> str: """Returns the name of the dataset.""" return str(self.__class__.__name__) - def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> None: + pass + + def _download_required_data(self) -> None: + pass + + def save_processed(self, data: pd.DataFrame, filename: str) -> None: """ - Preprocesses the raw data into a DataFrame. + Save the processed dataset to a pickle file. Args: - raw_data_path (str): Path to the raw data. - - Returns: - pd.DataFrame: The preprocessed data as a DataFrame. + data (pd.DataFrame): The processed dataset to be saved. + filename (str): The filename for the pickle file. """ - return pd.read_csv(raw_data_path, header=0) + if data is not None: + data.to_pickle(open(os.path.join(self.processed_dir_main, filename), "wb")) - def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: + def _get_data_size(self, input_file_path: str) -> None: + pass + + @abstractmethod + def _load_dict( + self, + input_file_path: str, + ) -> Generator[dict[str, Any], None, None]: """Loads data from a CSV file. Args: @@ -48,23 +55,10 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No Returns: List[Dict]: List of data dictionaries. """ - with open(input_file_path, "rb") as input_file: - df = pd.read_pickle(input_file) - - features = df[self.FEATURE_COLUMN_NAME].to_numpy() - if self.ID_COLUMN_NAME is not None and self.ID_COLUMN_NAME in df.columns: - idents = df[self.ID_COLUMN_NAME].to_numpy() - else: - idents = np.arange(len(df)) - labels = df[self.LABLES_COLUMNS].to_numpy() - - if self.GROUP_COLUMN_NAME in df.columns: - groups = df[self.GROUP_COLUMN_NAME].to_numpy() - for feat, labels, ident, group in zip(features, labels, idents, groups): - yield dict(features=feat, labels=labels, ident=ident, group=group) - else: - for feat, labels, ident in zip(features, labels, idents): - yield dict(features=feat, labels=labels, ident=ident) + pass + + def _get_data_splits(self) -> None: + pass @property def base_dir(self) -> str: @@ -76,10 +70,17 @@ def base_dir(self) -> str: """ return os.path.join("data", f"{self._name}:MNClassification") + @property + def raw_file_names_dict(self) -> None: + """Returns a dictionary of raw file names.""" + pass + class ClinTox(MoleculeNetDataExtractor, GroupSplitter): """Data module for ClinTox MoleculeNet dataset.""" + # Total: 1484, FDA_APPROVE is 1: 1390; CT_TOX is 1: 112 + # Multilabel splits?, stratified splits? LABLES_COLUMNS = [ "FDA_APPROVED", "CT_TOX", @@ -90,50 +91,53 @@ def raw_file_names_dict(self) -> dict: """Returns a dictionary of raw file names.""" return {"clintox": "clintox.csv"} - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/clintox.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["clintox"]), - "wt", - ) as fout: - fout.write(gfile.read().decode()) - return os.path.join(self.raw_dir, self.raw_file_names_dict["clintox"]) - class BBBP(MoleculeNetDataExtractor, GroupSplitter): """Data module for BBBP MoleculeNet dataset.""" - ID_COLUMN_NAME = "num" - LABLES_COLUMNS = [ - "p_np", - ] + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: + """Loads data from a CSV file. - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"bbbp": "bbbp.csv"} + Args: + input_file_path (str): Path to the CSV file. - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["bbbp"]), "ab" - ) as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv", - ) as src: - shutil.copyfileobj(src, dst) - return os.path.join(self.raw_dir, self.raw_file_names_dict["bbbp"]) + Returns: + List[Dict]: List of data dictionaries. + """ + splits = [] + tasks, datasets, transformers = dc.molnet.load_bbbp( + featurizer="Raw", splitter="scaffold" + ) + train: dc.data.DiskDataset = datasets[0] + valid: dc.data.DiskDataset = datasets[1] + test: dc.data.DiskDataset = datasets[2] + for split_name, data in [ + ("train", train), + ("valid", valid), + ("test", test), + ]: + for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): + yield dict( + features=smiles, + labels=labels, + ident=idx, + ) + splits.append( + { + "id": idx, + "split": split_name, + } + ) + splits_df = pd.DataFrame(splits) + splits_df.to_csv( + os.path.join(self.processed_dir_main, "splits.csv"), index=False + ) class Sider(MoleculeNetDataExtractor, GroupSplitter): """Data module for Sider MoleculeNet dataset.""" + # Total 1427, multilabel splits, stratified splits? LABLES_COLUMNS = [ "Hepatobiliary disorders", "Metabolism and nutrition disorders", @@ -187,6 +191,10 @@ def _download_required_data(self) -> str: class Bace(MoleculeNetDataExtractor, GeneralSplitter): """Data module for Bace MoleculeNet dataset.""" + # Scaffold? + # TODO: Train, val, test split already marked in data, which to use? + # BINARY CLASSIFICATION task, total 1513, Class 1: 691 + # what are other columns? ID_COLUMN_NAME = "CID" FEATURE_COLUMN_NAME = "mol" LABELS_COLUMNS = [ @@ -213,6 +221,9 @@ def _download_required_data(self) -> str: class HIV(MoleculeNetDataExtractor, GroupSplitter): """Data module for HIV MoleculeNet dataset.""" + # Scaffold? + # HIV: 82255, HIV_active is 1: 1443 + # Why activity not used CI: 39684, CM:1039, CA:404 columns? What are they? LABELS_COLUMNS = [ "HIV_active", ] @@ -237,6 +248,26 @@ def _download_required_data(self) -> str: class MUV(MoleculeNetDataExtractor, GroupSplitter): """Data module for MUV MoleculeNet dataset.""" + # remove row where all nan, or zeros ? + # To much Nan values, Total: 186175 + # 27.0 + # 29.0 + # 30.0 + # 30.0 + # 29.0 + # 29.0 + # 30.0 + # 28.0 + # 29.0 + # 28.0 + # 29.0 + # 29.0 + # 30.0 + # 30.0 + # 29.0 + # 29.0 + # 24.0 + # Multilabel splits ID_COLUMN_NAME = "mol_id" LABELS_COLUMNS = [ "MUV-466", @@ -281,6 +312,32 @@ def _download_required_data(self) -> str: if __name__ == "__main__": # Example usage - dataset = Sider() + dataset = BBBP() dataset.prepare_data() - # dataset.setup() + dataset.setup() + # TODO: add TOX 21, TOX CAST, PCBA, ToxChallenge + # import deepchem as dc + + # # https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + # tasks, datasets, transformers = dc.molnet.load_bbbp( + # featurizer="Raw", splitter="scaffold" + # ) + # # train, valid, test = datasets + # train: dc.data.DiskDataset = datasets[0] + # valid: dc.data.DiskDataset = datasets[1] + # test: dc.data.DiskDataset = datasets[2] + # for data in [train, valid, test]: + # for xi, yi, wi, idi in data.itersamples(): + # print( + # f"features={xi}", # SMILES + # f"labels={yi}", # label (0/1) + # f"ident={idi}", # molecule ID + # ) + + # dc.molnet.load_hiv() + # dc.molnet.load_bace_classification() + # dc.molnet.load_tox21() + # dc.molnet.load_toxcast() + # dc.molnet.load_sider() + # dc.molnet.load_clintox() + # dc.molnet.load_muv() From 20d18700f26a8e3d77e163fb1fd599b7f5ffb217 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 00:59:32 +0200 Subject: [PATCH 14/56] switch entirely to deepchem --- .../datasets/molecule_classification.py | 332 +++++------------- 1 file changed, 88 insertions(+), 244 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 585a9a12..be0fb33b 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -1,20 +1,22 @@ -import gzip import os -import shutil from abc import ABC, abstractmethod -from tempfile import NamedTemporaryFile from typing import Any, Generator -from urllib import request import deepchem as dc import pandas as pd +from deepchem.data import DiskDataset from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import _DynamicDataset -from chebai.preprocessing.splitters import GeneralSplitter, GroupSplitter class MoleculeNetDataExtractor(_DynamicDataset, ABC): + """ + Base class for MoleculeNet dataset extraction and preprocessing. + + Reference: https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + """ + READER = dr.ChemDataReader @property @@ -42,11 +44,7 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: def _get_data_size(self, input_file_path: str) -> None: pass - @abstractmethod - def _load_dict( - self, - input_file_path: str, - ) -> Generator[dict[str, Any], None, None]: + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: """Loads data from a CSV file. Args: @@ -55,6 +53,34 @@ def _load_dict( Returns: List[Dict]: List of data dictionaries. """ + splits = [] + train, valid, test = self._deep_chem_data_loader_api() + for split_name, data in [ + ("train", train), + ("valid", valid), + ("test", test), + ]: + for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): + yield dict( + features=smiles, + labels=labels, + ident=idx, + ) + splits.append( + { + "id": idx, + "split": split_name, + } + ) + splits_df = pd.DataFrame(splits) + splits_df.to_csv( + os.path.join(self.processed_dir_main, "splits.csv"), index=False + ) + + @abstractmethod + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: pass def _get_data_splits(self) -> None: @@ -76,238 +102,82 @@ def raw_file_names_dict(self) -> None: pass -class ClinTox(MoleculeNetDataExtractor, GroupSplitter): +class ClinTox(MoleculeNetDataExtractor): """Data module for ClinTox MoleculeNet dataset.""" - # Total: 1484, FDA_APPROVE is 1: 1390; CT_TOX is 1: 112 - # Multilabel splits?, stratified splits? - LABLES_COLUMNS = [ - "FDA_APPROVED", - "CT_TOX", - ] - - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"clintox": "clintox.csv"} + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_clintox( + featurizer="Raw", splitter="random" + ) + return datasets -class BBBP(MoleculeNetDataExtractor, GroupSplitter): +class BBBP(MoleculeNetDataExtractor): """Data module for BBBP MoleculeNet dataset.""" - def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - splits = [] + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_bbbp( featurizer="Raw", splitter="scaffold" ) - train: dc.data.DiskDataset = datasets[0] - valid: dc.data.DiskDataset = datasets[1] - test: dc.data.DiskDataset = datasets[2] - for split_name, data in [ - ("train", train), - ("valid", valid), - ("test", test), - ]: - for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): - yield dict( - features=smiles, - labels=labels, - ident=idx, - ) - splits.append( - { - "id": idx, - "split": split_name, - } - ) - splits_df = pd.DataFrame(splits) - splits_df.to_csv( - os.path.join(self.processed_dir_main, "splits.csv"), index=False - ) + return datasets -class Sider(MoleculeNetDataExtractor, GroupSplitter): +class Sider(MoleculeNetDataExtractor): """Data module for Sider MoleculeNet dataset.""" - # Total 1427, multilabel splits, stratified splits? - LABLES_COLUMNS = [ - "Hepatobiliary disorders", - "Metabolism and nutrition disorders", - "Product issues", - "Eye disorders", - "Investigations", - "Musculoskeletal and connective tissue disorders", - "Gastrointestinal disorders", - "Social circumstances", - "Immune system disorders", - "Reproductive system and breast disorders", - "Neoplasms benign, malignant and unspecified (incl cysts and polyps)", - "General disorders and administration site conditions", - "Endocrine disorders", - "Surgical and medical procedures", - "Vascular disorders", - "Blood and lymphatic system disorders", - "Skin and subcutaneous tissue disorders", - "Congenital, familial and genetic disorders", - "Infections and infestations", - "Respiratory, thoracic and mediastinal disorders", - "Psychiatric disorders", - "Renal and urinary disorders", - "Pregnancy, puerperium and perinatal conditions", - "Ear and labyrinth disorders", - "Cardiac disorders", - "Nervous system disorders", - "Injury, poisoning and procedural complications", - ] - - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"sider": "sider.csv"} - - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/sider.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["sider"]), "wt" - ) as fout: - fout.write(gfile.read().decode()) - return os.path.join(self.raw_dir, self.raw_file_names_dict["sider"]) - - -class Bace(MoleculeNetDataExtractor, GeneralSplitter): - """Data module for Bace MoleculeNet dataset.""" + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_sider( + featurizer="Raw", splitter="random" + ) + return datasets - # Scaffold? - # TODO: Train, val, test split already marked in data, which to use? - # BINARY CLASSIFICATION task, total 1513, Class 1: 691 - # what are other columns? - ID_COLUMN_NAME = "CID" - FEATURE_COLUMN_NAME = "mol" - LABELS_COLUMNS = [ - "Class", - ] - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"bace": "bace.csv"} +class Bace(MoleculeNetDataExtractor): + """Data module for Bace MoleculeNet dataset.""" - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["bace"]), "ab" - ) as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/bace.csv", - ) as src: - shutil.copyfileobj(src, dst) - return os.path.join(self.raw_dir, self.raw_file_names_dict["bace"]) + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_bace_classification( + featurizer="Raw", splitter="scaffold" + ) + return datasets -class HIV(MoleculeNetDataExtractor, GroupSplitter): +class HIV(MoleculeNetDataExtractor): """Data module for HIV MoleculeNet dataset.""" - # Scaffold? - # HIV: 82255, HIV_active is 1: 1443 - # Why activity not used CI: 39684, CM:1039, CA:404 columns? What are they? - LABELS_COLUMNS = [ - "HIV_active", - ] - - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"hiv": "hiv.csv"} - - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["hiv"]), "ab" - ) as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/HIV.csv", - ) as src: - shutil.copyfileobj(src, dst) - return os.path.join(self.raw_dir, self.raw_file_names_dict["hiv"]) + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_hiv( + featurizer="Raw", splitter="scaffold" + ) + return datasets -class MUV(MoleculeNetDataExtractor, GroupSplitter): +class MUV(MoleculeNetDataExtractor): """Data module for MUV MoleculeNet dataset.""" - # remove row where all nan, or zeros ? - # To much Nan values, Total: 186175 - # 27.0 - # 29.0 - # 30.0 - # 30.0 - # 29.0 - # 29.0 - # 30.0 - # 28.0 - # 29.0 - # 28.0 - # 29.0 - # 29.0 - # 30.0 - # 30.0 - # 29.0 - # 29.0 - # 24.0 - # Multilabel splits - ID_COLUMN_NAME = "mol_id" - LABELS_COLUMNS = [ - "MUV-466", - "MUV-548", - "MUV-600", - "MUV-644", - "MUV-652", - "MUV-689", - "MUV-692", - "MUV-712", - "MUV-713", - "MUV-733", - "MUV-737", - "MUV-810", - "MUV-832", - "MUV-846", - "MUV-852", - "MUV-858", - "MUV-859", - ] - - @property - def raw_file_names_dict(self) -> dict: - """Returns a dictionary of raw file names.""" - return {"muv": "muv.csv"} - - def _download_required_data(self) -> str: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/muv.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open( - os.path.join(self.raw_dir, self.raw_file_names_dict["muv"]), "wt" - ) as fout: - fout.write(gfile.read().decode()) - - return os.path.join(self.raw_dir, self.raw_file_names_dict["muv"]) + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_muv( + featurizer="Raw", splitter="random" + ) + return datasets if __name__ == "__main__": @@ -315,29 +185,3 @@ def _download_required_data(self) -> str: dataset = BBBP() dataset.prepare_data() dataset.setup() - # TODO: add TOX 21, TOX CAST, PCBA, ToxChallenge - # import deepchem as dc - - # # https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html - # tasks, datasets, transformers = dc.molnet.load_bbbp( - # featurizer="Raw", splitter="scaffold" - # ) - # # train, valid, test = datasets - # train: dc.data.DiskDataset = datasets[0] - # valid: dc.data.DiskDataset = datasets[1] - # test: dc.data.DiskDataset = datasets[2] - # for data in [train, valid, test]: - # for xi, yi, wi, idi in data.itersamples(): - # print( - # f"features={xi}", # SMILES - # f"labels={yi}", # label (0/1) - # f"ident={idi}", # molecule ID - # ) - - # dc.molnet.load_hiv() - # dc.molnet.load_bace_classification() - # dc.molnet.load_tox21() - # dc.molnet.load_toxcast() - # dc.molnet.load_sider() - # dc.molnet.load_clintox() - # dc.molnet.load_muv() From 9814e9f868a02c7ccf2d1f78a98d30777c013a04 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 01:14:26 +0200 Subject: [PATCH 15/56] use tox from dc, add new data --- .../datasets/molecule_classification.py | 47 ++++- chebai/preprocessing/datasets/tox21.py | 183 ------------------ 2 files changed, 45 insertions(+), 185 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index be0fb33b..8a8955b8 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -14,7 +14,11 @@ class MoleculeNetDataExtractor(_DynamicDataset, ABC): """ Base class for MoleculeNet dataset extraction and preprocessing. - Reference: https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + Reference: + - https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + - Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, + Aneesh S. Pappu, Karl Leswing, Vijay Pande; MoleculeNet: a benchmark for molecular + machine learning. Chem. Sci. 2018; 9 (2): 513–530. https://doi.org/10.1039/c7sc02664a """ READER = dr.ChemDataReader @@ -173,8 +177,47 @@ class MUV(MoleculeNetDataExtractor): def _deep_chem_data_loader_api( self, ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: - # Random splitting is recommended for this dataset. + # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_muv( + featurizer="Raw", splitter="scaffold" + ) + return datasets + + +class Tox21MolNet(MoleculeNetDataExtractor): + """Data module for Tox21MolNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_tox21( + featurizer="Raw", splitter="random" + ) + return datasets + + +class ToxCast(MoleculeNetDataExtractor): + """Data module for ToxCast MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_toxcast( + featurizer="Raw", splitter="random" + ) + return datasets + + +class PCBA(MoleculeNetDataExtractor): + """Data module for PCBA MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_pcba( featurizer="Raw", splitter="random" ) return datasets diff --git a/chebai/preprocessing/datasets/tox21.py b/chebai/preprocessing/datasets/tox21.py index f6298293..e0f306e6 100644 --- a/chebai/preprocessing/datasets/tox21.py +++ b/chebai/preprocessing/datasets/tox21.py @@ -16,183 +16,6 @@ from chebai.preprocessing.datasets.base import XYBaseDataModule -class Tox21MolNet(XYBaseDataModule): - """Data module for Tox21MolNet dataset.""" - - HEADERS = [ - "NR-AR", - "NR-AR-LBD", - "NR-AhR", - "NR-Aromatase", - "NR-ER", - "NR-ER-LBD", - "NR-PPAR-gamma", - "SR-ARE", - "SR-ATAD5", - "SR-HSE", - "SR-MMP", - "SR-p53", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Tox21MN" - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["tox21.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "tox21.csv"), "wt") as fout: - fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "tox21.csv"))) - groups = np.array([d.get("group") for d in data]) - - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if self._setup_data_flag != 1: - return - - self._setup_data_flag += 1 - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - # self._set_processed_data_props() - self._after_setup() - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - smiles = row["smiles"] - labels = [ - bool(int(float(label))) if len(label) >= 1 else None - for label in (row[k] for k in self.HEADERS) - ] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=row["mol_id"], - # group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=row["mol_id"])) - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - class Tox21Challenge(XYBaseDataModule): """Data module for Tox21Challenge dataset.""" @@ -381,9 +204,3 @@ class Tox21ChallengeChem(Tox21Challenge): """Chemical data reader for Tox21Challenge dataset.""" READER = dr.ChemDataReader - - -class Tox21MolNetChem(Tox21MolNet): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader From beebc7b3268ca493aea0511859a6ed3e233eec82 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 12:45:40 +0200 Subject: [PATCH 16/56] remove tox test --- configs/data/tox21/tox21_moleculenet.yml | 5 +- tests/integration/testTox21MolNetData.py | 164 ----------------------- 2 files changed, 2 insertions(+), 167 deletions(-) delete mode 100644 tests/integration/testTox21MolNetData.py diff --git a/configs/data/tox21/tox21_moleculenet.yml b/configs/data/tox21/tox21_moleculenet.yml index 1e8af70f..75a6c992 100644 --- a/configs/data/tox21/tox21_moleculenet.yml +++ b/configs/data/tox21/tox21_moleculenet.yml @@ -1,5 +1,4 @@ -class_path: chebai.preprocessing.datasets.tox21.Tox21MolNetChem +class_path: chebai.preprocessing.datasets.molecule_classification.Tox21MolNet init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 + diff --git a/tests/integration/testTox21MolNetData.py b/tests/integration/testTox21MolNetData.py deleted file mode 100644 index 36fcb431..00000000 --- a/tests/integration/testTox21MolNetData.py +++ /dev/null @@ -1,164 +0,0 @@ -import os -import unittest -from typing import Dict, List, Tuple - -import torch - -from chebai.preprocessing.datasets.tox21 import Tox21MolNetChem - - -class TestTox21Data(unittest.TestCase): - """ - Unit tests for Tox21 dataset preprocessing. - - Attributes: - tox21 (Tox21MolNetChem): Instance of Tox21 dataset handler. - overlaps_train_val (List): List of overlaps between training and validation sets based on SMILES features. - overlaps_train_test (List): List of overlaps between training and test sets based on SMILES features. - overlaps_val_test (List): List of overlaps between validation and test sets based on SMILES features. - overlaps_train_val_ids (List): List of overlaps between training and validation sets based on SMILES IDs. - overlaps_train_test_ids (List): List of overlaps between training and test sets based on SMILES IDs. - overlaps_val_test_ids (List): List of overlaps between validation and test sets based on SMILES IDs. - """ - - @classmethod - def setUpClass(cls) -> None: - """ - Set up Tox21 dataset and compute overlaps between data splits. - """ - cls.tox21 = Tox21MolNetChem() - cls.getDataSplitsOverlaps() - - @classmethod - def getDataSplitsOverlaps(cls) -> None: - """ - Get the overlap between data splits based on SMILES features and IDs. - """ - processed_path = os.path.join(os.getcwd(), cls.tox21.processed_dir) - print(f"Checking Data from - {processed_path}") - - train_set = torch.load( - os.path.join(processed_path, "train.pt"), weights_only=False - ) - val_set = torch.load( - os.path.join(processed_path, "validation.pt"), weights_only=False - ) - test_set = torch.load( - os.path.join(processed_path, "test.pt"), weights_only=False - ) - - train_smiles, train_smiles_ids = cls.get_features_ids(train_set) - val_smiles, val_smiles_ids = cls.get_features_ids(val_set) - test_smiles, test_smiles_ids = cls.get_features_ids(test_set) - - # Get overlaps based on SMILES features - cls.overlaps_train_val = cls.get_overlaps(train_smiles, val_smiles) - cls.overlaps_train_test = cls.get_overlaps(train_smiles, test_smiles) - cls.overlaps_val_test = cls.get_overlaps(val_smiles, test_smiles) - - # Get overlaps based on SMILES IDs - cls.overlaps_train_val_ids = cls.get_overlaps(train_smiles_ids, val_smiles_ids) - cls.overlaps_train_test_ids = cls.get_overlaps( - train_smiles_ids, test_smiles_ids - ) - cls.overlaps_val_test_ids = cls.get_overlaps(val_smiles_ids, test_smiles_ids) - - @staticmethod - def get_features_ids(data_split: List[Dict]) -> Tuple[List, List]: - """ - Returns SMILES features/tokens and SMILES IDs from the data. - - Args: - data_split (List[Dict]): List of dictionaries containing SMILES features and IDs. - - Returns: - Tuple[List, List]: Tuple of lists containing SMILES features and SMILES IDs. - """ - smiles_features, smiles_ids = [], [] - for entry in data_split: - smiles_features.append(entry["features"]) - smiles_ids.append(entry["ident"]) - - return smiles_features, smiles_ids - - @staticmethod - def get_overlaps(list_1: List, list_2: List) -> List: - """ - Get overlaps between two lists. - - Args: - list_1 (List): First list. - list_2 (List): Second list. - - Returns: - List: List of elements common to both input lists. - """ - overlap = [] - for element in list_1: - if element in list_2: - overlap.append(element) - return overlap - - def test_train_val_overlap_based_on_smiles(self) -> None: - """ - Check that train-val splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_train_val), - 0, - "Duplicate entities present in Train and Validation set based on SMILES", - ) - - def test_train_test_overlap_based_on_smiles(self) -> None: - """ - Check that train-test splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_train_test), - 0, - "Duplicate entities present in Train and Test set based on SMILES", - ) - - def test_val_test_overlap_based_on_smiles(self) -> None: - """ - Check that val-test splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_val_test), - 0, - "Duplicate entities present in Validation and Test set based on SMILES", - ) - - def test_train_val_overlap_based_on_ids(self) -> None: - """ - Check that train-val splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_train_val_ids), - 0, - "Duplicate entities present in Train and Validation set based on IDs", - ) - - def test_train_test_overlap_based_on_ids(self) -> None: - """ - Check that train-test splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_train_test_ids), - 0, - "Duplicate entities present in Train and Test set based on IDs", - ) - - def test_val_test_overlap_based_on_ids(self) -> None: - """ - Check that val-test splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_val_test_ids), - 0, - "Duplicate entities present in Validation and Test set based on IDs", - ) - - -if __name__ == "__main__": - unittest.main() From 50f033c534062e9647d6b400b7f87b60741be277 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 12:46:05 +0200 Subject: [PATCH 17/56] save raw and processed in our custom dir --- .../datasets/molecule_classification.py | 45 +++++++++++++++---- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 8a8955b8..89be4683 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -114,7 +114,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Random splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_clintox( - featurizer="Raw", splitter="random" + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -127,7 +130,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_bbbp( - featurizer="Raw", splitter="scaffold" + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -140,7 +146,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Random splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_sider( - featurizer="Raw", splitter="random" + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -153,7 +162,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_bace_classification( - featurizer="Raw", splitter="scaffold" + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -166,7 +178,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_hiv( - featurizer="Raw", splitter="scaffold" + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -179,7 +194,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Scaffold splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_muv( - featurizer="Raw", splitter="scaffold" + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -192,7 +210,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Random splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_tox21( - featurizer="Raw", splitter="random" + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -205,7 +226,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Random splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_toxcast( - featurizer="Raw", splitter="random" + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets @@ -218,7 +242,10 @@ def _deep_chem_data_loader_api( ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: # Random splitting is recommended for this dataset. tasks, datasets, transformers = dc.molnet.load_pcba( - featurizer="Raw", splitter="random" + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, ) return datasets From 455593506b015a1364d07a698011d88d68c6fb31 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 13:14:57 +0200 Subject: [PATCH 18/56] fix settings json --- .vscode/settings.json | 4 ++-- chebai/preprocessing/datasets/tox21.py | 3 --- configs/data/tox21/tox21_moleculenet.yml | 1 - 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index dbebc3c5..fe1bb4f2 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,7 +2,7 @@ "python.testing.unittestArgs": [ "-v", "-s", - "./tests", + "./tests/unit", "-p", "test*.py" ], @@ -13,4 +13,4 @@ "[python]": { "editor.defaultFormatter": "charliermarsh.ruff" } -} +} \ No newline at end of file diff --git a/chebai/preprocessing/datasets/tox21.py b/chebai/preprocessing/datasets/tox21.py index e0f306e6..da478c14 100644 --- a/chebai/preprocessing/datasets/tox21.py +++ b/chebai/preprocessing/datasets/tox21.py @@ -1,5 +1,4 @@ import csv -import gzip import os import shutil import zipfile @@ -7,10 +6,8 @@ from typing import Dict, Generator, List, Optional from urllib import request -import numpy as np import torch from rdkit import Chem -from sklearn.model_selection import GroupShuffleSplit, train_test_split from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import XYBaseDataModule diff --git a/configs/data/tox21/tox21_moleculenet.yml b/configs/data/tox21/tox21_moleculenet.yml index 75a6c992..33e59a86 100644 --- a/configs/data/tox21/tox21_moleculenet.yml +++ b/configs/data/tox21/tox21_moleculenet.yml @@ -1,4 +1,3 @@ class_path: chebai.preprocessing.datasets.molecule_classification.Tox21MolNet init_args: batch_size: 32 - From 1773c00e33d933a0b2d815b22d9c84598380a27b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 13:15:59 +0200 Subject: [PATCH 19/56] fix pre-commit --- .vscode/settings.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index fe1bb4f2..f81b15ba 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -13,4 +13,4 @@ "[python]": { "editor.defaultFormatter": "charliermarsh.ruff" } -} \ No newline at end of file +} From c6083704b370c9c8200afe0436a73513b34942f1 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 13:39:45 +0200 Subject: [PATCH 20/56] use mol as features instead of smiles as per https://github.com/ChEB-AI/python-chebai/pull/147/ --- chebai/preprocessing/datasets/molecule_classification.py | 2 +- chebai/preprocessing/reader.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py index 89be4683..2d320a77 100644 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ b/chebai/preprocessing/datasets/molecule_classification.py @@ -66,7 +66,7 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No ]: for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): yield dict( - features=smiles, + features=mol, labels=labels, ident=idx, ) diff --git a/chebai/preprocessing/reader.py b/chebai/preprocessing/reader.py index 0dae39cd..0c73a7ef 100644 --- a/chebai/preprocessing/reader.py +++ b/chebai/preprocessing/reader.py @@ -205,8 +205,12 @@ def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]: try: if isinstance(raw_data, str): mol = Chem.MolFromSmiles(raw_data.strip()) - else: + elif isinstance(raw_data, Chem.Mol): mol = raw_data + else: + raise ValueError( + f"Invalid input type: {type(raw_data)}. Expected str or Chem.Mol." + ) if mol is None: raise ValueError(f"Invalid input: {raw_data}") except ValueError as e: From d0804d6c50e4cbad1091168b3d2c7911ed66fd3f Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 14:20:58 +0200 Subject: [PATCH 21/56] fix test, and splitter docstrings --- chebai/preprocessing/datasets/base.py | 4 +- chebai/preprocessing/splitters/__init__.py | 4 +- chebai/preprocessing/splitters/group.py | 24 +-- .../splitters/{general.py => random.py} | 36 ++-- tests/unit/dataset_classes/testTox21MolNet.py | 185 ------------------ 5 files changed, 29 insertions(+), 224 deletions(-) rename chebai/preprocessing/splitters/{general.py => random.py} (68%) delete mode 100644 tests/unit/dataset_classes/testTox21MolNet.py diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 008047d9..d11ae368 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -66,8 +66,8 @@ class XYBaseDataModule(LightningDataModule): def __init__( self, batch_size: int = 1, - test_split: Optional[float] = 0.1, - validation_split: Optional[float] = 0.05, + test_split: float = 0.1, + validation_split: float = 0.05, reader_kwargs: Optional[dict] = None, prediction_kind: str = "test", data_limit: Optional[int] = None, diff --git a/chebai/preprocessing/splitters/__init__.py b/chebai/preprocessing/splitters/__init__.py index 84c176bf..49b0395e 100644 --- a/chebai/preprocessing/splitters/__init__.py +++ b/chebai/preprocessing/splitters/__init__.py @@ -1,5 +1,5 @@ -from .general import GeneralSplitter from .group import GroupSplitter from .multilabel import MultiLabelSplitter +from .random import RandomSplitter -__all__ = ["GroupSplitter", "MultiLabelSplitter", "GeneralSplitter"] +__all__ = ["GroupSplitter", "MultiLabelSplitter", "RandomSplitter"] diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py index 7f1338ec..70bd92dc 100644 --- a/chebai/preprocessing/splitters/group.py +++ b/chebai/preprocessing/splitters/group.py @@ -39,15 +39,15 @@ def create_group_splits( test_ratio: float = 0.1, seed: int | None = 42, ) -> dict[str, pd.DataFrame]: - """Create stratified train/validation/test splits for multilabel DataFrames. + """Create group-based train/validation/test splits for DataFrames. - Columns from index *label_start_col* onwards are treated as binary label - columns (one boolean column per label). The stratification strategy is - chosen automatically based on the number of label columns: - - - More than one label column: ``MultilabelStratifiedShuffleSplit`` from - the ``iterative-stratification`` package. - - Single label column: ``StratifiedShuffleSplit`` from ``scikit-learn``. + Splitting is done with ``GroupShuffleSplit`` using the ``group`` column, + so that all rows sharing the same group value are assigned to the same + split (no group leaks across train/val/test). This is **not** a + stratified split: label balance across splits is not guaranteed, even + though label columns are used to build the ``y`` array passed to the + splitter (``GroupShuffleSplit`` ignores label values and only inspects + the ``groups`` argument). Parameters ---------- @@ -55,7 +55,8 @@ def create_group_splits( Input data. Columns ``0`` to ``label_start_col - 1`` are treated as feature/metadata columns; all remaining columns are boolean label columns. A typical ChEBI DataFrame has columns - ``["chebi_id", "mol", "label1", "label2", ...]``. + ``["chebi_id", "mol", "label1", "label2", ...]``. A ``group`` column + must also be present and is used to keep related rows together. label_start_col : int Index of the first label column (default 2). train_ratio : float @@ -76,8 +77,9 @@ def create_group_splits( Raises ------ ValueError - If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, or - *label_start_col* is out of range. + If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, + *label_start_col* is out of range, the ``group`` column is missing, + or fewer than 2 unique groups are present. """ if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6: raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") diff --git a/chebai/preprocessing/splitters/general.py b/chebai/preprocessing/splitters/random.py similarity index 68% rename from chebai/preprocessing/splitters/general.py rename to chebai/preprocessing/splitters/random.py index 64230130..498a16f7 100644 --- a/chebai/preprocessing/splitters/general.py +++ b/chebai/preprocessing/splitters/random.py @@ -10,7 +10,7 @@ from chebai.preprocessing.datasets.base import _DynamicDataset -class GeneralSplitter(_DynamicDataset, ABC): +class RandomSplitter(_DynamicDataset, ABC): def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: """ Loads encoded/transformed data and generates training, validation, and test splits. @@ -20,9 +20,8 @@ def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: data = self.load_processed_data_from_file(filename) df_data = pd.DataFrame(data) - splits = create_general_splits( + splits = create_random_splits( df_data, - self._LABELS_START_IDX, 1 - self.validation_split - self.test_split, self.validation_split, self.test_split, @@ -31,32 +30,25 @@ def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: return splits["train"], splits["val"], splits["test"] -def create_general_splits( +def create_random_splits( df: pd.DataFrame, - label_start_col: int = 2, train_ratio: float = 0.8, val_ratio: float = 0.1, test_ratio: float = 0.1, + seed: int | None = 42, ) -> dict[str, pd.DataFrame]: - """Create stratified train/validation/test splits for multilabel DataFrames. + """Create random (non-stratified) train/validation/test splits. - Columns from index *label_start_col* onwards are treated as binary label - columns (one boolean column per label). The stratification strategy is - chosen automatically based on the number of label columns: - - - More than one label column: ``MultilabelStratifiedShuffleSplit`` from - the ``iterative-stratification`` package. - - Single label column: ``StratifiedShuffleSplit`` from ``scikit-learn``. + Rows are split purely at random using ``train_test_split`` from + scikit-learn, with no regard to label distribution or grouping. Parameters ---------- df : pd.DataFrame - Input data. Columns ``0`` to ``label_start_col - 1`` are treated as - feature/metadata columns; all remaining columns are boolean label - columns. A typical ChEBI DataFrame has columns - ``["chebi_id", "mol", "label1", "label2", ...]``. + Input data. label_start_col : int - Index of the first label column (default 2). + Index of the first label column (default 2). Unused by this + function; retained for consistency with related split functions. train_ratio : float Fraction of data for training (default 0.8). val_ratio : float @@ -82,17 +74,12 @@ def create_general_splits( raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]): raise ValueError("All ratios must be between 0 and 1") - if label_start_col >= len(df.columns): - raise ValueError( - f"label_start_col={label_start_col} is out of range for a DataFrame " - f"with {len(df.columns)} columns" - ) df_reset = df.reset_index(drop=True) # ── Step 1: carve out the test set ────────────────────────────────────── df_trainval, df_test = train_test_split( - df_reset, test_size=test_ratio, shuffle=True + df_reset, test_size=test_ratio, shuffle=True, random_state=seed ) # ── Step 2: split train/val from the remaining data ───────────────────── @@ -102,6 +89,7 @@ def create_general_splits( df_trainval, test_size=val_ratio_adjusted, shuffle=True, + random_state=seed, ) return { diff --git a/tests/unit/dataset_classes/testTox21MolNet.py b/tests/unit/dataset_classes/testTox21MolNet.py deleted file mode 100644 index 30383524..00000000 --- a/tests/unit/dataset_classes/testTox21MolNet.py +++ /dev/null @@ -1,185 +0,0 @@ -import unittest -from typing import List -from unittest.mock import MagicMock, mock_open, patch - -import torch - -from chebai.preprocessing.datasets.tox21 import Tox21MolNet -from chebai.preprocessing.reader import ChemDataReader -from tests.unit.mock_data.tox_mock_data import Tox21MolNetMockData - - -class TestTox21MolNet(unittest.TestCase): - @classmethod - @patch("os.makedirs", return_value=None) - def setUpClass(cls, mock_makedirs: MagicMock) -> None: - """ - Initialize a Tox21MolNet instance for testing. - - Args: - mock_makedirs (MagicMock): Mocked `os.makedirs` function. - """ - Tox21MolNet.READER = ChemDataReader - cls.data_module = Tox21MolNet() - - @patch( - "builtins.open", - new_callable=mock_open, - read_data=Tox21MolNetMockData.get_raw_data(), - ) - def test_load_data_from_file(self, mock_open_file: mock_open) -> None: - """ - Test the `_load_data_from_file` method for correct output. - - Args: - mock_open_file (mock_open): Mocked open function to simulate file reading. - """ - actual_data: list = self.data_module._load_data_from_file("fake/file/path.csv") - - first_instance = actual_data[0] - - # Check for required keys - required_keys = ["features", "labels", "ident"] - for key in required_keys: - self.assertIn( - key, first_instance, f"'{key}' key is missing in the output data." - ) - - self.assertTrue( - all(isinstance(feature, int) for feature in first_instance["features"]), - "Not all elements in 'features' are integers.", - ) - - # Check that 'features' can be converted to a tensor - features = first_instance["features"] - try: - tensor_features = torch.tensor(features) - self.assertTrue( - tensor_features.ndim > 0, - "'features' should be convertible to a non-empty tensor.", - ) - except Exception as e: - self.fail(f"'features' cannot be converted to a tensor: {str(e)}") - - @patch( - "builtins.open", - new_callable=mock_open, - read_data=Tox21MolNetMockData.get_raw_data(), - ) - @patch("torch.save") - def test_setup_processed_simple_split( - self, - mock_torch_save: MagicMock, - mock_open_file: mock_open, - ) -> None: - """ - Test the `setup_processed` method for basic data splitting and saving. - - Args: - mock_torch_save (MagicMock): Mocked `torch.save` function to avoid actual file writes. - mock_open_file (mock_open): Mocked `open` function to simulate file reading. - """ - self.data_module.setup_processed() - - # Verify if torch.save was called for each split (train, test, validation) - self.assertEqual( - mock_torch_save.call_count, 3, "Expected torch.save to be called 3 times." - ) - call_args_list = mock_torch_save.call_args_list - self.assertIn("test", call_args_list[0][0][1], "Missing 'test' split.") - self.assertIn("train", call_args_list[1][0][1], "Missing 'train' split.") - self.assertIn( - "validation", call_args_list[2][0][1], "Missing 'validation' split." - ) - - # Check for non-overlap between train, test, and validation splits - test_split: List[str] = [d["ident"] for d in call_args_list[0][0][0]] - train_split: List[str] = [d["ident"] for d in call_args_list[1][0][0]] - validation_split: List[str] = [d["ident"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split).isdisjoint(test_split), - "Overlap detected between the train and test splits.", - ) - self.assertTrue( - set(train_split).isdisjoint(validation_split), - "Overlap detected between the train and validation splits.", - ) - self.assertTrue( - set(test_split).isdisjoint(validation_split), - "Overlap detected between the test and validation splits.", - ) - - @patch.object( - Tox21MolNet, - "_load_data_from_file", - return_value=Tox21MolNetMockData.get_processed_grouped_data(), - ) - @patch("torch.save") - def test_setup_processed_with_group_split( - self, mock_torch_save: MagicMock, mock_load_file: MagicMock - ) -> None: - """ - Test the `setup_processed` method for group-based splitting and saving. - - Args: - mock_torch_save (MagicMock): Mocked `torch.save` function to avoid actual file writes. - mock_load_file (MagicMock): Mocked `_load_data_from_file` to provide custom data. - """ - # self.data_module.train_split = 0.5 - # To get the train split as 50%, set test and validation splits to 25% each - # Refer: https://github.com/ChEB-AI/python-chebai/pull/102 - self.data_module.test_split = 0.25 - self.data_module.validation_split = 0.25 - self.data_module.setup_processed() - - # Verify if torch.save was called for each split - self.assertEqual( - mock_torch_save.call_count, 3, "Expected torch.save to be called 3 times." - ) - call_args_list = mock_torch_save.call_args_list - self.assertIn("test", call_args_list[0][0][1], "Missing 'test' split.") - self.assertIn("train", call_args_list[1][0][1], "Missing 'train' split.") - self.assertIn( - "validation", call_args_list[2][0][1], "Missing 'validation' split." - ) - - # Check for non-overlap between train, test, and validation splits (based on 'ident') - test_split: List[str] = [d["ident"] for d in call_args_list[0][0][0]] - train_split: List[str] = [d["ident"] for d in call_args_list[1][0][0]] - validation_split: List[str] = [d["ident"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split).isdisjoint(test_split), - "Overlap detected between the train and test splits (based on 'ident').", - ) - self.assertTrue( - set(train_split).isdisjoint(validation_split), - "Overlap detected between the train and validation splits (based on 'ident').", - ) - self.assertTrue( - set(test_split).isdisjoint(validation_split), - "Overlap detected between the test and validation splits (based on 'ident').", - ) - - # Check for non-overlap between train, test, and validation splits (based on 'group') - test_split_grp: List[str] = [d["group"] for d in call_args_list[0][0][0]] - train_split_grp: List[str] = [d["group"] for d in call_args_list[1][0][0]] - validation_split_grp: List[str] = [d["group"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split_grp).isdisjoint(test_split_grp), - "Overlap detected between the train and test splits (based on 'group').", - ) - self.assertTrue( - set(train_split_grp).isdisjoint(validation_split_grp), - "Overlap detected between the train and validation splits (based on 'group').", - ) - self.assertTrue( - set(test_split_grp).isdisjoint(validation_split_grp), - "Overlap detected between the test and validation splits (based on 'group').", - ) - - -if __name__ == "__main__": - unittest.main() From 3da48764cdd146d8dc02ab33d4c9b34339897817 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 14:42:35 +0200 Subject: [PATCH 22/56] rename to molNet_classification --- .../{molecule_classification.py => molNet_classification.py} | 0 configs/data/moleculenet/bace_moleculenet.yml | 2 +- configs/data/moleculenet/bbbp_moleculenet.yml | 2 +- configs/data/moleculenet/clintox_moleculenet.yml | 2 +- configs/data/moleculenet/hiv_moleculenet.yml | 2 +- configs/data/moleculenet/muv_moleculenet.yml | 2 +- configs/data/moleculenet/sider_moleculenet.yml | 2 +- configs/data/tox21/tox21_moleculenet.yml | 2 +- 8 files changed, 7 insertions(+), 7 deletions(-) rename chebai/preprocessing/datasets/{molecule_classification.py => molNet_classification.py} (100%) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molNet_classification.py similarity index 100% rename from chebai/preprocessing/datasets/molecule_classification.py rename to chebai/preprocessing/datasets/molNet_classification.py diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index f801cc9f..89bbde24 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.Bace +class_path: chebai.preprocessing.datasets.molNet_classification.Bace init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index 2ef99142..aebe1938 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.BBBP +class_path: chebai.preprocessing.datasets.molNet_classification.BBBP init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index 870a6445..0be762fc 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.ClinTox +class_path: chebai.preprocessing.datasets.molNet_classification.ClinTox init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index f3499e26..cbb85024 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.HIV +class_path: chebai.preprocessing.datasets.molNet_classification.HIV init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index d276a76e..a5460fc4 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.MUV +class_path: chebai.preprocessing.datasets.molNet_classification.MUV init_args: batch_size: 32 validation_split: 0.05 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 3b774843..6cf1c9b5 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,4 +1,4 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.Sider +class_path: chebai.preprocessing.datasets.molNet_classification.Sider init_args: batch_size: 10 validation_split: 0.05 diff --git a/configs/data/tox21/tox21_moleculenet.yml b/configs/data/tox21/tox21_moleculenet.yml index 33e59a86..a0b6cd99 100644 --- a/configs/data/tox21/tox21_moleculenet.yml +++ b/configs/data/tox21/tox21_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.Tox21MolNet +class_path: chebai.preprocessing.datasets.molNet_classification.Tox21MolNet init_args: batch_size: 32 From 833727f238d045305e6ef6e3b17f707a292faf1b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 19 Jul 2026 14:50:25 +0200 Subject: [PATCH 23/56] remove split size from config --- configs/data/moleculenet/bace_moleculenet.yml | 2 -- configs/data/moleculenet/bbbp_moleculenet.yml | 2 -- configs/data/moleculenet/clintox_moleculenet.yml | 2 -- configs/data/moleculenet/hiv_moleculenet.yml | 2 -- configs/data/moleculenet/muv_moleculenet.yml | 2 -- configs/data/moleculenet/sider_moleculenet.yml | 2 -- configs/data/{tox21 => moleculenet}/tox21_moleculenet.yml | 2 +- 7 files changed, 1 insertion(+), 13 deletions(-) rename configs/data/{tox21 => moleculenet}/tox21_moleculenet.yml (88%) diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index 89bbde24..9023f37c 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.Bace init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index aebe1938..343e65f2 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.BBBP init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index 0be762fc..7bd7d86a 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.ClinTox init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index cbb85024..3dc9cbff 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.HIV init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index a5460fc4..93fda109 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.MUV init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 6cf1c9b5..09f2f0fa 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,5 +1,3 @@ class_path: chebai.preprocessing.datasets.molNet_classification.Sider init_args: batch_size: 10 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/tox21/tox21_moleculenet.yml b/configs/data/moleculenet/tox21_moleculenet.yml similarity index 88% rename from configs/data/tox21/tox21_moleculenet.yml rename to configs/data/moleculenet/tox21_moleculenet.yml index a0b6cd99..a34c391a 100644 --- a/configs/data/tox21/tox21_moleculenet.yml +++ b/configs/data/moleculenet/tox21_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.Tox21MolNet +class_path: chebai.preprocessing.datasets.molNet_classification.Tox21 init_args: batch_size: 32 From 99ebd143d8dfd0e48ce3cac9e96965c0816219bd Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 22 Jul 2026 11:43:47 +0200 Subject: [PATCH 24/56] move chebi version parameter from base data class to chebi class --- chebai/preprocessing/datasets/base.py | 4 ---- chebai/preprocessing/datasets/chebi.py | 3 +++ 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index d11ae368..67ee2cc8 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -33,7 +33,6 @@ class XYBaseDataModule(LightningDataModule): label_filter (Optional[int]): The index of the label to filter. Default is None. balance_after_filter (Optional[float]): The ratio of negative samples to positive samples after filtering. Default is None. num_workers (int): The number of worker processes for data loading. Default is 1. - chebi_version (int): The version of ChEBI to use. Default is 200. inner_k_folds (int): The number of folds for inner cross-validation. Use -1 to disable inner cross-validation. Default is -1. fold_index (Optional[int]): The index of the fold to use for training and validation. Default is None. base_dir (Optional[str]): The base directory for storing processed and raw data. Default is None. @@ -50,7 +49,6 @@ class XYBaseDataModule(LightningDataModule): label_filter (Optional[int]): The index of the label to filter. balance_after_filter (Optional[float]): The ratio of negative samples to positive samples after filtering. num_workers (int): The number of worker processes for data loading. - chebi_version (int): The version of ChEBI to use. inner_k_folds (int): The number of folds for inner cross-validation. If it is less than to, no cross-validation will be performed. fold_index (Optional[int]): The index of the fold to use for training and validation (only relevant for cross-validation). _base_dir (Optional[str]): The base directory for storing processed and raw data. @@ -75,7 +73,6 @@ def __init__( balance_after_filter: Optional[float] = None, num_workers: int = 1, persistent_workers: bool = True, - chebi_version: int = 200, inner_k_folds: int = -1, # use inner cross-validation if > 1 fold_index: Optional[int] = None, base_dir: Optional[str] = None, @@ -99,7 +96,6 @@ def __init__( self.balance_after_filter = balance_after_filter self.num_workers = num_workers self.persistent_workers: bool = bool(persistent_workers) - self.chebi_version = chebi_version assert type(inner_k_folds) is int self.inner_k_folds = inner_k_folds self.use_inner_cross_validation = ( diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index 9f209bc2..5af23285 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -52,6 +52,7 @@ class _ChEBIDataExtractor(MultiLabelSplitter, _DynamicDataset, ABC): def __init__( self, + chebi_version: int = 241, chebi_version_train: Optional[int] = None, single_class: Optional[int] = None, subset: Optional[Literal["2_STAR", "3_STAR"]] = None, @@ -81,6 +82,8 @@ def __init__( self.subset = subset super(_ChEBIDataExtractor, self).__init__(**kwargs) + self.chebi_version = chebi_version + # use different version of chebi for training and validation (if not None) # (still uses self.chebi_version for test set) self.chebi_version_train = chebi_version_train From bc72d04d64d1807d4571d5f505f408244226e764 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 22 Jul 2026 11:59:55 +0200 Subject: [PATCH 25/56] fix test --- chebai/preprocessing/datasets/chebi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index 5af23285..6845d00d 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -60,6 +60,7 @@ def __init__( aug_smiles_variations: Optional[int] = None, **kwargs, ): + self.chebi_version = chebi_version if bool(augment_smiles): assert int(aug_smiles_variations) > 0, ( "Number of variations must be greater than 0" @@ -82,7 +83,6 @@ def __init__( self.subset = subset super(_ChEBIDataExtractor, self).__init__(**kwargs) - self.chebi_version = chebi_version # use different version of chebi for training and validation (if not None) # (still uses self.chebi_version for test set) From 57881adf1b9fb78fb5a80f3fb2d597db27ce7b05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Simon=20Fl=C3=BCgel?= <43573433+sfluegel05@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:18:13 +0200 Subject: [PATCH 26/56] Rename 'val' to 'validation' in data splits (#175) * Rename 'val' to 'validation' in data splits * Update chebi-utils version requirement in pyproject.toml --- pyproject.toml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 89f3856d..6c00552f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,10 +40,7 @@ dev = [ "omegaconf", "deepsmiles", "torchmetrics", - "chebi-utils>=0.1.1", - # In case of urllib.error.URLError: =0.3", ] linters = [ From 4f3b32597b585e87755bcb4a0c761f9f2a2ffcc5 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Wed, 22 Jul 2026 20:52:50 +0200 Subject: [PATCH 27/56] fix validation key for all splitters --- chebai/preprocessing/splitters/group.py | 6 +++--- chebai/preprocessing/splitters/multilabel.py | 2 +- chebai/preprocessing/splitters/random.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py index 70bd92dc..e4ddf118 100644 --- a/chebai/preprocessing/splitters/group.py +++ b/chebai/preprocessing/splitters/group.py @@ -28,7 +28,7 @@ def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: self.test_split, self.dynamic_data_split_seed, ) - return splits["train"], splits["val"], splits["test"] + return splits["train"], splits["validation"], splits["test"] def create_group_splits( @@ -71,7 +71,7 @@ def create_group_splits( Returns ------- dict - Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each + Dictionary with keys ``'train'``, ``'validation'``, ``'test'``, each containing a DataFrame. Raises @@ -133,6 +133,6 @@ def create_group_splits( return { "train": df_train.reset_index(drop=True), - "val": df_val.reset_index(drop=True), + "validation": df_val.reset_index(drop=True), "test": df_test.reset_index(drop=True), } diff --git a/chebai/preprocessing/splitters/multilabel.py b/chebai/preprocessing/splitters/multilabel.py index 77366ea8..a7c0ced2 100644 --- a/chebai/preprocessing/splitters/multilabel.py +++ b/chebai/preprocessing/splitters/multilabel.py @@ -29,4 +29,4 @@ def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: self.test_split, self.dynamic_data_split_seed, ) - return splits["train"], splits["val"], splits["test"] + return splits["train"], splits["validation"], splits["test"] diff --git a/chebai/preprocessing/splitters/random.py b/chebai/preprocessing/splitters/random.py index 498a16f7..bf407fdc 100644 --- a/chebai/preprocessing/splitters/random.py +++ b/chebai/preprocessing/splitters/random.py @@ -27,7 +27,7 @@ def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: self.test_split, self.dynamic_data_split_seed, ) - return splits["train"], splits["val"], splits["test"] + return splits["train"], splits["validation"], splits["test"] def create_random_splits( @@ -61,7 +61,7 @@ def create_random_splits( Returns ------- dict - Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each + Dictionary with keys ``'train'``, ``'validation'``, ``'test'``, each containing a DataFrame. Raises @@ -94,6 +94,6 @@ def create_random_splits( return { "train": df_train.reset_index(drop=True), - "val": df_val.reset_index(drop=True), + "validation": df_val.reset_index(drop=True), "test": df_test.reset_index(drop=True), } From e563149ed3162d4d982c4b89db84427b3bfa8f36 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 24 Jul 2026 11:48:41 +0200 Subject: [PATCH 28/56] assert if split retrieval fails --- chebai/preprocessing/datasets/base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 67ee2cc8..f036ff86 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -1139,6 +1139,15 @@ def _retrieve_splits_from_csv(self) -> None: self._dynamic_df_train = df_data[df_data["ident"].isin(train_ids)] self._dynamic_df_val = df_data[df_data["ident"].isin(validation_ids)] self._dynamic_df_test = df_data[df_data["ident"].isin(test_ids)] + assert len(self._dynamic_df_train) > 0, ( + "No training data found after applying splits" + ) + assert len(self._dynamic_df_val) > 0, ( + "No validation data found after applying splits" + ) + assert len(self._dynamic_df_test) > 0, ( + "No test data found after applying splits" + ) # ------------------------------ Phase: DataLoaders ----------------------------------- def load_processed_data( From beff841a80424fe5a24c712fcda73d80081701bf Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 14:29:22 +0200 Subject: [PATCH 29/56] rename to moleculenet_classification --- .../{molNet_classification.py => molculenet_classification.py} | 0 configs/data/moleculenet/bace_moleculenet.yml | 2 +- configs/data/moleculenet/bbbp_moleculenet.yml | 2 +- configs/data/moleculenet/clintox_moleculenet.yml | 2 +- configs/data/moleculenet/hiv_moleculenet.yml | 2 +- configs/data/moleculenet/muv_moleculenet.yml | 2 +- configs/data/moleculenet/sider_moleculenet.yml | 2 +- configs/data/moleculenet/tox21_moleculenet.yml | 2 +- pyproject.toml | 3 +++ 9 files changed, 10 insertions(+), 7 deletions(-) rename chebai/preprocessing/datasets/{molNet_classification.py => molculenet_classification.py} (100%) diff --git a/chebai/preprocessing/datasets/molNet_classification.py b/chebai/preprocessing/datasets/molculenet_classification.py similarity index 100% rename from chebai/preprocessing/datasets/molNet_classification.py rename to chebai/preprocessing/datasets/molculenet_classification.py diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index 9023f37c..97c6c6f1 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.Bace +class_path: chebai.preprocessing.datasets.molculenet_classification.Bace init_args: batch_size: 32 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index 343e65f2..b7d2ab87 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.BBBP +class_path: chebai.preprocessing.datasets.molculenet_classification.BBBP init_args: batch_size: 32 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index 7bd7d86a..948ab4c3 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.ClinTox +class_path: chebai.preprocessing.datasets.molculenet_classification.ClinTox init_args: batch_size: 32 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index 3dc9cbff..d4f3ea00 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.HIV +class_path: chebai.preprocessing.datasets.molculenet_classification.HIV init_args: batch_size: 32 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index 93fda109..e3cc8905 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.MUV +class_path: chebai.preprocessing.datasets.molculenet_classification.MUV init_args: batch_size: 32 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 09f2f0fa..4682c83d 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.Sider +class_path: chebai.preprocessing.datasets.molculenet_classification.Sider init_args: batch_size: 10 diff --git a/configs/data/moleculenet/tox21_moleculenet.yml b/configs/data/moleculenet/tox21_moleculenet.yml index a34c391a..4759ffd7 100644 --- a/configs/data/moleculenet/tox21_moleculenet.yml +++ b/configs/data/moleculenet/tox21_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molNet_classification.Tox21 +class_path: chebai.preprocessing.datasets.molculenet_classification.Tox21 init_args: batch_size: 32 diff --git a/pyproject.toml b/pyproject.toml index 6c00552f..26abb251 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,9 @@ dev = [ "deepsmiles", "torchmetrics", "chebi-utils>=0.3", + # In case of urllib.error.URLError: Date: Sat, 25 Jul 2026 14:46:17 +0200 Subject: [PATCH 30/56] rename to molecule_net_classification --- ...culenet_classification.py => molecule_net_classification.py} | 2 +- configs/data/moleculenet/bace_moleculenet.yml | 2 +- configs/data/moleculenet/bbbp_moleculenet.yml | 2 +- configs/data/moleculenet/clintox_moleculenet.yml | 2 +- configs/data/moleculenet/hiv_moleculenet.yml | 2 +- configs/data/moleculenet/muv_moleculenet.yml | 2 +- configs/data/moleculenet/sider_moleculenet.yml | 2 +- configs/data/moleculenet/tox21_moleculenet.yml | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) rename chebai/preprocessing/datasets/{molculenet_classification.py => molecule_net_classification.py} (99%) diff --git a/chebai/preprocessing/datasets/molculenet_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py similarity index 99% rename from chebai/preprocessing/datasets/molculenet_classification.py rename to chebai/preprocessing/datasets/molecule_net_classification.py index 2d320a77..cc1c9855 100644 --- a/chebai/preprocessing/datasets/molculenet_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -202,7 +202,7 @@ def _deep_chem_data_loader_api( return datasets -class Tox21MolNet(MoleculeNetDataExtractor): +class Tox21(MoleculeNetDataExtractor): """Data module for Tox21MolNet dataset.""" def _deep_chem_data_loader_api( diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index 97c6c6f1..13f01309 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.Bace +class_path: chebai.preprocessing.datasets.molecule_net_classification.Bace init_args: batch_size: 32 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index b7d2ab87..8ad18668 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.BBBP +class_path: chebai.preprocessing.datasets.molecule_net_classification.BBBP init_args: batch_size: 32 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index 948ab4c3..a389f725 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.ClinTox +class_path: chebai.preprocessing.datasets.molecule_net_classification.ClinTox init_args: batch_size: 32 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index d4f3ea00..d1febe83 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.HIV +class_path: chebai.preprocessing.datasets.molecule_net_classification.HIV init_args: batch_size: 32 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index e3cc8905..9e02496d 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.MUV +class_path: chebai.preprocessing.datasets.molecule_net_classification.MUV init_args: batch_size: 32 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 4682c83d..c33beb1d 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.Sider +class_path: chebai.preprocessing.datasets.molecule_net_classification.Sider init_args: batch_size: 10 diff --git a/configs/data/moleculenet/tox21_moleculenet.yml b/configs/data/moleculenet/tox21_moleculenet.yml index 4759ffd7..8ce308e9 100644 --- a/configs/data/moleculenet/tox21_moleculenet.yml +++ b/configs/data/moleculenet/tox21_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molculenet_classification.Tox21 +class_path: chebai.preprocessing.datasets.molecule_net_classification.Tox21 init_args: batch_size: 32 From 2a624dadc3f006c1158a6a1156927903e777296f Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 15:06:33 +0200 Subject: [PATCH 31/56] CLASS name captilazie --- chebai/preprocessing/datasets/molecule_net_classification.py | 4 ++-- configs/data/moleculenet/bace_moleculenet.yml | 2 +- configs/data/moleculenet/sider_moleculenet.yml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index cc1c9855..06311ee0 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -138,7 +138,7 @@ def _deep_chem_data_loader_api( return datasets -class Sider(MoleculeNetDataExtractor): +class SIDER(MoleculeNetDataExtractor): """Data module for Sider MoleculeNet dataset.""" def _deep_chem_data_loader_api( @@ -154,7 +154,7 @@ def _deep_chem_data_loader_api( return datasets -class Bace(MoleculeNetDataExtractor): +class BACE(MoleculeNetDataExtractor): """Data module for Bace MoleculeNet dataset.""" def _deep_chem_data_loader_api( diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index 13f01309..da9736c5 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_net_classification.Bace +class_path: chebai.preprocessing.datasets.molecule_net_classification.BACE init_args: batch_size: 32 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index c33beb1d..2ae64c2e 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,3 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_net_classification.Sider +class_path: chebai.preprocessing.datasets.molecule_net_classification.SIDER init_args: batch_size: 10 From 4c95eb36bea5032b011e0eb41a3433f4f32d2855 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 16:12:52 +0200 Subject: [PATCH 32/56] add multilabel avg precision --- configs/metrics/binary-f1-roc-auc.yml | 2 +- configs/metrics/micro-macro-f1-roc-auc.yml | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/metrics/binary-f1-roc-auc.yml b/configs/metrics/binary-f1-roc-auc.yml index 05834343..d87bb04f 100644 --- a/configs/metrics/binary-f1-roc-auc.yml +++ b/configs/metrics/binary-f1-roc-auc.yml @@ -1,6 +1,6 @@ class_path: torchmetrics.MetricCollection init_args: - metrics: + metrics: # Use this for: BACE, BBBP, HIV f1: class_path: torchmetrics.classification.BinaryF1Score roc-auc: diff --git a/configs/metrics/micro-macro-f1-roc-auc.yml b/configs/metrics/micro-macro-f1-roc-auc.yml index c659b877..9e7b0d45 100644 --- a/configs/metrics/micro-macro-f1-roc-auc.yml +++ b/configs/metrics/micro-macro-f1-roc-auc.yml @@ -1,6 +1,6 @@ class_path: torchmetrics.MetricCollection init_args: - metrics: + metrics: # Use this for: SIDER, ClinTox, Tox21, ToxCast micro-f1: class_path: torchmetrics.classification.MultilabelF1Score init_args: @@ -9,3 +9,5 @@ init_args: class_path: chebai.callbacks.epoch_metrics.MacroF1 roc-auc: class_path: torchmetrics.classification.MultilabelAUROC + pr-auc: # Especially used for MUV, PCBA + class_path: torchmetrics.classification.MultilabelAveragePrecision From d059192c547a280a11579454a48bb993db1a538c Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 7 Aug 2026 13:21:07 +0200 Subject: [PATCH 33/56] add data type prop for data specific token files --- chebai/preprocessing/datasets/base.py | 8 +++++++ chebai/preprocessing/datasets/chebi.py | 23 +++++++++++++++++++ .../datasets/molecule_net_classification.py | 8 +++++++ 3 files changed, 39 insertions(+) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index f036ff86..2622ca4d 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -642,6 +642,14 @@ def classes_txt_file_path(self) -> Optional[str]: # - chebai/cli.py: to link this property to `model.init_args.classes_txt_file_path` return None + @property + def data_type(self) -> str: + """ + Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. + This property is used to create a separate tokens directory for each data type. + """ + raise NotImplementedError + class MergedDataset(XYBaseDataModule): MERGED = [] diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index 6845d00d..1e8dfcfe 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -530,6 +530,14 @@ def classes_txt_file_path(self) -> str: # - chebai/cli.py: to link this property to `model.init_args.classes_txt_file_path` return os.path.join(self.processed_dir_main, "classes.txt") + @property + def data_type(self) -> str: + """ + Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. + This property is used to create a separate tokens directory for each data type. + """ + return "chebi" + class ChEBIFromList(_ChEBIDataExtractor): """ @@ -648,6 +656,21 @@ class ChEBIOver50(ChEBIOverX): THRESHOLD: int = 50 +class ChEBIOver50_ChemDataReader(ChEBIOverX): + """ + A class for extracting data from the ChEBI dataset with a threshold of 50 for selecting classes. + + Inherits from ChEBIOverX. + + Attributes: + THRESHOLD (int): The threshold for selecting classes (50). + """ + + READER = dr.ChemDataReader + + THRESHOLD: int = 50 + + class ChEBIOver25(ChEBIOverX): """ A class for extracting data from the ChEBI dataset with a threshold of 25 for selecting classes. diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 06311ee0..442deb35 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -105,6 +105,14 @@ def raw_file_names_dict(self) -> None: """Returns a dictionary of raw file names.""" pass + @property + def data_type(self) -> str: + """ + Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. + This property is used to create a separate tokens directory for each data type. + """ + return self.__class__.__name__ + class ClinTox(MoleculeNetDataExtractor): """Data module for ClinTox MoleculeNet dataset.""" From 5f76acb9a0f23ed5223b433ef07a1ae89a582e3a Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 7 Aug 2026 14:08:35 +0200 Subject: [PATCH 34/56] suggested changes --- chebai/preprocessing/datasets/base.py | 8 +++----- .../preprocessing/datasets/molecule_net_classification.py | 4 ++-- chebai/preprocessing/splitters/group.py | 2 +- chebai/preprocessing/splitters/random.py | 2 +- pyproject.toml | 1 + 5 files changed, 8 insertions(+), 9 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 8eda561c..ebd41a6a 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -7,11 +7,11 @@ import lightning as pl import numpy as np import pandas as pd -from rdkit import Chem import torch import tqdm from lightning.pytorch.core.datamodule import LightningDataModule from lightning_utilities.core.rank_zero import rank_zero_info +from rdkit import Chem from torch.utils.data import DataLoader from chebai.preprocessing import reader as dr @@ -280,8 +280,7 @@ def dataloader(self, kind: str, **kwargs) -> DataLoader: if len(dataset) == 0: raise ValueError( - f"Dataset is empty for {kind} data.", - "Please check the data preparation and filtering steps.", + f"Dataset is empty for {kind} data.\nPlease check the data preparation and filtering steps.", ) return DataLoader( @@ -618,7 +617,6 @@ def raw_file_names(self) -> List[str]: return list(self.raw_file_names_dict.values()) @property - @abstractmethod def raw_file_names_dict(self) -> dict: """ Returns the dictionary of raw file names (i.e., files that are directly obtained from an external source). @@ -953,7 +951,7 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: data (pd.DataFrame): The processed dataset to be saved. filename (str): The filename for the pickle file. """ - data.to_pickle(open(os.path.join(self.processed_dir_main, filename), "wb")) + data.to_pickle(os.path.join(self.processed_dir_main, filename)) def get_processed_pickled_df_file(self, filename: str) -> Optional[pd.DataFrame]: """ diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 442deb35..284c5ea4 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -43,7 +43,7 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: filename (str): The filename for the pickle file. """ if data is not None: - data.to_pickle(open(os.path.join(self.processed_dir_main, filename), "wb")) + data.to_pickle(os.path.join(self.processed_dir_main, filename)) def _get_data_size(self, input_file_path: str) -> None: pass @@ -61,7 +61,7 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No train, valid, test = self._deep_chem_data_loader_api() for split_name, data in [ ("train", train), - ("valid", valid), + ("validation", valid), ("test", test), ]: for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py index e4ddf118..7659b1ac 100644 --- a/chebai/preprocessing/splitters/group.py +++ b/chebai/preprocessing/splitters/group.py @@ -1,4 +1,4 @@ -"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" +"""Generate group-based train/validation/test splits from DataFrames.""" from __future__ import annotations diff --git a/chebai/preprocessing/splitters/random.py b/chebai/preprocessing/splitters/random.py index bf407fdc..bd668d3a 100644 --- a/chebai/preprocessing/splitters/random.py +++ b/chebai/preprocessing/splitters/random.py @@ -1,4 +1,4 @@ -"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" + """Generate random (non-stratified) train/validation/test splits from DataFrames.""" from __future__ import annotations diff --git a/pyproject.toml b/pyproject.toml index 6a0ee94d..5143888e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dev = [ # In case of urllib.error.URLError: Date: Fri, 7 Aug 2026 14:44:09 +0200 Subject: [PATCH 35/56] fix ident --- chebai/preprocessing/splitters/random.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chebai/preprocessing/splitters/random.py b/chebai/preprocessing/splitters/random.py index bd668d3a..b84490ba 100644 --- a/chebai/preprocessing/splitters/random.py +++ b/chebai/preprocessing/splitters/random.py @@ -1,4 +1,4 @@ - """Generate random (non-stratified) train/validation/test splits from DataFrames.""" +"""Generate random (non-stratified) train/validation/test splits from DataFrames.""" from __future__ import annotations From 2f18c58bb219d8d8867aa64aac5f5959956a2aea Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 19:57:10 +0200 Subject: [PATCH 36/56] fix pr-auc link error --- chebai/cli.py | 1 + 1 file changed, 1 insertion(+) diff --git a/chebai/cli.py b/chebai/cli.py index d65dd51e..fccf4de5 100644 --- a/chebai/cli.py +++ b/chebai/cli.py @@ -75,6 +75,7 @@ def call_data_methods(data: Type[XYBaseDataModule]): "mse", "rmse", "r2", + "pr-auc", ): # When using lightning > 2.5.1 then need to uncomment all metrics that are not used # for average in ("mse", "rmse","r2"): # for regression From b6b5ed0b7afd7abe4843af99902e154bb213b740 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 19:57:31 +0200 Subject: [PATCH 37/56] data type for each molenet class --- .../datasets/molecule_net_classification.py | 43 +++++++++++++++---- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 284c5ea4..e2a597d0 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -105,14 +105,6 @@ def raw_file_names_dict(self) -> None: """Returns a dictionary of raw file names.""" pass - @property - def data_type(self) -> str: - """ - Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. - This property is used to create a separate tokens directory for each data type. - """ - return self.__class__.__name__ - class ClinTox(MoleculeNetDataExtractor): """Data module for ClinTox MoleculeNet dataset.""" @@ -129,6 +121,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "clin_tox" + class BBBP(MoleculeNetDataExtractor): """Data module for BBBP MoleculeNet dataset.""" @@ -145,6 +141,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "bbbp" + class SIDER(MoleculeNetDataExtractor): """Data module for Sider MoleculeNet dataset.""" @@ -161,6 +161,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "sider" + class BACE(MoleculeNetDataExtractor): """Data module for Bace MoleculeNet dataset.""" @@ -177,6 +181,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "bace" + class HIV(MoleculeNetDataExtractor): """Data module for HIV MoleculeNet dataset.""" @@ -193,6 +201,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "hiv" + class MUV(MoleculeNetDataExtractor): """Data module for MUV MoleculeNet dataset.""" @@ -209,6 +221,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "muv" + class Tox21(MoleculeNetDataExtractor): """Data module for Tox21MolNet dataset.""" @@ -225,6 +241,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "tox21" + class ToxCast(MoleculeNetDataExtractor): """Data module for ToxCast MoleculeNet dataset.""" @@ -241,6 +261,9 @@ def _deep_chem_data_loader_api( ) return datasets + def data_type(self) -> str: + return "toxcast" + class PCBA(MoleculeNetDataExtractor): """Data module for PCBA MoleculeNet dataset.""" @@ -257,6 +280,10 @@ def _deep_chem_data_loader_api( ) return datasets + @property + def data_type(self) -> str: + return "pcba" + if __name__ == "__main__": # Example usage From 18a2f662510cdccede3aabf276a0c89361e4da4e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 22:07:47 +0200 Subject: [PATCH 38/56] fix diff dir for molenet issue --- .../datasets/molecule_net_classification.py | 44 ++++++++++++++++--- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index e2a597d0..41ace1e1 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -23,11 +23,6 @@ class MoleculeNetDataExtractor(_DynamicDataset, ABC): READER = dr.ChemDataReader - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return str(self.__class__.__name__) - def _preprocess_data_into_dataframe(self, raw_data_path: str) -> None: pass @@ -125,6 +120,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "clin_tox" + @property + def _name(self) -> str: + return "ClinTox" + class BBBP(MoleculeNetDataExtractor): """Data module for BBBP MoleculeNet dataset.""" @@ -145,6 +144,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "bbbp" + @property + def _name(self) -> str: + return "BBBP" + class SIDER(MoleculeNetDataExtractor): """Data module for Sider MoleculeNet dataset.""" @@ -165,6 +168,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "sider" + @property + def _name(self) -> str: + return "SIDER" + class BACE(MoleculeNetDataExtractor): """Data module for Bace MoleculeNet dataset.""" @@ -185,6 +192,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "bace" + @property + def _name(self) -> str: + return "BACE" + class HIV(MoleculeNetDataExtractor): """Data module for HIV MoleculeNet dataset.""" @@ -205,6 +216,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "hiv" + @property + def _name(self) -> str: + return "HIV" + class MUV(MoleculeNetDataExtractor): """Data module for MUV MoleculeNet dataset.""" @@ -225,6 +240,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "muv" + @property + def _name(self) -> str: + return "MUV" + class Tox21(MoleculeNetDataExtractor): """Data module for Tox21MolNet dataset.""" @@ -245,6 +264,10 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "tox21" + @property + def _name(self) -> str: + return "Tox21" + class ToxCast(MoleculeNetDataExtractor): """Data module for ToxCast MoleculeNet dataset.""" @@ -261,9 +284,14 @@ def _deep_chem_data_loader_api( ) return datasets + @property def data_type(self) -> str: return "toxcast" + @property + def _name(self) -> str: + return "ToxCast" + class PCBA(MoleculeNetDataExtractor): """Data module for PCBA MoleculeNet dataset.""" @@ -284,9 +312,13 @@ def _deep_chem_data_loader_api( def data_type(self) -> str: return "pcba" + @property + def _name(self) -> str: + return "PCBA" + if __name__ == "__main__": # Example usage - dataset = BBBP() + dataset = BACE() dataset.prepare_data() dataset.setup() From 6f6e310b5774fd43641a8e36b3ec44d7bbf67031 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 23:37:20 +0200 Subject: [PATCH 39/56] fix data empty error + prompt user to use split file path --- chebai/preprocessing/datasets/base.py | 3 +- .../datasets/molecule_net_classification.py | 31 ++++++++++++++++--- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index ebd41a6a..3b39a9da 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -951,7 +951,8 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: data (pd.DataFrame): The processed dataset to be saved. filename (str): The filename for the pickle file. """ - data.to_pickle(os.path.join(self.processed_dir_main, filename)) + if data is not None and not data.empty: + data.to_pickle(os.path.join(self.processed_dir_main, filename)) def get_processed_pickled_df_file(self, filename: str) -> Optional[pd.DataFrame]: """ diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 41ace1e1..86b1bcfe 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -21,6 +21,22 @@ class MoleculeNetDataExtractor(_DynamicDataset, ABC): machine learning. Chem. Sci. 2018; 9 (2): 513–530. https://doi.org/10.1039/c7sc02664a """ + def __init__( + self, + test_split: float | None = None, + validation_split: float | None = None, + **kwargs, + ): + if test_split is not None or validation_split is not None: + raise ValueError( + "Custom splits are not supported for MoleculeNet datasets. " + "Please use the predefined splits provided by the deepchem community" + "by using `--splits_file_path=`" + ) + super().__init__( + test_split=test_split, validation_split=validation_split, **kwargs + ) + READER = dr.ChemDataReader def _preprocess_data_into_dataframe(self, raw_data_path: str) -> None: @@ -71,10 +87,10 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No "split": split_name, } ) - splits_df = pd.DataFrame(splits) - splits_df.to_csv( - os.path.join(self.processed_dir_main, "splits.csv"), index=False - ) + splits_file_path = os.path.join(self.processed_dir_main, "splits.csv") + if not os.path.exists(splits_file_path): + splits_df = pd.DataFrame(splits) + splits_df.to_csv(splits_file_path, index=False) @abstractmethod def _deep_chem_data_loader_api( @@ -85,6 +101,13 @@ def _deep_chem_data_loader_api( def _get_data_splits(self) -> None: pass + def _generate_dynamic_splits(self) -> None: + raise ValueError( + "Custom splits are not supported for MoleculeNet datasets. " + "Please use the predefined splits provided by the deepchem community" + "by using `--splits_file_path=`" + ) + @property def base_dir(self) -> str: """ From 05492ad3198e1af0658b1bcf03eae7747baea61d Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 23:37:58 +0200 Subject: [PATCH 40/56] pre-commit run -a --- chebai/loss/asymmetric_loss.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/chebai/loss/asymmetric_loss.py b/chebai/loss/asymmetric_loss.py index 9ed6f170..8b795907 100644 --- a/chebai/loss/asymmetric_loss.py +++ b/chebai/loss/asymmetric_loss.py @@ -9,10 +9,10 @@ class AsymmetricLoss(nn.Module): Asymmetric Loss for multi-label and single-label classification tasks. Implementation adapted from: https://github.com/Alibaba-MIIL/ASL - + Asymmetric Loss from: "Asymmetric Loss For Multi-Label Classification" https://openaccess.thecvf.com/content/ICCV2021/papers/Ben-Baruch_Asymmetric_Loss_For_Multi-Label_Classification_ICCV_2021_paper.pdf - + Args: gamma_neg (float): Negative focusing parameter. Default is 1. gamma_pos (float): Positive focusing parameter. Default is 1. @@ -47,12 +47,12 @@ def __init__( def forward(self, inputs, targets, **kwargs): """ Forward pass to compute the Asymmetric Loss. - + Args: inputs: Predictions (logits) from the model. targets: Ground truth labels. **kwargs: Additional keyword arguments (for compatibility with training framework). - + Returns: Loss tensor with the reduction option applied. """ @@ -68,7 +68,7 @@ def forward(self, inputs, targets, **kwargs): def _multi_label_asymmetric_loss(self, x, y): """ Standard asymmetric loss for multi-label classification. - + Parameters ---------- x: input logits @@ -112,12 +112,14 @@ def _multi_label_asymmetric_loss(self, x, y): def _single_label_asymmetric_loss(self, inputs, target): """ Asymmetric loss for single-label classification problems. - + "input" dimensions: - (batch_size, number_classes) "target" dimensions: - (batch_size) """ log_preds = self.logsoftmax(inputs) - self.targets_classes = torch.zeros_like(inputs).scatter_(1, target.long().unsqueeze(1), 1) + self.targets_classes = torch.zeros_like(inputs).scatter_( + 1, target.long().unsqueeze(1), 1 + ) # ASL weights targets = self.targets_classes @@ -125,13 +127,15 @@ def _single_label_asymmetric_loss(self, inputs, target): xs_pos = torch.exp(log_preds) xs_neg = 1 - xs_pos grad_ctx = ( - torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext() - ) + torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext() + ) with grad_ctx: xs_pos = xs_pos * targets xs_neg = xs_neg * anti_targets - asymmetric_w = torch.pow(1 - xs_pos - xs_neg, - self.gamma_pos * targets + self.gamma_neg * anti_targets) + asymmetric_w = torch.pow( + 1 - xs_pos - xs_neg, + self.gamma_pos * targets + self.gamma_neg * anti_targets, + ) log_preds = log_preds * asymmetric_w # loss calculation From 7a93792f8aad0d9f37dd840858c9cbc5ba58926e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 11 Aug 2026 23:59:08 +0200 Subject: [PATCH 41/56] fix duplicate entries in molnet datasets --- chebai/preprocessing/datasets/base.py | 7 +++++++ .../preprocessing/datasets/molecule_net_classification.py | 4 +++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 3b39a9da..0ea7f453 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -341,6 +341,13 @@ def _load_data_from_file(self, path: str) -> List[Dict[str, Any]]: if d["features"] is not None ] + number_of_unique_ids = len(set(d["ident"] for d in data)) + assert len(data) == number_of_unique_ids, ( + "Duplicate entries found in the dataset. " + f"Unique entries {number_of_unique_ids}. " + f"Total entries {len(data)}. " + ) + data = [val for val in data if self._filter_to_token_limit(val)] return data diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 86b1bcfe..6f1ad771 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -70,12 +70,13 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No """ splits = [] train, valid, test = self._deep_chem_data_loader_api() + idx = 0 for split_name, data in [ ("train", train), ("validation", valid), ("test", test), ]: - for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): + for mol, labels, wi, smiles in data.itersamples(): yield dict( features=mol, labels=labels, @@ -87,6 +88,7 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No "split": split_name, } ) + idx += 1 splits_file_path = os.path.join(self.processed_dir_main, "splits.csv") if not os.path.exists(splits_file_path): splits_df = pd.DataFrame(splits) From e3df68128976e1f0d7e56886e643e91a7188fa07 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 13 Aug 2026 17:38:25 +0200 Subject: [PATCH 42/56] convert labels to bool dtype --- .../datasets/molecule_net_classification.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 6f1ad771..e4af274c 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -77,17 +77,8 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No ("test", test), ]: for mol, labels, wi, smiles in data.itersamples(): - yield dict( - features=mol, - labels=labels, - ident=idx, - ) - splits.append( - { - "id": idx, - "split": split_name, - } - ) + yield dict(features=mol, labels=labels.astype(bool), ident=idx) + splits.append({"id": idx, "split": split_name}) idx += 1 splits_file_path = os.path.join(self.processed_dir_main, "splits.csv") if not os.path.exists(splits_file_path): From f778893251a5fac37ba315689803dba88e728746 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 15 Aug 2026 14:26:16 +0200 Subject: [PATCH 43/56] ignore labels with 0 weights --- .../datasets/molecule_net_classification.py | 26 ++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index e4af274c..3beef041 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -76,8 +76,19 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No ("validation", valid), ("test", test), ]: - for mol, labels, wi, smiles in data.itersamples(): - yield dict(features=mol, labels=labels.astype(bool), ident=idx) + for mol, labels, w, smiles in data.itersamples(): + # https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + # Note that the “w” matrix represents the weight of each sample. Some assays may have missing values, in which case the weight is 0. + # Otherwise, the weight is 1. This is when `transformers` are set to `[]` in the deepchem data loader API. + # By default `transformers` is set to ['balancing'], which doesn't hand you the raw 0/1 weight matrix. + # Instead meaning DeepChem automatically applies a BalancingTransformer before returning the dataset. + # That transformer reweights the observed labels per task so positive and negative examples end up with equal total weight — it upweights the rarer class. + # Currently, transformers are set to `[]` in the deepchem data loader API, so we can get the raw 0/1 weight matrix. + labels = [ + bool(label) if weight != 0 else None + for label, weight in zip(labels, w) + ] + yield dict(features=mol, labels=labels, ident=idx) splits.append({"id": idx, "split": split_name}) idx += 1 splits_file_path = os.path.join(self.processed_dir_main, "splits.csv") @@ -129,6 +140,7 @@ def _deep_chem_data_loader_api( splitter="random", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -153,6 +165,7 @@ def _deep_chem_data_loader_api( splitter="scaffold", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -177,6 +190,7 @@ def _deep_chem_data_loader_api( splitter="random", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -201,6 +215,7 @@ def _deep_chem_data_loader_api( splitter="scaffold", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -225,6 +240,7 @@ def _deep_chem_data_loader_api( splitter="scaffold", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -249,6 +265,7 @@ def _deep_chem_data_loader_api( splitter="scaffold", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -273,6 +290,7 @@ def _deep_chem_data_loader_api( splitter="random", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -297,6 +315,7 @@ def _deep_chem_data_loader_api( splitter="random", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -321,6 +340,7 @@ def _deep_chem_data_loader_api( splitter="random", data_dir=self.raw_dir, save_dir=self.processed_dir_main, + transformers=[], ) return datasets @@ -335,6 +355,6 @@ def _name(self) -> str: if __name__ == "__main__": # Example usage - dataset = BACE() + dataset = SIDER() dataset.prepare_data() dataset.setup() From af8645dc7d7474b17d37dc4a7aec0b721ab2189b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 15 Aug 2026 21:53:56 +0200 Subject: [PATCH 44/56] fix None label issue --- .../datasets/molecule_net_classification.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py index 3beef041..0eb0ac0b 100644 --- a/chebai/preprocessing/datasets/molecule_net_classification.py +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -85,7 +85,7 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No # That transformer reweights the observed labels per task so positive and negative examples end up with equal total weight — it upweights the rarer class. # Currently, transformers are set to `[]` in the deepchem data loader API, so we can get the raw 0/1 weight matrix. labels = [ - bool(label) if weight != 0 else None + bool(label) if int(weight) != 0 else None for label, weight in zip(labels, w) ] yield dict(features=mol, labels=labels, ident=idx) @@ -355,6 +355,20 @@ def _name(self) -> str: if __name__ == "__main__": # Example usage - dataset = SIDER() - dataset.prepare_data() - dataset.setup() + for dataset_class in [ + ClinTox, + BBBP, + SIDER, + BACE, + Tox21, + ToxCast, + # PCBA, + # HIV, + # MUV, + ]: + dataset = dataset_class() + dataset.prepare_data() + dataset.setup() + # dataset = SIDER() + # dataset.prepare_data() + # dataset.setup() From bac2cd3c345def3599a408b25a77b43fc7e43a5f Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 15 Aug 2026 21:57:48 +0200 Subject: [PATCH 45/56] move missing label mask to collator func and handle cases where all labels are None [None, None] --- chebai/preprocessing/collate.py | 46 +++++++----- chebai/preprocessing/reader.py | 5 -- tests/unit/collators/testRaggedCollator.py | 84 ++++++++++++++++++---- tests/unit/mock_data/tox_mock_data.py | 8 --- 4 files changed, 101 insertions(+), 42 deletions(-) diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index 06b4ff34..93d146a6 100644 --- a/chebai/preprocessing/collate.py +++ b/chebai/preprocessing/collate.py @@ -61,10 +61,13 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: """ Collate ragged data samples (i.e., samples of unequal size, such as molecular sequences) into a batch. - Handles both fully and partially labeled data, where some samples may have `None` as their label. The indices - of non-null labels are stored in the `non_null_labels` field, which is used to filter out predictions for - unlabeled data during evaluation (e.g., F1, MSE). For models supporting partially labeled data, this method - ensures alignment between features and labels. Missing labels are passed as a loss keyword. + Handles both fully and partially labeled data by use of the following fields in the returned XYData: + + `non_null_labels`: Stores batch row indices of samples where the whole `labels` field is not None, like [0, 2]. + - Example: [[True, False], None, [False, None]] would result in `non_null_labels` = [0, 2]. + - This is used to filter out predictions for unlabeled samples during evaluation. + + `missing_labels`: Stores a per-sample, per-label-position boolean mask for unknown entries inside a label row, like the None in [1, None, 0]. Args: data (List[Union[Dict, Tuple]]): List of ragged data samples. Each sample can be a dictionary or tuple @@ -81,30 +84,41 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: if isinstance(data[0], tuple): # For legacy data x, y, idents = zip(*data) - missing_labels = None else: x, y, idents = zip( *((d["features"], d["labels"], d.get("ident")) for d in data) ) - missing_labels = [ - d.get( - "missing_labels", - [False for _ in y[0]] if y[0] is not None else [False], - ) - for d in data - ] + # Compute the per-sample, per-label-position boolean mask for unknown entries + # (e.g., the None in [1, None, 0]) on the *original* labels, before any + # filtering/padding is applied to `y`. Rows whose entire label is None are + # represented as all-False rows of the maximum label length. + if any(labels is not None for labels in y): + max_label_len = max(len(labels) for labels in y if labels is not None) + missing_labels = pad_sequence( + [ + torch.tensor([label is None for label in labels]) + if labels is not None + else torch.zeros(max_label_len, dtype=torch.bool) + for labels in y + ], + batch_first=True, + ) + else: + missing_labels = torch.tensor([]) + + # Typical y: ([True, False], None, [True, None], [True]) if any(x is not None for x in y): - # If any label is not None: (None, None, `1`, None) + # If any label is not None: (None, None, `[True, None]`, None) if any(x is None for x in y): - # If any label is None: (`None`, `None`, 1, `None`) + # If any label is None: (`None`, [True, False], [True], [False]) non_null_labels = [i for i, r in enumerate(y) if r is not None] y = self.process_label_rows( tuple(ye for i, ye in enumerate(y) if i in non_null_labels) ) loss_kwargs["non_null_labels"] = non_null_labels else: - # If all labels are not None: (`0`, `2`, `1`, `3`) + # If all labels are not None: (`[True, False]`, `[False, True, True]`, `[False]`, `[True]`) y = self.process_label_rows(y) else: @@ -112,7 +126,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: y = None loss_kwargs["non_null_labels"] = [] - loss_kwargs["missing_labels"] = torch.tensor(missing_labels) + loss_kwargs["missing_labels"] = missing_labels # Calculate the lengths of each sequence, create a binary mask for valid (non-padded) positions lens = torch.tensor(list(map(len, x))) model_kwargs["mask"] = torch.arange(max(lens))[None, :] < lens[:, None] diff --git a/chebai/preprocessing/reader.py b/chebai/preprocessing/reader.py index 664a8d8f..19e492e6 100644 --- a/chebai/preprocessing/reader.py +++ b/chebai/preprocessing/reader.py @@ -99,11 +99,6 @@ def _read_components(self, row: Dict[str, Any]) -> Dict[str, Any]: under the additional `missing_labels` keyword.""" labels = self._get_raw_label(row) additional_kwargs = self._get_additional_kwargs(row) - if labels is not None: - if any(label is None for label in labels): - additional_kwargs["missing_labels"] = [ - label is None for label in labels - ] return dict( features=self._get_raw_data(row), labels=labels, diff --git a/tests/unit/collators/testRaggedCollator.py b/tests/unit/collators/testRaggedCollator.py index d9ab2b1d..3c68825b 100644 --- a/tests/unit/collators/testRaggedCollator.py +++ b/tests/unit/collators/testRaggedCollator.py @@ -73,20 +73,50 @@ def test_call_with_missing_entire_labels(self) -> None: data: List[Dict] = [ {"features": [1, 2], "labels": [True, False], "ident": "sample1"}, {"features": [3, 4, 5], "labels": None, "ident": "sample2"}, - {"features": [6], "labels": [True], "ident": "sample3"}, + {"features": [7], "labels": [True, None], "ident": "sample3"}, + {"features": [6], "labels": [True], "ident": "sample4"}, + {"features": [8, 9], "labels": [None, None], "ident": "sample5"}, ] result: XYData = self.collator(data) # https://github.com/ChEB-AI/python-chebai/pull/48#issuecomment-2324393829 - expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]]) + expected_x = torch.tensor( + [ + [1, 2, 0], + [3, 4, 5], + [7, 0, 0], + [6, 0, 0], + [8, 9, 0], + ] + ) expected_y = torch.tensor( - [[True, False], [True, False]] + [ + [True, False], + [True, False], + [True, False], + [False, False], + ] ) # True -> 1, False -> 0 expected_mask_for_x = torch.tensor( - [[True, True, False], [True, True, True], [True, False, False]] + [ + [True, True, False], + [True, True, True], + [True, False, False], + [True, False, False], + [True, True, False], + ] + ) + expected_lens_for_x = torch.tensor([2, 3, 1, 1, 2]) + expected_missing_labels = torch.tensor( + [ + [False, False], # sample1 has no missing labels + [False, False], # sample2 has no missing labels (entire label is None) + [False, True], # sample3 has a missing label at index 1 + [False, False], # sample4 has no missing labels + [True, True], # sample5 has missing labels at both indices + ] ) - expected_lens_for_x = torch.tensor([2, 3, 1]) self.assertTrue( torch.equal(result.x, expected_x), @@ -110,19 +140,26 @@ def test_call_with_missing_entire_labels(self) -> None: ) self.assertEqual( result.additional_fields["loss_kwargs"]["non_null_labels"], - [0, 2], + [0, 2, 3, 4], "The non-null labels list does not match the expected output.", ) self.assertEqual( len(result.additional_fields["loss_kwargs"]["non_null_labels"]), - result.y.shape[1], + result.y.shape[0], "The length of non null labels list must match with target label variable size", ) self.assertEqual( result.additional_fields["idents"], - ("sample1", "sample2", "sample3"), + ("sample1", "sample2", "sample3", "sample4", "sample5"), "The identifiers do not match the expected output when labels are missing.", ) + self.assertTrue( + torch.equal( + result.additional_fields["loss_kwargs"]["missing_labels"], + expected_missing_labels, + ), + "The missing labels tensor does not match the expected output when labels are missing.", + ) def test_call_with_none_in_labels(self) -> None: """ @@ -132,18 +169,32 @@ def test_call_with_none_in_labels(self) -> None: {"features": [1, 2], "labels": [None, True], "ident": "sample1"}, {"features": [3, 4, 5], "labels": [True, False], "ident": "sample2"}, {"features": [6], "labels": [True], "ident": "sample3"}, + {"features": [7, 8], "labels": [None, None], "ident": "sample4"}, ] result: XYData = self.collator(data) - expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]]) + expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0], [7, 8, 0]]) expected_y = torch.tensor( - [[False, True], [True, False], [True, False]] + [[False, True], [True, False], [True, False], [False, False]] ) # None -> False expected_mask_for_x = torch.tensor( - [[True, True, False], [True, True, True], [True, False, False]] + [ + [True, True, False], + [True, True, True], + [True, False, False], + [True, True, False], + ] + ) + expected_lens_for_x = torch.tensor([2, 3, 1, 2]) + expected_missing_labels = torch.tensor( + [ + [True, False], # sample1 has a missing label at index 0 + [False, False], # sample2 has no missing labels + [False, False], # sample3 has no missing labels + [True, True], # sample4 has missing labels at both indices + ] ) - expected_lens_for_x = torch.tensor([2, 3, 1]) self.assertTrue( torch.equal(result.x, expected_x), @@ -167,9 +218,16 @@ def test_call_with_none_in_labels(self) -> None: ) self.assertEqual( result.additional_fields["idents"], - ("sample1", "sample2", "sample3"), + ("sample1", "sample2", "sample3", "sample4"), "The identifiers do not match the expected output when labels contain None.", ) + self.assertTrue( + torch.equal( + result.additional_fields["loss_kwargs"]["missing_labels"], + expected_missing_labels, + ), + "The missing labels tensor does not match the expected output when labels contain None.", + ) def test_call_with_empty_data(self) -> None: """ diff --git a/tests/unit/mock_data/tox_mock_data.py b/tests/unit/mock_data/tox_mock_data.py index fcf5633f..7567d6b2 100644 --- a/tests/unit/mock_data/tox_mock_data.py +++ b/tests/unit/mock_data/tox_mock_data.py @@ -394,10 +394,6 @@ def data_in_dict_format() -> List[Dict]: for dict_ in data_list: dict_["features"] = Tox21ChallengeMockData.FEATURE_OF_SMILES dict_["group"] = None - if any(label is None for label in dict_["labels"]): - dict_["missing_labels"] = [ - True if label is None else False for label in dict_["labels"] - ] return data_list @@ -509,9 +505,5 @@ def get_setup_processed_output_data() -> List[Dict]: "group": None, } ) - if any(label is None for label in dict_["labels"]): - complete_list[-1]["missing_labels"] = [ - True if label is None else False for label in dict_["labels"] - ] return complete_list From 112fc7ce82fd23a8d7cdc4c83db1e7859c43cb2f Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 18 Aug 2026 14:21:18 +0200 Subject: [PATCH 46/56] add script to compute avg score of best model across different wandb runs --- chebai/result/compute_avg_performance.py | 259 +++++++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 chebai/result/compute_avg_performance.py diff --git a/chebai/result/compute_avg_performance.py b/chebai/result/compute_avg_performance.py new file mode 100644 index 00000000..78fb11dd --- /dev/null +++ b/chebai/result/compute_avg_performance.py @@ -0,0 +1,259 @@ +""" +For each of one or more local W&B run files (run-*.wandb), find the step +with the best validation macro-F1 score, print that score along with the +corresponding validation micro-F1 (and any other metrics logged at that +same step). At the end, print the average of the best macro-F1 (and its +corresponding metrics) across all the given files. + +Usage: + python find_best_f1.py run1.wandb run2.wandb run3.wandb + python find_best_f1.py *.wandb --macro-metric val/macro_f1 --micro-metric val/micro_f1 + python find_best_f1.py *.wandb --max-epoch 200 + +If --macro-metric / --micro-metric aren't given, the script auto-detects +them per file (case-insensitive match on "f1"+"macro" / "f1"+"micro", +preferring keys that also mention "val"/"eval"/"test"). +""" + +import argparse +import json +import sys +from pathlib import Path +from statistics import mean + +try: + # Newer wandb versions (>=0.16 or so) + from wandb.sdk.internal.datastore import DataStore +except ImportError: + # Older wandb versions + from wandb.old.datastore import DataStore + +from wandb.proto import wandb_internal_pb2 as pb + +INTERNAL_KEY_PREFIXES = ("_",) # e.g. _step, _timestamp, _runtime + + +def iter_history_rows(wandb_path: str): + """ + Yields dicts of {key: value} for every 'history' record logged in the run, + parsed straight out of the binary .wandb file (no network / API calls). + """ + ds = DataStore() + ds.open_for_scan(wandb_path) + + while True: + data = ds.scan_data() + if data is None: + break + record = pb.Record() + record.ParseFromString(data) + + if record.WhichOneof("record_type") == "history": + row = {} + for item in record.history.item: + key = item.key if item.key else ".".join(item.nested_key) + try: + val = json.loads(item.value_json) + except Exception: + val = item.value_json + row[key] = val + yield row + + +def detect_metric_key(all_keys, must_contain, explicit=None): + """Find a logged key matching all substrings in must_contain (case-insensitive), + preferring ones that also look like validation metrics.""" + if explicit: + return explicit + + candidates = [k for k in all_keys if all(s in k.lower() for s in must_contain)] + val_candidates = [ + k + for k in candidates + if any(tag in k.lower() for tag in ("val", "eval", "test")) + ] + chosen = val_candidates if val_candidates else candidates + return chosen[0] if chosen else None + + +def get_epoch(row, epoch_key=None): + if epoch_key: + return row.get(epoch_key) + if "epoch" in row: + return row["epoch"] + return row.get("_step") + + +def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): + rows = list(iter_history_rows(str(path))) + if not rows: + print(f" [!] No history records found in {path}, skipping.") + return None + + all_keys = set() + for row in rows: + all_keys.update(row.keys()) + + macro_key = detect_metric_key(all_keys, ("f1", "macro"), macro_metric) + if macro_key is None: + print(f" [!] Could not find a macro-F1 metric in {path}. Available keys:") + print(" " + ", ".join(sorted(all_keys))) + return None + + micro_key = detect_metric_key(all_keys, ("f1", "micro"), micro_metric) + + best_row = None + best_epoch = None + best_val = None + for row in rows: + if macro_key not in row or row[macro_key] is None: + continue + epoch = get_epoch(row, epoch_key) + if epoch is not None and max_epoch is not None and epoch > max_epoch: + continue + try: + val = float(row[macro_key]) + except (TypeError, ValueError): + continue + if best_val is None or val > best_val: + best_val = val + best_row = row + best_epoch = epoch + + if best_row is None: + print( + f" [!] No numeric values for '{macro_key}' within epoch <= {max_epoch} in {path}." + ) + return None + + return { + "file": str(path), + "macro_key": macro_key, + "micro_key": micro_key, + "epoch": best_epoch, + "row": best_row, + } + + +def print_result(result): + row = result["row"] + macro_key = result["macro_key"] + micro_key = result["micro_key"] + + print(f"File: {result['file']}") + print(f" Best epoch/step: {result['epoch']}") + print(f" {macro_key}: {row[macro_key]:.4f}") + + if micro_key and micro_key in row and row[micro_key] is not None: + print(f" {micro_key} (corresponding): {row[micro_key]:.4f}") + elif micro_key: + print(f" {micro_key} (corresponding): N/A") + else: + print(" (no micro-F1 metric found)") + + shown = {macro_key, micro_key, "epoch", "_step"} + other_keys = sorted( + k + for k in row.keys() + if k not in shown and not k.startswith(INTERNAL_KEY_PREFIXES) + ) + if other_keys: + print(" Other metrics at this step:") + for k in other_keys: + v = row[k] + if isinstance(v, float): + print(f" {k}: {v:.4f}") + else: + print(f" {k}: {v}") + print() + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "wandb_files", nargs="+", help="Paths to one or more run-*.wandb files" + ) + parser.add_argument( + "--macro-metric", default=None, help="Exact key for macro-F1 (skip auto-detect)" + ) + parser.add_argument( + "--micro-metric", default=None, help="Exact key for micro-F1 (skip auto-detect)" + ) + parser.add_argument( + "--epoch-key", + default=None, + help="Key to use as the epoch/step number (default: auto epoch/_step)", + ) + parser.add_argument( + "--max-epoch", + type=int, + default=200, + help="Only consider steps/epochs up to this value", + ) + args = parser.parse_args() + + results = [] + for f in args.wandb_files: + path = Path(f) + if not path.exists(): + print(f" [!] File not found: {path}, skipping.") + continue + print(f"Processing {path.name} ...") + result = process_file( + path, args.macro_metric, args.micro_metric, args.epoch_key, args.max_epoch + ) + if result: + print_result(result) + results.append(result) + + if not results: + sys.exit("No valid results across the given files.") + + if len(results) < len(args.wandb_files): + print( + f"({len(results)}/{len(args.wandb_files)} files produced a valid result)\n" + ) + + print("=" * 50) + print(f"Average across {len(results)} file(s):") + + # Average the macro-F1 across files + macro_vals = [r["row"][r["macro_key"]] for r in results] + print(f" Average best macro-F1: {mean(macro_vals):.4f} (n={len(macro_vals)})") + + # Average the corresponding micro-F1 across files (where present) + micro_vals = [ + r["row"][r["micro_key"]] + for r in results + if r["micro_key"] + and r["micro_key"] in r["row"] + and r["row"][r["micro_key"]] is not None + ] + if micro_vals: + print( + f" Average corresponding micro-F1: {mean(micro_vals):.4f} (n={len(micro_vals)}/{len(results)})" + ) + + # Average every other numeric key found in the best rows (union across files) + shown = ( + {"epoch", "_step"} + | {r["macro_key"] for r in results} + | {r["micro_key"] for r in results if r["micro_key"]} + ) + other_key_values = {} + for r in results: + for k, v in r["row"].items(): + if k in shown or k.startswith(INTERNAL_KEY_PREFIXES): + continue + if isinstance(v, (int, float)): + other_key_values.setdefault(k, []).append(v) + + for k in sorted(other_key_values): + vals = other_key_values[k] + print(f" Average {k}: {mean(vals):.4f} (n={len(vals)}/{len(results)})") + + +if __name__ == "__main__": + main() From 892eeb435adfef55e5afb175cb23b358ee497968 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 18 Aug 2026 14:23:34 +0200 Subject: [PATCH 47/56] compute sample std deviation too --- chebai/result/compute_avg_performance.py | 28 +++++++++++++++++------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/chebai/result/compute_avg_performance.py b/chebai/result/compute_avg_performance.py index 78fb11dd..fb756537 100644 --- a/chebai/result/compute_avg_performance.py +++ b/chebai/result/compute_avg_performance.py @@ -2,8 +2,9 @@ For each of one or more local W&B run files (run-*.wandb), find the step with the best validation macro-F1 score, print that score along with the corresponding validation micro-F1 (and any other metrics logged at that -same step). At the end, print the average of the best macro-F1 (and its -corresponding metrics) across all the given files. +same step). At the end, print the average +/- sample standard deviation +(ddof=1, the standard convention for reporting results across seeds) of +the best macro-F1 and its corresponding metrics across all the given files. Usage: python find_best_f1.py run1.wandb run2.wandb run3.wandb @@ -19,7 +20,7 @@ import json import sys from pathlib import Path -from statistics import mean +from statistics import mean, stdev try: # Newer wandb versions (>=0.16 or so) @@ -135,6 +136,17 @@ def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): } +def format_mean_std(vals): + """Mean +/- sample standard deviation (ddof=1), the convention used in + research for reporting performance across seeds/runs. Falls back to + 'no tolerance' when only one value is available (stdev is undefined).""" + m = mean(vals) + if len(vals) > 1: + s = stdev(vals) # sample std (n-1 denominator) + return f"{m:.4f} \u00b1 {s:.4f}" + return f"{m:.4f} (n=1, no std)" + + def print_result(result): row = result["row"] macro_key = result["macro_key"] @@ -219,11 +231,11 @@ def main(): print("=" * 50) print(f"Average across {len(results)} file(s):") - # Average the macro-F1 across files + # Average (+/- sample std) the macro-F1 across files macro_vals = [r["row"][r["macro_key"]] for r in results] - print(f" Average best macro-F1: {mean(macro_vals):.4f} (n={len(macro_vals)})") + print(f" Best macro-F1: {format_mean_std(macro_vals)} (n={len(macro_vals)})") - # Average the corresponding micro-F1 across files (where present) + # Average (+/- sample std) the corresponding micro-F1 across files (where present) micro_vals = [ r["row"][r["micro_key"]] for r in results @@ -233,7 +245,7 @@ def main(): ] if micro_vals: print( - f" Average corresponding micro-F1: {mean(micro_vals):.4f} (n={len(micro_vals)}/{len(results)})" + f" Corresponding micro-F1: {format_mean_std(micro_vals)} (n={len(micro_vals)}/{len(results)})" ) # Average every other numeric key found in the best rows (union across files) @@ -252,7 +264,7 @@ def main(): for k in sorted(other_key_values): vals = other_key_values[k] - print(f" Average {k}: {mean(vals):.4f} (n={len(vals)}/{len(results)})") + print(f" {k}: {format_mean_std(vals)} (n={len(vals)}/{len(results)})") if __name__ == "__main__": From df704a1ebb633bb6a9fe41e3ff412482cf9e8d7b Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 18 Aug 2026 14:41:45 +0200 Subject: [PATCH 48/56] use wandb identifiers instead of filepaths --- chebai/result/compute_avg_performance.py | 30 +++++++++++++++++------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/chebai/result/compute_avg_performance.py b/chebai/result/compute_avg_performance.py index fb756537..bf115fcf 100644 --- a/chebai/result/compute_avg_performance.py +++ b/chebai/result/compute_avg_performance.py @@ -7,7 +7,7 @@ the best macro-F1 and its corresponding metrics across all the given files. Usage: - python find_best_f1.py run1.wandb run2.wandb run3.wandb + python find_best_f1.py 2cb51q4o 0nwo7wrt s4w2w2cx python find_best_f1.py *.wandb --macro-metric val/macro_f1 --micro-metric val/micro_f1 python find_best_f1.py *.wandb --max-epoch 200 @@ -185,7 +185,9 @@ def main(): description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument( - "wandb_files", nargs="+", help="Paths to one or more run-*.wandb files" + "wandb_run_ids", + nargs="+", + help="Identifiers of local W&B run files (run-*.wandb) to process", ) parser.add_argument( "--macro-metric", default=None, help="Exact key for macro-F1 (skip auto-detect)" @@ -207,11 +209,23 @@ def main(): args = parser.parse_args() results = [] - for f in args.wandb_files: - path = Path(f) + for wandb_id in args.wandb_run_ids: + file_name = f"run-{wandb_id}.wandb" + matches = list(Path(".").rglob(file_name)) + + if len(matches) == 0: + raise FileNotFoundError(f"Could not find {file_name}") + + if len(matches) > 1: + raise RuntimeError( + f"Found multiple files named {file_name}:\n" + + "\n".join(str(p.resolve()) for p in matches) + ) + file_path = matches[0].resolve() + path = Path(file_path) if not path.exists(): - print(f" [!] File not found: {path}, skipping.") - continue + raise FileNotFoundError(f" [!] File not found: {path}, skipping.") + print(f"Processing {path.name} ...") result = process_file( path, args.macro_metric, args.micro_metric, args.epoch_key, args.max_epoch @@ -223,9 +237,9 @@ def main(): if not results: sys.exit("No valid results across the given files.") - if len(results) < len(args.wandb_files): + if len(results) < len(args.wandb_run_ids): print( - f"({len(results)}/{len(args.wandb_files)} files produced a valid result)\n" + f"({len(results)}/{len(args.wandb_run_ids)} files produced a valid result)\n" ) print("=" * 50) From 2750293a30f5a533b7ae49ce1fd7c1fdb496b593 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 18 Aug 2026 14:48:30 +0200 Subject: [PATCH 49/56] raise error for any discrepancy --- chebai/result/compute_avg_performance.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/chebai/result/compute_avg_performance.py b/chebai/result/compute_avg_performance.py index bf115fcf..2fb63b5a 100644 --- a/chebai/result/compute_avg_performance.py +++ b/chebai/result/compute_avg_performance.py @@ -88,8 +88,7 @@ def get_epoch(row, epoch_key=None): def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): rows = list(iter_history_rows(str(path))) if not rows: - print(f" [!] No history records found in {path}, skipping.") - return None + raise ValueError(f" [!] No history records found in {path}, skipping.") all_keys = set() for row in rows: @@ -97,9 +96,9 @@ def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): macro_key = detect_metric_key(all_keys, ("f1", "macro"), macro_metric) if macro_key is None: - print(f" [!] Could not find a macro-F1 metric in {path}. Available keys:") - print(" " + ", ".join(sorted(all_keys))) - return None + raise ValueError( + f" [!] Could not find a macro-F1 metric in {path}. Available keys: {', '.join(sorted(all_keys))}" + ) micro_key = detect_metric_key(all_keys, ("f1", "micro"), micro_metric) @@ -122,10 +121,9 @@ def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): best_epoch = epoch if best_row is None: - print( + raise ValueError( f" [!] No numeric values for '{macro_key}' within epoch <= {max_epoch} in {path}." ) - return None return { "file": str(path), @@ -161,7 +159,7 @@ def print_result(result): elif micro_key: print(f" {micro_key} (corresponding): N/A") else: - print(" (no micro-F1 metric found)") + raise ValueError(f" [!] No micro-F1 metric found in {result['file']}.") shown = {macro_key, micro_key, "epoch", "_step"} other_keys = sorted( @@ -224,7 +222,7 @@ def main(): file_path = matches[0].resolve() path = Path(file_path) if not path.exists(): - raise FileNotFoundError(f" [!] File not found: {path}, skipping.") + raise FileNotFoundError(f" [!] File not found: {path}.") print(f"Processing {path.name} ...") result = process_file( @@ -237,10 +235,14 @@ def main(): if not results: sys.exit("No valid results across the given files.") - if len(results) < len(args.wandb_run_ids): + if len(results) == len(args.wandb_run_ids): print( f"({len(results)}/{len(args.wandb_run_ids)} files produced a valid result)\n" ) + else: + raise ValueError( + f"({len(results)} and {len(args.wandb_run_ids)} do not match)\n" + ) print("=" * 50) print(f"Average across {len(results)} file(s):") From 1bed1b4dc4a0be5571959749a1646ab116f956b7 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Wed, 19 Aug 2026 16:23:58 +0200 Subject: [PATCH 50/56] config to save best 3 based on macro, model from last epoch --- configs/training/default_callbacks.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/configs/training/default_callbacks.yml b/configs/training/default_callbacks.yml index 7635b88c..a20b9315 100644 --- a/configs/training/default_callbacks.yml +++ b/configs/training/default_callbacks.yml @@ -1,12 +1,11 @@ - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: - monitor: val_micro-f1 + monitor: val_macro-f1 mode: 'max' - filename: 'best_micro_f1_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' + filename: 'best_macro_f1_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' every_n_epochs: 1 save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: filename: 'per_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' - every_n_epochs: 25 - save_top_k: -1 + save_top_k: 1 From 154f7a990f8a1b3a2dc88e04d701dfabb54d9bc7 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 19 Aug 2026 16:50:17 +0200 Subject: [PATCH 51/56] binary callbacks save top 3 models --- configs/training/binary_callbacks.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/configs/training/binary_callbacks.yml b/configs/training/binary_callbacks.yml index 5f6369ac..9bad0262 100644 --- a/configs/training/binary_callbacks.yml +++ b/configs/training/binary_callbacks.yml @@ -4,16 +4,15 @@ mode: 'max' filename: 'best_f1_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' every_n_epochs: 1 - save_top_k: 1 + save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: monitor: val_roc-auc mode: 'max' filename: 'best_roc-auc_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' every_n_epochs: 1 - save_top_k: 1 + save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: filename: 'per_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' - every_n_epochs: 25 - save_top_k: -1 + save_top_k: 1 From 7729262c0766990557d072091fa610200ec83d61 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 20 Aug 2026 13:09:41 +0200 Subject: [PATCH 52/56] add base loss class to handle missing labels --- chebai/loss/base.py | 29 +++++++++++++++++++++++++++++ chebai/loss/bce_weighted.py | 10 +++++----- chebai/models/base.py | 3 +++ 3 files changed, 37 insertions(+), 5 deletions(-) create mode 100644 chebai/loss/base.py diff --git a/chebai/loss/base.py b/chebai/loss/base.py new file mode 100644 index 00000000..94834eec --- /dev/null +++ b/chebai/loss/base.py @@ -0,0 +1,29 @@ +from typing import Optional + +import torch + + +class BCELogitLossWithValidLabels(torch.nn.BCEWithLogitsLoss): + def __init__(self, **kwargs): + kwargs["reduction"] = "none" + super().__init__(**kwargs) + + def forward( + self, + input: torch.Tensor, + target: torch.Tensor, + valid_label_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + loss_mat = super().forward(input, target) + + if valid_label_mask is None: + return loss_mat.mean() + + loss_mat = torch.where( + valid_label_mask, + loss_mat, + torch.zeros_like(loss_mat), + ) + + return loss_mat.sum() / valid_label_mask.sum().clamp_min(1) diff --git a/chebai/loss/bce_weighted.py b/chebai/loss/bce_weighted.py index 440495cf..7453321b 100644 --- a/chebai/loss/bce_weighted.py +++ b/chebai/loss/bce_weighted.py @@ -3,11 +3,12 @@ import torch +from chebai.loss.base import BCELogitLossWithValidLabels from chebai.preprocessing.datasets.base import XYBaseDataModule from chebai.preprocessing.datasets.chebi import _ChEBIDataExtractor -class BCEWeighted(torch.nn.BCEWithLogitsLoss): +class BCEWeighted(BCELogitLossWithValidLabels): """ BCEWithLogitsLoss with weights automatically computed according to the beta parameter. @@ -104,10 +105,9 @@ def forward( torch.Tensor: The computed loss. """ self.set_pos_weight(input) - return super().forward(input, target) + return super().forward(input, target, **kwargs) -class UnWeightedBCEWithLogitsLoss(torch.nn.BCEWithLogitsLoss): +class UnWeightedBCEWithLogitsLoss(BCELogitLossWithValidLabels): def forward(self, input, target, **kwargs): - # As the custom passed kwargs are not used in BCEWithLogitsLoss, we can ignore them - return super().forward(input, target) + return super().forward(input, target, **kwargs) diff --git a/chebai/models/base.py b/chebai/models/base.py index 58d387a9..48412b16 100644 --- a/chebai/models/base.py +++ b/chebai/models/base.py @@ -295,6 +295,9 @@ def _execute( model_output, labels, data.get("loss_kwargs", dict()) ) loss_kwargs = dict() + loss_kwargs["valid_label_mask"] = loss_kwargs_candidates[ + "valid_label_mask" + ] if self.pass_loss_kwargs: loss_kwargs = loss_kwargs_candidates loss_kwargs["current_epoch"] = self.trainer.current_epoch From 8737cec5982e2a6fa15cbca64508ef2251e4f960 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 20 Aug 2026 13:37:40 +0200 Subject: [PATCH 53/56] use valid label mask instead of missing labels --- chebai/preprocessing/collate.py | 10 +++---- tests/unit/collators/testRaggedCollator.py | 32 +++++++++++----------- 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index 93d146a6..fcd2d61a 100644 --- a/chebai/preprocessing/collate.py +++ b/chebai/preprocessing/collate.py @@ -67,7 +67,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: - Example: [[True, False], None, [False, None]] would result in `non_null_labels` = [0, 2]. - This is used to filter out predictions for unlabeled samples during evaluation. - `missing_labels`: Stores a per-sample, per-label-position boolean mask for unknown entries inside a label row, like the None in [1, None, 0]. + `valid_label_mask`: Stores a per-sample, per-label-position boolean mask for valid entries inside a label row, like the True in [1, None, 0]. Args: data (List[Union[Dict, Tuple]]): List of ragged data samples. Each sample can be a dictionary or tuple @@ -95,9 +95,9 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: # represented as all-False rows of the maximum label length. if any(labels is not None for labels in y): max_label_len = max(len(labels) for labels in y if labels is not None) - missing_labels = pad_sequence( + valid_label_mask = pad_sequence( [ - torch.tensor([label is None for label in labels]) + torch.tensor([label is not None for label in labels]) if labels is not None else torch.zeros(max_label_len, dtype=torch.bool) for labels in y @@ -105,7 +105,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: batch_first=True, ) else: - missing_labels = torch.tensor([]) + valid_label_mask = None # Typical y: ([True, False], None, [True, None], [True]) if any(x is not None for x in y): @@ -126,7 +126,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: y = None loss_kwargs["non_null_labels"] = [] - loss_kwargs["missing_labels"] = missing_labels + loss_kwargs["valid_label_mask"] = valid_label_mask # Calculate the lengths of each sequence, create a binary mask for valid (non-padded) positions lens = torch.tensor(list(map(len, x))) model_kwargs["mask"] = torch.arange(max(lens))[None, :] < lens[:, None] diff --git a/tests/unit/collators/testRaggedCollator.py b/tests/unit/collators/testRaggedCollator.py index 3c68825b..aa0258cd 100644 --- a/tests/unit/collators/testRaggedCollator.py +++ b/tests/unit/collators/testRaggedCollator.py @@ -108,13 +108,13 @@ def test_call_with_missing_entire_labels(self) -> None: ] ) expected_lens_for_x = torch.tensor([2, 3, 1, 1, 2]) - expected_missing_labels = torch.tensor( + expected_valid_label_mask = torch.tensor( [ - [False, False], # sample1 has no missing labels + [True, True], # sample1 has no missing labels [False, False], # sample2 has no missing labels (entire label is None) - [False, True], # sample3 has a missing label at index 1 - [False, False], # sample4 has no missing labels - [True, True], # sample5 has missing labels at both indices + [True, False], # sample3 has a missing label at index 1 + [True, False], # sample4 has no missing labels + [False, False], # sample5 has missing labels at both indices ] ) @@ -155,10 +155,10 @@ def test_call_with_missing_entire_labels(self) -> None: ) self.assertTrue( torch.equal( - result.additional_fields["loss_kwargs"]["missing_labels"], - expected_missing_labels, + result.additional_fields["loss_kwargs"]["valid_label_mask"], + expected_valid_label_mask, ), - "The missing labels tensor does not match the expected output when labels are missing.", + "The valid label mask tensor does not match the expected output when labels are missing.", ) def test_call_with_none_in_labels(self) -> None: @@ -187,12 +187,12 @@ def test_call_with_none_in_labels(self) -> None: ] ) expected_lens_for_x = torch.tensor([2, 3, 1, 2]) - expected_missing_labels = torch.tensor( + expected_valid_label_mask = torch.tensor( [ - [True, False], # sample1 has a missing label at index 0 - [False, False], # sample2 has no missing labels - [False, False], # sample3 has no missing labels - [True, True], # sample4 has missing labels at both indices + [False, True], # sample1 has a missing label at index 0 + [True, True], # sample2 has no missing labels + [True, False], # sample3 has no missing labels + [False, False], # sample4 has missing labels at both indices ] ) @@ -223,10 +223,10 @@ def test_call_with_none_in_labels(self) -> None: ) self.assertTrue( torch.equal( - result.additional_fields["loss_kwargs"]["missing_labels"], - expected_missing_labels, + result.additional_fields["loss_kwargs"]["valid_label_mask"], + expected_valid_label_mask, ), - "The missing labels tensor does not match the expected output when labels contain None.", + "The valid label mask tensor does not match the expected output when labels contain None.", ) def test_call_with_empty_data(self) -> None: From 7d8a12f6644a82e5abdbef24355ba258a0728ef4 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 20 Aug 2026 15:28:03 +0200 Subject: [PATCH 54/56] valid mask method --- chebai/preprocessing/collate.py | 37 ++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index fcd2d61a..d7af511a 100644 --- a/chebai/preprocessing/collate.py +++ b/chebai/preprocessing/collate.py @@ -89,23 +89,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: *((d["features"], d["labels"], d.get("ident")) for d in data) ) - # Compute the per-sample, per-label-position boolean mask for unknown entries - # (e.g., the None in [1, None, 0]) on the *original* labels, before any - # filtering/padding is applied to `y`. Rows whose entire label is None are - # represented as all-False rows of the maximum label length. - if any(labels is not None for labels in y): - max_label_len = max(len(labels) for labels in y if labels is not None) - valid_label_mask = pad_sequence( - [ - torch.tensor([label is not None for label in labels]) - if labels is not None - else torch.zeros(max_label_len, dtype=torch.bool) - for labels in y - ], - batch_first=True, - ) - else: - valid_label_mask = None + valid_label_mask = self._get_valid_label_mask(y) # Typical y: ([True, False], None, [True, None], [True]) if any(x is not None for x in y): @@ -160,3 +144,22 @@ def process_label_rows(self, labels: Tuple) -> torch.Tensor: ], batch_first=True, ) + + def _get_valid_label_mask(self, y: Tuple) -> torch.Tensor | None: + # Compute the per-sample, per-label-position boolean mask for unknown entries + # (e.g., the None in [1, None, 0]) on the *original* labels, before any + # filtering/padding is applied to `y`. Rows whose entire label is None are + # represented as all-False rows of the maximum label length. + if any(labels is not None for labels in y): + max_label_len = max(len(labels) for labels in y if labels is not None) + return pad_sequence( + [ + torch.tensor([label is not None for label in labels]) + if labels is not None + else torch.zeros(max_label_len, dtype=torch.bool) + for labels in y + ], + batch_first=True, + ) + + return None From 084c0497ad908622bac3eb424b8db9cc358a22a4 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 21 Aug 2026 00:10:32 +0200 Subject: [PATCH 55/56] add himol specific marco auc-roc --- chebai/callbacks/epoch_metrics.py | 83 ++++++++++++++++++++++ chebai/cli.py | 1 + configs/metrics/micro-macro-f1-roc-auc.yml | 8 +++ 3 files changed, 92 insertions(+) diff --git a/chebai/callbacks/epoch_metrics.py b/chebai/callbacks/epoch_metrics.py index 76d6a8fd..3edb20b6 100644 --- a/chebai/callbacks/epoch_metrics.py +++ b/chebai/callbacks/epoch_metrics.py @@ -1,5 +1,8 @@ +import warnings + import torch import torchmetrics +from sklearn.metrics import roc_auc_score def custom_reduce_fx(input: torch.Tensor) -> torch.Tensor: @@ -179,3 +182,83 @@ def compute(self) -> torch.Tensor: balanced_acc = (tpr + tnr) / 2 return torch.mean(balanced_acc) + + +class HiMolMacroAUROC(torchmetrics.Metric): + """ + Macro-averaged multilabel AUROC that exactly replicates HiMol's eval(): + - missing labels (== ignore_index) are excluded per-task before scoring + - any task that, after exclusion, has only one class present is + dropped from BOTH the sum and the divisor of the macro average + (instead of being scored as 0.0 and diluting the mean, which is + torchmetrics' default MultilabelAUROC behavior) + + Assumes preds/target are shape (N, num_labels), target values in {0, 1}, + with `ignore_index` marking missing/unlabeled entries. + + References: + https://github.com/ZangXuan/HiMol/blob/ffdcb247b361a1f85ddb741862cff25e4a3b3341/finetune/optimization.py#L95-L104 + """ + + full_state_update = False + is_differentiable = False + higher_is_better = True + + def __init__(self, num_labels: int, ignore_index: int = 0, **kwargs): + super().__init__(**kwargs) + self.num_labels = num_labels + self.ignore_index = ignore_index + + self.add_state("preds", default=[], dist_reduce_fx="cat") + self.add_state("target", default=[], dist_reduce_fx="cat") + + def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: + if preds.shape != target.shape: + raise ValueError( + f"preds/target shape mismatch: {preds.shape} vs {target.shape}" + ) + if preds.ndim != 2 or preds.shape[1] != self.num_labels: + raise ValueError( + f"expected shape (N, {self.num_labels}), got {tuple(preds.shape)}" + ) + + self.preds.append(preds.detach().cpu()) + self.target.append(target.detach().cpu()) + + def compute(self) -> torch.Tensor: + preds = torch.cat(self.preds, dim=0).cpu().numpy() + target = torch.cat(self.target, dim=0).cpu().numpy() + + roc_list = [] + n_skipped = 0 + + for i in range(self.num_labels): + col_target = target[:, i] + col_preds = preds[:, i] + + valid = col_target != self.ignore_index + col_target = col_target[valid] + col_preds = col_preds[valid] + + # need at least one of each class to define AUC + if ( + len(col_target) == 0 + or (col_target == 0).sum() == 0 + or (col_target == 1).sum() == 0 + ): + n_skipped += 1 + continue + + roc_list.append(roc_auc_score(col_target, col_preds)) + + if n_skipped > 0: + warnings.warn( + f"{n_skipped}/{self.num_labels} labels skipped (missing or single-class " + f"after masking). Macro AUROC computed over {len(roc_list)} labels.", + stacklevel=2, + ) + + if len(roc_list) == 0: + return torch.tensor(float("nan")) + + return torch.tensor(sum(roc_list) / len(roc_list)) diff --git a/chebai/cli.py b/chebai/cli.py index fccf4de5..df355a98 100644 --- a/chebai/cli.py +++ b/chebai/cli.py @@ -76,6 +76,7 @@ def call_data_methods(data: Type[XYBaseDataModule]): "rmse", "r2", "pr-auc", + "himol-marco-roc-auc", ): # When using lightning > 2.5.1 then need to uncomment all metrics that are not used # for average in ("mse", "rmse","r2"): # for regression diff --git a/configs/metrics/micro-macro-f1-roc-auc.yml b/configs/metrics/micro-macro-f1-roc-auc.yml index 9e7b0d45..88bdcad3 100644 --- a/configs/metrics/micro-macro-f1-roc-auc.yml +++ b/configs/metrics/micro-macro-f1-roc-auc.yml @@ -9,5 +9,13 @@ init_args: class_path: chebai.callbacks.epoch_metrics.MacroF1 roc-auc: class_path: torchmetrics.classification.MultilabelAUROC + init_args: + ignore_index: -1 pr-auc: # Especially used for MUV, PCBA class_path: torchmetrics.classification.MultilabelAveragePrecision + init_args: + ignore_index: -1 + himol-marco-roc-auc: + class_path: chebai.callbacks.epoch_metrics.HiMolMacroAUROC + init_args: + ignore_index: -1 From f4fabae32523ada591b2bcb85e1f8eab9425a18b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 21 Aug 2026 13:23:48 +0200 Subject: [PATCH 56/56] also return None if valid_label_mask has all valid labels --- chebai/preprocessing/collate.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index d7af511a..3bfb3d70 100644 --- a/chebai/preprocessing/collate.py +++ b/chebai/preprocessing/collate.py @@ -152,7 +152,7 @@ def _get_valid_label_mask(self, y: Tuple) -> torch.Tensor | None: # represented as all-False rows of the maximum label length. if any(labels is not None for labels in y): max_label_len = max(len(labels) for labels in y if labels is not None) - return pad_sequence( + valid_label_mask = pad_sequence( [ torch.tensor([label is not None for label in labels]) if labels is not None @@ -161,5 +161,9 @@ def _get_valid_label_mask(self, y: Tuple) -> torch.Tensor | None: ], batch_first=True, ) + if (~valid_label_mask).sum() != 0: + # If there are any invalid labels, return the valid_label_mask + # Else, return None to indicate that all labels are valid (no None entries). + return valid_label_mask return None