Spaces:
Running
Running
| """Fit preprocessing assets and the endpoint-specific TabICLv2 estimators.""" | |
| import os | |
| import json | |
| import random | |
| import logging | |
| import argparse | |
| import numpy as np | |
| from sklearn.preprocessing import StandardScaler | |
| from datetime import datetime | |
| from src.model import Tox21TabICL | |
| from src.utils import TASKS, KNOWN_DESCR, create_dir, write_pickle, normalize_config | |
| parser = argparse.ArgumentParser(description="TabICLv2 training script for Tox21") | |
| parser.add_argument( | |
| "--config", | |
| type=str, | |
| default="config/config.json", | |
| ) | |
| parser.add_argument( | |
| "--seed", | |
| type=int, | |
| default=None, | |
| ) | |
| def main(config): | |
| print("Preprocess train molecules") | |
| timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") | |
| # setup logger | |
| logger = logging.getLogger(__name__) | |
| script_name = os.path.splitext(os.path.basename(__file__))[0] | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| handlers=[ | |
| logging.FileHandler( | |
| os.path.join( | |
| config["log_folder"], | |
| f"{script_name}_{timestamp}.log", | |
| ) | |
| ), | |
| logging.StreamHandler(), | |
| ], | |
| ) | |
| logger.info(config) | |
| # seeding | |
| random.seed(config["seed"]) | |
| np.random.seed(config["seed"]) | |
| train_data = np.load( | |
| os.path.join( | |
| config["data_folder"], f"tox21_train_cv{config['cvfold']}.npz" | |
| ) | |
| ) | |
| train_X = np.concatenate([train_data[descr] for descr in KNOWN_DESCR], axis=1) | |
| train_y = train_data["labels"] | |
| val_data = np.load( | |
| os.path.join( | |
| config["data_folder"], f"tox21_validation_cv{config['cvfold']}.npz" | |
| ) | |
| ) | |
| val_X = np.concatenate([val_data[descr] for descr in KNOWN_DESCR], axis=1) | |
| val_y = val_data["labels"] | |
| train_X = np.concatenate([train_X, val_X], axis=0) | |
| train_y = np.concatenate([train_y, val_y], axis=0) | |
| # remove molecules that couldn't be sanitized | |
| mask = ~np.isnan(train_X).any(axis=1) | |
| train_X = train_X[mask] | |
| train_y = train_y[mask] | |
| # select 500 features with most variance | |
| variances_ = np.nanvar(train_X, axis=0) | |
| order = np.argsort(variances_)[ | |
| : -(config["feature_selection"]["max_features"] + 1) : -1 | |
| ] | |
| logger.info( | |
| f"Save used feature order as: {os.path.join(config['data_folder'], 'feature_order.npy')}" | |
| ) | |
| np.save(os.path.join(config["data_folder"], "feature_order.npy"), order) | |
| train_X = train_X[:, order] | |
| max_samples = ( | |
| config["max_samples"] if config["max_samples"] > 0 else train_X.shape[0] | |
| ) | |
| resample_idxs = np.random.choice( | |
| np.arange(train_X.shape[0]), size=(max_samples,), replace=True | |
| ) | |
| train_X = train_X[resample_idxs] | |
| train_y = train_y[resample_idxs] | |
| # create and save scaler | |
| scaler = StandardScaler() | |
| train_X = scaler.fit_transform(train_X) | |
| logger.info( | |
| f"Save fitted scaler as: {os.path.join(config['data_folder'], 'scaler.pkl')}" | |
| ) | |
| write_pickle(os.path.join(config["data_folder"], "scaler.pkl"), scaler) | |
| if config["ckpt_path"]: | |
| logger.info( | |
| f"Fitted TabICLClassifier state will be saved in: {config['ckpt_path']}" | |
| ) | |
| else: | |
| logger.info("Fitted TabICLClassifier state will NOT be saved.") | |
| model = Tox21TabICL( | |
| seed=config["seed"], | |
| device=config["device"], | |
| model_config=config["model_config"], | |
| ) | |
| logger.info("Start training.") | |
| for i, task in enumerate(TASKS): | |
| logger.info(f"Fitting task: {task}") | |
| # select task-specific data | |
| task_train_y = train_y[:, i] | |
| task_mask = ~np.isnan(task_train_y) | |
| task_train_y = task_train_y[task_mask].astype(int) | |
| task_train_X = train_X[task_mask] | |
| print(f"Fit task {task} using {sum(task_mask)} samples") | |
| model.fit(task, task_train_X, task_train_y) | |
| log_text = f"Finished training." | |
| logger.info(log_text) | |
| if config["ckpt_path"]: | |
| model.save_model(config["ckpt_path"]) | |
| logger.info(f"Save model under: {config['ckpt_path']}") | |
| 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["log_folder"]) | |
| if args.seed is not None: | |
| config["seed"] = args.seed | |
| main(config) | |