rk-random/PACT-Net

1import pandas as pd2from sklearn.model_selection import train_test_split3import importlib.resources as pkg_resources4import polyatomic_complexes5import numpy as np6from typing import Tuple7from pathlib import Path8 9 10def load_dataset(name) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:11    if name.lower() == "esol":12        data_path = (13            pkg_resources.files("polyatomic_complexes.dataset.esol") / "ESOL.csv"14        )15        df = pd.read_csv(str(data_path))16        target_col = "measured log solubility in mols per litre"17    elif name.lower() == "freesolv":18        data_path = (19            pkg_resources.files("polyatomic_complexes.dataset.free_solv")20            / "FreeSolv.csv"21        )22        df = pd.read_csv(str(data_path))23        target_col = "expt"24    elif name.lower() == "lipophil":25        data_path = (26            pkg_resources.files("polyatomic_complexes.dataset.lipophilicity")27            / "Lipophilicity.csv"28        )29        df = pd.read_csv(str(data_path))30        target_col = "exp"31    elif name.lower() == "boilingpoint":32        data_path = (33            Path(__file__).parent.parent / "benchmark_csv/boiling_point.csv".__str__()34        )35        df = pd.read_csv(data_path)36        target_col = "boiling_point_K"37    elif name.lower() == "qm9":38        data_path = (39            Path(__file__).parent.parent / "benchmark_csv/qm9_subset.csv".__str__()40        )41        df = pd.read_csv(data_path)42        target_col = "cv"43    elif name.lower() == "ic50":44        data_path = (45            Path(__file__).parent.parent / "benchmark_csv/ic_50_subset.csv".__str__()46        )47        df = pd.read_csv(data_path)48        target_col = "pIC50"49    elif name.lower() == "bindingdb":50        data_path = (51            Path(__file__).parent.parent / "benchmark_csv/bindingdb.csv".__str__()52        )53        df = pd.read_csv(data_path)54        target_col = "pIC50"55    else:56        raise ValueError(f"Unknown dataset: {name}")57 58    df.dropna(subset=["smiles", target_col], inplace=True)59    smiles = df["smiles"]60    targets = df[target_col]61    X_train, X_test, y_train, y_test = train_test_split(62        smiles, targets, test_size=0.2, random_state=4263    )64    return X_train.to_numpy(), X_test.to_numpy(), y_train.to_numpy(), y_test.to_numpy()65