UniKP / scripts /UniKP_pH_s2.py
anzhi2710gmailcom's picture
Upload folder using huggingface_hub
90d94ec verified
Raw
History Blame Contribute Delete
7.73 kB
import _bootstrap
import torch
from model.build_vocab import WordVocab
from model.pretrain_trfm import TrfmSeq2seq
from model.utils import split
import json
from transformers import T5EncoderModel, T5Tokenizer
import re
import gc
from sklearn import metrics
from sklearn.ensemble import ExtraTreesRegressor
from sklearn.linear_model import LinearRegression
import numpy as np
import pandas as pd
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
import random
import pickle
import math
from project_paths import (
DEGREE_SMILES_PATH, KCAT_DATASET_PATH, KCAT_KM_SAMPLES_PATH, KM_TEST_PATH,
PH_SMILES_PATH, PROT_T5_MODEL, TRFM_PATH, UNIKP_MODEL_DIR, VOCAB_PATH,
)
def smiles_to_vec(Smiles):
pad_index = 0
unk_index = 1
eos_index = 2
sos_index = 3
mask_index = 4
vocab = WordVocab.load_vocab(VOCAB_PATH)
def get_inputs(sm):
seq_len = 220
sm = sm.split()
if len(sm)>218:
print('SMILES is too long ({:d})'.format(len(sm)))
sm = sm[:109]+sm[-109:]
ids = [vocab.stoi.get(token, unk_index) for token in sm]
ids = [sos_index] + ids + [eos_index]
seg = [1]*len(ids)
padding = [pad_index]*(seq_len - len(ids))
ids.extend(padding), seg.extend(padding)
return ids, seg
def get_array(smiles):
x_id, x_seg = [], []
for sm in smiles:
a,b = get_inputs(sm)
x_id.append(a)
x_seg.append(b)
return torch.tensor(x_id), torch.tensor(x_seg)
trfm = TrfmSeq2seq(len(vocab), 256, len(vocab), 4)
trfm.load_state_dict(torch.load(TRFM_PATH))
trfm.eval()
x_split = [split(sm) for sm in Smiles]
xid, xseg = get_array(x_split)
X = trfm.encode(torch.t(xid))
return X
def Seq_to_vec(Sequence):
for i in range(len(Sequence)):
if len(Sequence[i]) > 1000:
Sequence[i] = Sequence[i][:500] + Sequence[i][-500:]
sequences_Example = []
for i in range(len(Sequence)):
zj = ''
for j in range(len(Sequence[i]) - 1):
zj += Sequence[i][j] + ' '
zj += Sequence[i][-1]
sequences_Example.append(zj)
tokenizer = T5Tokenizer.from_pretrained(PROT_T5_MODEL, do_lower_case=False)
model = T5EncoderModel.from_pretrained(PROT_T5_MODEL)
gc.collect()
print(torch.cuda.is_available())
# 'cuda:0' if torch.cuda.is_available() else
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model = model.eval()
features = []
for i in range(len(sequences_Example)):
print('For sequence ', str(i+1))
sequences_Example_i = sequences_Example[i]
sequences_Example_i = [re.sub(r"[UZOB]", "X", sequences_Example_i)]
ids = tokenizer.batch_encode_plus(sequences_Example_i, add_special_tokens=True, padding=True)
input_ids = torch.tensor(ids['input_ids']).to(device)
attention_mask = torch.tensor(ids['attention_mask']).to(device)
with torch.no_grad():
embedding = model(input_ids=input_ids, attention_mask=attention_mask)
embedding = embedding.last_hidden_state.cpu().numpy()
for seq_num in range(len(embedding)):
seq_len = (attention_mask[seq_num] == 1).sum()
seq_emd = embedding[seq_num][:seq_len - 1]
features.append(seq_emd)
features_normalize = np.zeros([len(features), len(features[0][0])], dtype=float)
for i in range(len(features)):
for k in range(len(features[0][0])):
for j in range(len(features[i])):
features_normalize[i][k] += features[i][j][k]
features_normalize[i][k] /= len(features[i])
return features_normalize
def Kcat_predict(feature, pH, sequence, smiles, Label):
# Generate index
Train_Validation_index = random.sample(range(len(feature)), int(len(feature)*0.8))
Test_index = []
for i in range(len(feature)):
if i not in Train_Validation_index:
Test_index.append(i)
Validation_index = random.sample(Train_Validation_index, int(len(Train_Validation_index)*0.2))
Train_index = []
for i in range(len(feature)):
if i not in Validation_index and i not in Test_index:
Train_index.append(i)
print(len(Train_index), len(Validation_index), len(Test_index))
Training_Validation_Test = []
for i in range(len(feature)):
if i in Train_index:
Training_Validation_Test.append(0)
elif i in Validation_index:
Training_Validation_Test.append(1)
else:
Training_Validation_Test.append(2)
Train_index = np.array(Train_index)
Validation_index = np.array(Validation_index)
Test_index = np.array(Test_index)
print(Train_index.shape, Validation_index.shape, Test_index.shape)
# First model
print(feature[Train_index].shape, pH[Train_index].shape)
model_1_input = np.concatenate((feature[Train_index], pH[Train_index]), axis=1)
model_first = ExtraTreesRegressor()
model_first.fit(model_1_input, Label[Train_index])
# Second model
with open("PreKcat_new/0_model.pkl", "rb") as f:
model_base = pickle.load(f)
Kcat_baseline = model_base.predict(feature[Validation_index]).reshape([len(Validation_index), 1])
model_1_2_input = np.concatenate((feature[Validation_index], pH[Validation_index]), axis=1)
Kcat_calibrated = model_first.predict(model_1_2_input).reshape([len(Validation_index), 1])
kcat_fused = np.concatenate((Kcat_baseline, Kcat_calibrated), axis=1)
model_second = LinearRegression()
model_second.fit(kcat_fused, Label[Validation_index])
# Final prediction
model_1_3_input = np.concatenate((feature, pH), axis=1)
Kcat_calibrated_3 = model_first.predict(model_1_3_input).reshape([len(feature), 1])
Kcat_baseline_3 = model_base.predict(feature).reshape([len(feature), 1])
kcat_fused_3 = np.concatenate((Kcat_baseline_3, Kcat_calibrated_3), axis=1)
Predicted_value = model_second.predict(kcat_fused_3).reshape([len(feature)])
Training_Validation_Test = np.array(Training_Validation_Test).reshape([len(feature)])
pH = np.array(pH).reshape([len(Label)])
Kcat_baseline_3 = np.array(Kcat_baseline_3).reshape([len(feature)])
Kcat_calibrated_3 = np.array(Kcat_calibrated_3).reshape([len(feature)])
print(Training_Validation_Test.shape)
# save
res = pd.DataFrame({'Value': Label,
'sequence': sequence,
'smiles': smiles,
'pH': pH,
'Prediction_first_base': Kcat_baseline_3,
'Prediction_first_pH': Kcat_calibrated_3,
'Prediction_second': Predicted_value,
'Training_Validation_Test': Training_Validation_Test})
res.to_excel('pH/s2_pH_Kcat.xlsx')
if __name__ == '__main__':
# Dataset Load
database = np.array(pd.read_excel(PH_SMILES_PATH)).T
sequence = database[1]
smiles = database[3]
pH = database[5]
Label = database[4]
for i in range(len(Label)):
Label[i] = math.log(Label[i], 10)
print(max(Label), min(Label))
pH = np.array(pH).reshape([len(Label), 1])
with open("pH/features_636_pH_PreKcat.pkl", "rb") as f:
feature = pickle.load(f)
# Modelling
Kcat_predict(feature, pH, sequence, smiles, Label)