|
|
| import numpy as np |
| import pandas as pd |
| import joblib |
| import argparse |
| import os |
| import sys |
| import torch |
| from rdkit import Chem |
| from rdkit.Chem import AllChem, MACCSkeys |
|
|
| BASE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, BASE) |
|
|
| from build_vocab import WordVocab |
| from pretrain_trfm import TrfmSeq2seq |
| from utils import split |
|
|
| cache_data = np.load(os.path.join(BASE, "prott5_cache.npz"), allow_pickle=True) |
| prott5_cache = cache_data["cache"].item() |
| scaler = joblib.load(os.path.join(BASE, "scaler.pkl")) |
| smiles_cache = joblib.load(os.path.join(BASE, "smiles_cache.pkl")) |
|
|
| N_FOLDS = 5 |
| et_kcat_models = [joblib.load(os.path.join(BASE, f"et_kcat_fold{i}.pkl")) for i in range(N_FOLDS)] |
| et_km_models = [joblib.load(os.path.join(BASE, f"et_km_fold{i}.pkl")) for i in range(N_FOLDS)] |
|
|
| vocab = WordVocab.load_vocab(os.path.join(BASE, "vocab.pkl")) |
| pad_index = 0; unk_index = 1; eos_index = 2; sos_index = 3 |
|
|
| trfm = TrfmSeq2seq(len(vocab), 256, len(vocab), 4) |
| trfm.load_state_dict(torch.load(os.path.join(BASE, "trfm_12_23000.pkl"), map_location="cpu")) |
| trfm.eval() |
|
|
| def smiles_to_vec(smiles): |
| x_split = [split(smiles)] |
| xid = [] |
| for tokens in x_split: |
| seq_len = 220 |
| if len(tokens) > 218: |
| tokens = tokens[:109] + tokens[-109:] |
| ids = [vocab.stoi.get(t, unk_index) for t in tokens] |
| ids = [sos_index] + ids + [eos_index] |
| pad_len = seq_len - len(ids) |
| ids += [pad_index] * pad_len |
| xid.append(ids) |
| xid = torch.tensor(xid).T |
| with torch.no_grad(): |
| out = trfm.encode(xid) |
| return np.array(out, dtype=np.float32).squeeze(0) |
|
|
| def get_maccs(smiles): |
| mol = Chem.MolFromSmiles(smiles) |
| if mol is None: return np.zeros(167, dtype=np.float32) |
| return np.array(MACCSkeys.GenMACCSKeys(mol), dtype=np.float32) |
|
|
| def get_morgan_fp(smiles, radius=2, nbits=2048): |
| mol = Chem.MolFromSmiles(smiles) |
| if mol is None: return np.zeros(nbits, dtype=np.float32) |
| fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=nbits) |
| return np.array(fp, dtype=np.float32) |
|
|
| def predict(input_path, output_path): |
| df = pd.read_parquet(input_path) |
| rows = [] |
| for _, row in df.iterrows(): |
| prot_emb = prott5_cache.get(row["sequence"], np.zeros(1024, dtype=np.float32)) |
| smiles_emb = smiles_cache.get(row["smiles"]) |
| if smiles_emb is None: |
| smiles_emb = smiles_to_vec(row["smiles"]) |
| maccs = get_maccs(row["smiles"]) |
| morgan = get_morgan_fp(row["smiles"]) |
| rows.append(np.concatenate([prot_emb, smiles_emb, maccs, morgan])) |
| X = np.array(rows, dtype=np.float32) |
| X = scaler.transform(X) |
| pred_kcat = np.mean([m.predict(X) for m in et_kcat_models], axis=0) |
| pred_km = np.mean([m.predict(X) for m in et_km_models], axis=0) |
| out = pd.DataFrame({ |
| "id" : df["id"].values, |
| "log10_kcat" : pred_kcat, |
| "log10_km" : pred_km, |
| }) |
| out.to_csv(output_path, index=False) |
| print(f"Saved {len(out)} predictions to {output_path}") |
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--input", required=True) |
| parser.add_argument("--output", required=True) |
| args = parser.parse_args() |
| predict(args.input, args.output) |
|
|
