Spaces:
Running
Running
File size: 3,209 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 | """
This files includes a predict function for the Tox21.
As an input it takes a list of SMILES and it outputs a nested dictionary with
SMILES and target names as keys.
"""
# ---------------------------------------------------------------------------------------
# Dependencies
import os
import json
from collections import defaultdict
import numpy as np
from src.model import Tox21TabICL
from src.data import create_descriptors
from src.utils import load_pickle, KNOWN_DESCR, normalize_config
# ---------------------------------------------------------------------------------------
from tqdm import tqdm
CONFIG_FILE = "config/config.json"
def predict(
smiles_list: list[str], default_prediction=0.5
) -> dict[str, dict[str, float]]:
"""Applies the classifier to a list of SMILES strings. Returns prediction=0.0 for
any molecule that could not be cleaned.
Args:
smiles_list (list[str]): list of SMILES strings
Returns:
dict: nested prediction dictionary, following {'<smiles>': {'<target>': <pred>}}
"""
print(f"Received {len(smiles_list)} SMILES strings")
with open(CONFIG_FILE, "r") as f:
config = json.load(f)
config = normalize_config(config)
# preprocessing pipeline
scaler_path = os.path.join(config["data_folder"], "scaler.pkl")
ecdfs_path = os.path.join(config["data_folder"], "ecdfs.pkl")
feature_selection_path = os.path.join(config["data_folder"], "feat_selection.npz")
feature_order_path = os.path.join(config["data_folder"], "feature_order.npy")
scaler = load_pickle(scaler_path)
ecdfs = load_pickle(ecdfs_path)
feature_selection = np.load(feature_selection_path)
feature_order = np.load(feature_order_path)
print(f"Loaded scaler from: {scaler_path}")
print(f"Loaded ecdfs from: {ecdfs_path}")
print(f"Loaded feature selection from: {feature_selection_path}")
print(f"Loaded feature order from: {feature_order_path}")
features = create_descriptors(
smiles_list,
ecdfs=ecdfs,
feature_selection=feature_selection,
radius=config["ecfp"]["radius"],
fpsize=config["ecfp"]["fpsize"],
)["features"]
features = np.concatenate([features[descr] for descr in KNOWN_DESCR], axis=1)
features = scaler.transform(features[:, feature_order])
print(f"Created descriptors for molecules.")
print(f"{(np.isnan(features).all(axis=1).sum())} molecules removed during cleaning")
is_clean = ~np.isnan(features).all(axis=1)
# setup model
model = Tox21TabICL(
seed=config["seed"],
device=config["device"],
model_config=config["model_config"],
)
model.load_model(config["ckpt_path"])
print(f"Loaded model from: {config['ckpt_path']}")
# make predicitons
predictions = defaultdict(dict)
print(f"Create predictions:")
preds = []
for target in tqdm(model.tasks):
preds = np.empty_like(is_clean, dtype=float)
preds[~is_clean] = default_prediction
preds[is_clean] = model.predict(target, features[is_clean])
for smiles, pred in zip(smiles_list, preds):
predictions[smiles][target] = float(pred)
return predictions
|