Spaces:
Running
Running
File size: 3,314 Bytes
3286d5c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # pipeline taken from https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py
"""
This files includes a the data processing for Tox21.
As an input it takes a list of SMILES and it outputs a nested dictionary with
SMILES and target names as keys.
"""
import os
import json
import argparse
import numpy as np
from src.data import create_descriptors, get_tox21_split
from src.utils import TASKS, HF_TOKEN, write_pickle, create_dir, normalize_config
parser = argparse.ArgumentParser(
description="Data preprocessing script for the Tox21 dataset"
)
parser.add_argument(
"--config",
type=str,
default="config/config.json",
)
def main(config):
"""Preprocess the training and validation data for TabICLv2.
1. Download Tox21 train/val data from HF
2. Preprocess dataset splits
"""
ds = get_tox21_split(HF_TOKEN, cvfold=config["cvfold"])
feature_creation_kwargs = {
"radius": config["ecfp"]["radius"],
"fpsize": config["ecfp"]["fpsize"],
"min_var": config["feature_selection"]["min_var"],
"max_corr": config["feature_selection"]["max_corr"],
}
splits = ["train", "validation"]
for split in splits:
print(f"Preprocess {split} molecules")
ds_split = ds[split]
smiles = list(ds_split["smiles"])
if split == "train":
output = create_descriptors(
smiles,
return_feature_selection=True,
return_ecdfs=True,
**feature_creation_kwargs,
)
features = output.pop("features")
feature_selection = output.pop("feature_selection")
ecdfs = output.pop("ecdfs")
feature_selection_path = os.path.join(
config["data_folder"], "feat_selection.npz"
)
np.savez(
feature_selection_path,
ecfps_selec=feature_selection["ecfps_selec"],
tox_selec=feature_selection["tox_selec"],
)
print(f"Saved feature selection under {feature_selection_path}")
ecdfs_path = os.path.join(config["data_folder"], "ecdfs.pkl")
write_pickle(ecdfs_path, ecdfs)
print(f"Saved ECDFs under {ecdfs_path}")
else:
features = create_descriptors(
smiles,
ecdfs=ecdfs,
feature_selection=feature_selection,
**feature_creation_kwargs,
)["features"]
labels = []
for task in TASKS:
labels.append(ds_split[task].to_numpy())
labels = np.stack(labels, axis=1)
save_path = os.path.join(
config["data_folder"], f"tox21_{split}_cv{config['cvfold']}.npz"
)
with open(save_path, "wb") as f:
np.savez(
f,
labels=labels,
**features,
)
print(f"Saved preprocessed {split} split under {config['data_folder']}")
print("Preprocessing finished successfully")
if __name__ == "__main__":
args = parser.parse_args()
with open(args.config, "r") as f:
config = json.load(f)
config = normalize_config(config)
create_dir(config["data_folder"])
main(config)
|