--- a
+++ b/HINT/model.py
@@ -0,0 +1,915 @@
+from sklearn.metrics import roc_auc_score, f1_score, average_precision_score, precision_score, recall_score, accuracy_score
+import matplotlib.pyplot as plt
+from copy import deepcopy 
+import numpy as np 
+from tqdm import tqdm 
+import torch 
+torch.manual_seed(0)
+from torch import nn 
+from torch.autograd import Variable
+import torch.nn.functional as F
+from HINT.module import Highway, GCN 
+from functools import reduce 
+import pickle
+
+
+class Interaction(nn.Sequential):
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 
+					device, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0, 
+					):
+		super(Interaction, self).__init__()
+		self.molecule_encoder = molecule_encoder 
+		self.disease_encoder = disease_encoder 
+		self.protocol_encoder = protocol_encoder 
+		self.global_embed_size = global_embed_size 
+		self.highway_num_layer = highway_num_layer 
+		self.feature_dim = self.molecule_encoder.embedding_size + self.disease_encoder.embedding_size + self.protocol_encoder.embedding_size
+		self.epoch = epoch 
+		self.lr = lr 
+		self.weight_decay = weight_decay 
+		self.save_name = prefix_name + '_interaction'
+
+		self.f = F.relu
+		self.loss = nn.BCEWithLogitsLoss()
+
+		##### NN 
+		self.encoder2interaction_fc = nn.Linear(self.feature_dim, self.global_embed_size).to(device)
+		self.encoder2interaction_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)
+		self.pred_nn = nn.Linear(self.global_embed_size, 1)
+
+		self.device = device 
+		self = self.to(device)
+
+	def feed_lst_of_module(self, input_feature, lst_of_module):
+		x = input_feature
+		for single_module in lst_of_module:
+			x = self.f(single_module(x))
+		return x
+
+	def forward_get_three_encoders(self, smiles_lst2, icdcode_lst3, criteria_lst):
+		molecule_embed = self.molecule_encoder.forward_smiles_lst_lst(smiles_lst2)
+		icd_embed = self.disease_encoder.forward_code_lst3(icdcode_lst3)
+		protocol_embed = self.protocol_encoder.forward(criteria_lst)
+		return molecule_embed, icd_embed, protocol_embed	
+
+	def forward_encoder_2_interaction(self, molecule_embed, icd_embed, protocol_embed):
+		encoder_embedding = torch.cat([molecule_embed, icd_embed, protocol_embed], 1)
+		# interaction_embedding = self.feed_lst_of_module(encoder_embedding, [self.encoder2interaction_fc, self.encoder2interaction_highway])
+		h = self.encoder2interaction_fc(encoder_embedding)
+		h = self.f(h)
+		h = self.encoder2interaction_highway(h)
+		interaction_embedding = self.f(h)
+		return interaction_embedding 
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):
+		molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)
+		interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)
+		output = self.pred_nn(interaction_embedding)
+		return output ### 32, 1
+
+	def evaluation(self, predict_all, label_all, threshold = 0.5):
+		import pickle, os
+		from sklearn.metrics import roc_curve, precision_recall_curve
+		with open("predict_label.txt", 'w') as fout:
+			for i,j in zip(predict_all, label_all):
+				fout.write(str(i)[:6] + '\t' + str(j)[:4]+'\n')
+		auc_score = roc_auc_score(label_all, predict_all)
+		figure_folder = "figure"
+		#### ROC-curve 
+		fpr, tpr, thresholds = roc_curve(label_all, predict_all, pos_label=1)
+		# roc_curve =plt.figure()
+		# plt.plot(fpr,tpr,'-',label=self.save_name + ' ROC Curve ')
+		# plt.legend(fontsize = 15)
+		# plt.savefig(os.path.join(figure_folder,self.save_name+"_roc_curve.png"))
+		#### PR-curve
+		precision, recall, thresholds = precision_recall_curve(label_all, predict_all)
+		# plt.plot(recall,precision, label = self.save_name + ' PR Curve')
+		# plt.legend(fontsize = 15)
+		# plt.savefig(os.path.join(figure_folder,self.save_name + "_pr_curve.png"))
+		label_all = [int(i) for i in label_all]
+		float2binary = lambda x:0 if x < threshold else 1
+		predict_all = list(map(float2binary, predict_all))
+		f1score = f1_score(label_all, predict_all)
+		prauc_score = average_precision_score(label_all, predict_all)
+		# print(predict_all)
+		precision = precision_score(label_all, predict_all)
+		recall = recall_score(label_all, predict_all)
+		accuracy = accuracy_score(label_all, predict_all)
+		predict_1_ratio = sum(predict_all) / len(predict_all)
+		label_1_ratio = sum(label_all) / len(label_all)
+		return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 
+
+	def testloader_to_lst(self, dataloader):
+		nctid_lst, label_lst, smiles_lst2, icdcode_lst3, criteria_lst = [], [], [], [], []
+		for nctid, label, smiles, icdcode, criteria in dataloader:
+			nctid_lst.extend(nctid)
+			label_lst.extend([i.item() for i in label])
+			smiles_lst2.extend(smiles)
+			icdcode_lst3.extend(icdcode)
+			criteria_lst.extend(criteria)
+		length = len(nctid_lst)
+		assert length == len(smiles_lst2) and length == len(icdcode_lst3)
+		return nctid_lst, label_lst, smiles_lst2, icdcode_lst3, criteria_lst, length 
+
+	def generate_predict(self, dataloader):
+		whole_loss = 0 
+		label_all, predict_all, nctid_all = [], [], []
+		for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:
+			nctid_all.extend(nctid_lst)
+			label_vec = label_vec.to(self.device)
+			output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst).view(-1)  
+			loss = self.loss(output, label_vec.float())
+			whole_loss += loss.item()
+			predict_all.extend([i.item() for i in torch.sigmoid(output)])
+			label_all.extend([i.item() for i in label_vec])
+
+		return whole_loss, predict_all, label_all, nctid_all
+
+	def bootstrap_test(self, dataloader, valid_loader = None, sample_num = 20):
+		best_threshold = 0.5
+		# if validloader is not None:
+		# 	best_threshold = self.select_threshold_for_binary(valid_loader)
+		# 	print(f"best_threshold: {best_threshold}")
+		self.eval()
+		whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader)
+		from HINT.utils import plot_hist
+		plt.clf()
+		prefix_name = "./figure/" + self.save_name 
+		plot_hist(prefix_name, predict_all, label_all)		
+		def bootstrap(length, sample_num):
+			idx = [i for i in range(length)]
+			from random import choices 
+			bootstrap_idx = [choices(idx, k = length) for i in range(sample_num)]
+			return bootstrap_idx 
+		results_lst = []
+		bootstrap_idx_lst = bootstrap(len(predict_all), sample_num = sample_num)
+		for bootstrap_idx in bootstrap_idx_lst: 
+			bootstrap_label = [label_all[idx] for idx in bootstrap_idx]		
+			bootstrap_predict = [predict_all[idx] for idx in bootstrap_idx]
+			results = self.evaluation(bootstrap_predict, bootstrap_label, threshold = best_threshold)
+			results_lst.append(results)
+		self.train() 
+		auc = [results[0] for results in results_lst]
+		f1score = [results[1] for results in results_lst]
+		prauc_score = [results[2] for results in results_lst]
+		print("PR-AUC   mean: "+str(np.mean(prauc_score))[:6], "std: "+str(np.std(prauc_score))[:6])
+		print("F1       mean: "+str(np.mean(f1score))[:6], "std: "+str(np.std(f1score))[:6])
+		print("ROC-AUC  mean: "+str(np.mean(auc))[:6], "std: "+str(np.std(auc))[:6])
+
+		for nctid, label, predict in zip(nctid_all, label_all, predict_all):
+			if (predict > 0.5 and label == 0) or (predict < 0.5 and label == 1):
+				print(nctid, label, str(predict)[:6])
+
+		nctid2predict = {nctid:predict for nctid, predict in zip(nctid_all, predict_all)} 
+		pickle.dump(nctid2predict, open('results/nctid2predict.pkl', 'wb'))
+		return nctid_all, predict_all 
+
+	def ongoing_test(self, dataloader, sample_num = 20):
+		self.eval()
+		best_threshold = 0.5 
+		whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader) 
+		self.train() 
+		return nctid_all, predict_all 
+		
+	def test(self, dataloader, return_loss = True, validloader=None):
+		# if validloader is not None:
+		# 	best_threshold = self.select_threshold_for_binary(validloader)
+		self.eval()
+		best_threshold = 0.5 
+		whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader)
+		# from HINT.utils import plot_hist
+		# plt.clf()
+		# prefix_name = "./figure/" + self.save_name 
+		# plot_hist(prefix_name, predict_all, label_all)
+		self.train()
+		if return_loss:
+			return whole_loss, predict_all, label_all
+		else:
+			print_num = 6
+			auc_score, f1score, prauc_score, precision, recall, accuracy, \
+			predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)
+			print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \
+				 + "\nPR-AUC: " + str(prauc_score)[:print_num] \
+				 + "\nPrecision: " + str(precision)[:print_num] \
+				 + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \
+				 + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \
+				 + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])
+			return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 
+
+	def learn(self, train_loader, valid_loader, test_loader):
+		opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)
+		train_loss_record = [] 
+		valid_loss, valid_predict, valid_label = self.test(valid_loader, return_loss=True)
+		valid_loss_record = [valid_loss]
+		best_valid_loss = valid_loss
+		best_model = deepcopy(self)
+		train_output = []
+		valid_output = []
+		for ep in tqdm(range(self.epoch)):
+			for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:
+				label_vec = label_vec.to(self.device)
+				output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,
+				loss = self.loss(output, label_vec.float())
+				train_loss_record.append(loss.item())
+				train_output.append((loss.item(), output, label_vec))
+				opt.zero_grad()
+				loss.backward()
+				opt.step()
+			valid_loss, valid_predict, valid_label = self.test(valid_loader, return_loss=True)
+			valid_loss_record.append(valid_loss)
+			valid_output.append((valid_loss, valid_predict, valid_label))
+
+			print(f"valid_loss: {valid_loss}")
+			print(best_valid_loss)
+			if valid_loss < best_valid_loss:
+				best_valid_loss = valid_loss 
+				best_model = deepcopy(self)
+
+		self.plot_learning_curve(train_loss_record, valid_loss_record)
+		self = deepcopy(best_model)
+		auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)
+		return train_output, valid_output
+
+	def plot_learning_curve(self, train_loss_record, valid_loss_record):
+		plt.plot(train_loss_record)
+		plt.savefig("./figure/" + self.save_name + '_train_loss.jpg')
+		plt.clf() 
+		plt.plot(valid_loss_record)
+		plt.savefig("./figure/" + self.save_name + '_valid_loss.jpg')
+		plt.clf() 
+
+	def select_threshold_for_binary(self, validloader):
+		_, prediction, label_all, nctid_all = self.generate_predict(validloader)
+		best_f1 = 0
+		for threshold in prediction:
+			float2binary = lambda x:0 if x<threshold else 1
+			predict_all = list(map(float2binary, prediction))
+			f1score = precision_score(label_all, predict_all)        
+			if f1score > best_f1:
+				best_f1 = f1score 
+				best_threshold = threshold
+		return best_threshold 
+
+
+class HINTModel_multi(Interaction):
+
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 
+					device, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0, 
+					):
+		super(HINTModel_multi, self).__init__(molecule_encoder = molecule_encoder, 
+								   disease_encoder = disease_encoder, 
+								   protocol_encoder = protocol_encoder, 
+								   device = device, 
+								   prefix_name = prefix_name, 
+								   global_embed_size = global_embed_size, 
+								   highway_num_layer = highway_num_layer,
+								   epoch = epoch,
+								   lr = lr, 
+								   weight_decay = weight_decay)
+		self.pred_nn = nn.Linear(self.global_embed_size, 4)
+		self.loss = nn.CrossEntropyLoss()
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):
+		molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)
+		interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)
+		output = self.pred_nn(interaction_embedding)
+		return output ### 32, 4
+
+	def generate_predict(self, dataloader):
+		whole_loss = 0 
+		label_all, predict_all = [], []
+		for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:
+			label_vec = label_vec.to(self.device)
+			output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst) 
+			loss = self.loss(output, label_vec)
+			whole_loss += loss.item()
+			predict_all.extend(torch.argmax(output, 1).tolist())
+			# predict_all.extend([i.item() for i in torch.sigmoid(output)])
+			label_all.extend([i.item() for i in label_vec])
+
+		accuracy = len(list(filter(lambda x:x[0]==x[1], zip(predict_all, label_all)))) / len(label_all)
+		return whole_loss, predict_all, label_all, accuracy
+
+	def test(self, dataloader, return_loss = True, validloader=None):
+		# if validloader is not None:
+		# 	best_threshold = self.select_threshold_for_binary(validloader)
+		self.eval()
+		whole_loss, predict_all, label_all, accuracy = self.generate_predict(dataloader)
+		self.train()
+		return whole_loss, predict_all, label_all, accuracy
+		# # from HINT.utils import plot_hist
+		# # plt.clf()
+		# # prefix_name = "./figure/" + self.save_name 
+		# # plot_hist(prefix_name, predict_all, label_all)
+		# self.train()
+		# if return_loss:
+		# 	return whole_loss
+		# else:
+		# 	print_num = 5
+		# 	auc_score, f1score, prauc_score, precision, recall, accuracy, \
+		# 	predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)
+		# 	print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \
+		# 		 + "\nPR-AUC: " + str(prauc_score)[:print_num] \
+		# 		 + "\nPrecision: " + str(precision)[:print_num] \
+		# 		 + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \
+		# 		 + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \
+		# 		 + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])
+		# 	return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 
+
+	def learn(self, train_loader, valid_loader, test_loader):
+		opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)
+		train_loss_record = []
+		valid_loss, predict_all, label_all, accuracy = self.test(valid_loader, return_loss=True)
+		print('accuracy', accuracy)
+		# valid_loss_record = [valid_loss]
+		# best_valid_loss = valid_loss
+		best_model = deepcopy(self)
+		for ep in tqdm(range(self.epoch)):
+			self.train() 
+			for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:
+				label_vec = label_vec.to(self.device)
+				output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst)  #### 32, 1 -> 32, ||  label_vec 32,
+				# print(label_vec.shape, output.shape, label_vec, output)
+				loss = self.loss(output, label_vec)
+				train_loss_record.append(loss.item())
+				opt.zero_grad() 
+				loss.backward() 
+				opt.step()
+			valid_loss, predict_all, label_all, accuracy = self.test(valid_loader, return_loss=True)
+			print('accuracy', accuracy)
+		return predict_all, label_all
+		# 	valid_loss_record.append(valid_loss)
+		# 	if valid_loss < best_valid_loss:
+		# 		best_valid_loss = valid_loss 
+		# 		best_model = deepcopy(self)
+
+		# self.plot_learning_curve(train_loss_record, valid_loss_record)
+		# self = deepcopy(best_model)
+		# auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)
+
+
+class HINT_nograph(Interaction):
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, device, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0, ):
+		super(HINT_nograph, self).__init__(molecule_encoder = molecule_encoder, 
+								   disease_encoder = disease_encoder, 
+								   protocol_encoder = protocol_encoder,
+								   device = device,  
+								   global_embed_size = global_embed_size, 
+								   prefix_name = prefix_name, 
+								   highway_num_layer = highway_num_layer,
+								   epoch = epoch,
+								   lr = lr, 
+								   weight_decay = weight_decay, 
+								   ) 
+		self.save_name = prefix_name + '_HINT_nograph'
+		'''	### interaction model 
+		self.molecule_encoder = molecule_encoder 
+		self.disease_encoder = disease_encoder 
+		self.protocol_encoder = protocol_encoder 
+		self.global_embed_size = global_embed_size 
+		self.highway_num_layer = highway_num_layer 
+		self.feature_dim = self.molecule_encoder.embedding_size + self.disease_encoder.embedding_size + self.protocol_encoder.embedding_size
+		self.epoch = epoch 
+		self.lr = lr 
+		self.weight_decay = weight_decay 
+		self.save_name = save_name
+
+		self.f = F.relu
+		self.loss = nn.BCEWithLogitsLoss()
+
+		##### NN 
+		self.encoder2interaction_fc = nn.Linear(self.feature_dim, self.global_embed_size)
+		self.encoder2interaction_highway = Highway(self.global_embed_size, self.highway_num_layer)
+		self.pred_nn = nn.Linear(self.global_embed_size, 1)
+		'''
+
+		#### risk of disease 
+		self.risk_disease_fc = nn.Linear(self.disease_encoder.embedding_size, self.global_embed_size)
+		self.risk_disease_higway = Highway(self.global_embed_size, self.highway_num_layer)
+
+		#### augment interaction 
+		self.augment_interaction_fc = nn.Linear(self.global_embed_size*2, self.global_embed_size)
+		self.augment_interaction_highway = Highway(self.global_embed_size, self.highway_num_layer)
+
+		#### ADMET 
+		self.admet_model = []
+		for i in range(5):
+			admet_fc = nn.Linear(self.molecule_encoder.embedding_size, self.global_embed_size).to(device)
+			admet_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)
+			self.admet_model.append(nn.ModuleList([admet_fc, admet_highway])) 
+		self.admet_model = nn.ModuleList(self.admet_model)
+
+		#### PK 
+		self.pk_fc = nn.Linear(self.global_embed_size*5, self.global_embed_size)
+		self.pk_highway = Highway(self.global_embed_size, self.highway_num_layer)
+
+		#### trial node 
+		self.trial_fc = nn.Linear(self.global_embed_size*2, self.global_embed_size)
+		self.trial_highway = Highway(self.global_embed_size, self.highway_num_layer)
+
+		## self.pred_nn = nn.Linear(self.global_embed_size, 1)
+
+		self.device = device 
+		self = self.to(device)
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst, if_gnn = False):
+		### encoder for molecule, disease and protocol
+		molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)
+		### interaction 
+		interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)
+		### risk of disease 
+		risk_of_disease_embedding = self.feed_lst_of_module(input_feature = icd_embed, 
+															lst_of_module = [self.risk_disease_fc, self.risk_disease_higway])
+		### augment interaction   
+		augment_interaction_input = torch.cat([interaction_embedding, risk_of_disease_embedding], 1)
+		augment_interaction_embedding = self.feed_lst_of_module(input_feature = augment_interaction_input, 
+																lst_of_module = [self.augment_interaction_fc, self.augment_interaction_highway])
+		### admet 
+		admet_embedding_lst = []
+		for idx in range(5):
+			admet_embedding = self.feed_lst_of_module(input_feature = molecule_embed, 
+													  lst_of_module = self.admet_model[idx])
+			admet_embedding_lst.append(admet_embedding)
+		### pk 
+		pk_input = torch.cat(admet_embedding_lst, 1)
+		pk_embedding = self.feed_lst_of_module(input_feature = pk_input, 
+											   lst_of_module = [self.pk_fc, self.pk_highway])
+		### trial 
+		trial_input = torch.cat([pk_embedding, augment_interaction_embedding], 1)
+		trial_embedding = self.feed_lst_of_module(input_feature = trial_input, 
+												  lst_of_module = [self.trial_fc, self.trial_highway])
+		output = self.pred_nn(trial_embedding)
+		if if_gnn == False:
+			return output 
+		else:
+			embedding_lst = [molecule_embed, icd_embed, protocol_embed, interaction_embedding, risk_of_disease_embedding, \
+							 augment_interaction_embedding] + admet_embedding_lst + [pk_embedding, trial_embedding]
+			return embedding_lst
+
+
+class HINTModel(HINT_nograph):
+
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 
+					device, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					gnn_hidden_size, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0,):
+		super(HINTModel, self).__init__(molecule_encoder = molecule_encoder, 
+								   disease_encoder = disease_encoder, 
+								   protocol_encoder = protocol_encoder, 
+								   device = device, 
+								   prefix_name = prefix_name, 
+								   global_embed_size = global_embed_size, 
+								   highway_num_layer = highway_num_layer,
+								   epoch = epoch,
+								   lr = lr, 
+								   weight_decay = weight_decay)
+		self.save_name = prefix_name 
+		self.gnn_hidden_size = gnn_hidden_size 
+		#### GNN 
+		self.adj = self.generate_adj()          
+		self.gnn = GCN(
+            nfeat = self.global_embed_size,
+            nhid = self.gnn_hidden_size,
+            nclass = 1,
+            dropout = 0.6,
+		    init = 'uniform') 
+		### gnn's attention 		
+		self.node_size = self.adj.shape[0]
+		'''
+		self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() \
+        								if self.adj[i,j]==1 else None  \
+        								for j in range(self.node_size)]) \
+        								for i in range(self.node_size)])
+        '''
+		self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() if self.adj[i,j]==1 else None for j in range(self.node_size)]) for i in range(self.node_size)])
+		# self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() if self.adj[i,j]==1 else None for j in range(self.node_size)]) for i in range(self.node_size)])
+
+		'''
+nn.ModuleList([ nn.ModuleList([nn.Linear(3,2) for j in range(5)] + [None]) for i in range(3)])
+		'''
+
+		self.device = device 
+		self = self.to(device)
+
+	def generate_adj(self):        								
+		##### consistent with HINT_nograph.forward
+		lst = ["molecule", "disease", "criteria", 'INTERACTION', 'risk_disease', 'augment_interaction', 'A', 'D', 'M', 'E', 'T', 'PK', "final"]
+		edge_lst = [("disease", "molecule"), ("disease", "criteria"), ("molecule", "criteria"), 
+					("disease", "INTERACTION"), ("molecule", "INTERACTION"),  ("criteria", "INTERACTION"), 
+					("disease", "risk_disease"), ('risk_disease', 'augment_interaction'), ('INTERACTION', 'augment_interaction'),
+					("molecule", "A"), ("molecule", "D"), ("molecule", "M"), ("molecule", "E"), ("molecule", "T"),
+					('A', 'PK'), ('D', 'PK'), ('M', 'PK'), ('E', 'PK'), ('T', 'PK'), 
+					('augment_interaction', 'final'), ('PK', 'final')]
+		adj = torch.zeros(len(lst), len(lst))
+		adj = torch.eye(len(lst)) * len(lst)
+		num2str = {k:v for k,v in enumerate(lst)}
+		str2num = {v:k for k,v in enumerate(lst)}
+		for i,j in edge_lst:
+			n1,n2 = str2num[i], str2num[j]
+			adj[n1,n2] = 1
+			adj[n2,n1] = 1
+		return adj.to(self.device) 
+
+	def generate_attention_matrx(self, node_feature_mat):
+		attention_mat = torch.zeros(self.node_size, self.node_size).to(self.device)
+		for i in range(self.node_size):
+			for j in range(self.node_size):
+				if self.adj[i,j]!=1:
+					continue 
+				feature = torch.cat([node_feature_mat[i].view(1,-1), node_feature_mat[j].view(1,-1)], 1)
+				attention_model = self.graph_attention_model_mat[i][j]
+				attention_mat[i,j] = torch.sigmoid(self.feed_lst_of_module(input_feature=feature, lst_of_module=attention_model))
+		return attention_mat 
+
+	##### self.global_embed_size*2 -> 1 
+	def gnn_attention(self):
+		highway_nn = Highway(size = self.global_embed_size*2, num_layers = self.highway_num_layer).to(self.device)
+		highway_fc = nn.Linear(self.global_embed_size*2, 1).to(self.device)
+		return nn.ModuleList([highway_nn, highway_fc])	
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst, return_attention_matrix = False):
+		embedding_lst = HINT_nograph.forward(self, smiles_lst2, icdcode_lst3, criteria_lst, if_gnn = True)
+		### length is 13, each is 32,50 
+		batch_size = embedding_lst[0].shape[0]
+		output_lst = []
+		if return_attention_matrix:
+			attention_mat_lst = []
+		for i in range(batch_size):
+			node_feature_lst = [embedding[i].view(1,-1) for embedding in embedding_lst]
+			node_feature_mat = torch.cat(node_feature_lst, 0) ### 13, 50 
+			attention_mat = self.generate_attention_matrx(node_feature_mat)
+			output = self.gnn(node_feature_mat, self.adj * attention_mat)
+			output = output[-1].view(1,-1)
+			output_lst.append(output)
+			if return_attention_matrix:
+				attention_mat_lst.append(attention_mat)
+		output_mat = torch.cat(output_lst, 0)
+		if not return_attention_matrix:
+			return output_mat 
+		else:
+			return output_mat, attention_mat_lst
+
+	def interpret(self, complete_dataloader):
+		from graph_visualize_interpret import data2graph 
+		from HINT.utils import replace_strange_symbol
+		for nctid_lst, status_lst, why_stop_lst, label_vec, phase_lst, \
+			diseases_lst, icdcode_lst3, drugs_lst, smiles_lst2, criteria_lst in complete_dataloader: 
+			output, attention_mat_lst = self.forward(smiles_lst2, icdcode_lst3, criteria_lst, return_attention_matrix=True)
+			output = output.view(-1)
+			batch_size = len(nctid_lst)
+			for i in range(batch_size):
+				name = '__'.join([nctid_lst[i], status_lst[i], why_stop_lst[i], \
+														str(label_vec[i].item()), str(torch.sigmoid(output[i]).item())[:5], \
+														phase_lst[i], diseases_lst[i], drugs_lst[i]])
+				if len(name) > 150:
+					name = name[:250]
+				name = replace_strange_symbol(name)
+				name = name.replace('__', '_')
+				name = name.replace('  ', ' ')
+				name = 'interpret_result/' + name + '.png'
+				print(name)
+				data2graph(attention_matrix = attention_mat_lst[i], adj = self.adj, save_name = name)
+
+	def init_pretrain(self, admet_model):
+		self.molecule_encoder = admet_model.molecule_encoder
+
+	### generate attention matrix 
+
+
+class Only_Molecule(Interaction):
+
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0):
+		super(Only_Molecule, self).__init__(molecule_encoder=molecule_encoder, 
+											disease_encoder=disease_encoder, 
+											protocol_encoder=protocol_encoder, 
+											global_embed_size = global_embed_size, 
+											highway_num_layer = highway_num_layer,
+											prefix_name = prefix_name, 
+											epoch = epoch,
+											lr = lr, 
+											weight_decay = weight_decay,)
+		self.molecule2out = nn.Linear(self.global_embed_size,1)
+
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):
+		molecule_embed = self.molecule_encoder.forward_smiles_lst_lst(smiles_lst2)
+		return self.molecule2out(molecule_embed)
+
+class Only_Disease(Only_Molecule):
+
+	def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 
+					global_embed_size, 
+					highway_num_layer,
+					prefix_name, 
+					epoch = 20,
+					lr = 3e-4, 
+					weight_decay = 0):
+		super(Only_Disease, self).__init__(molecule_encoder = molecule_encoder, 
+											disease_encoder=disease_encoder, 
+											protocol_encoder=protocol_encoder, 
+											global_embed_size = global_embed_size, 
+											highway_num_layer = highway_num_layer,
+											prefix_name = prefix_name, 
+											epoch = epoch,
+											lr = lr, 
+											weight_decay = weight_decay,)
+		self.disease2out = self.molecule2out 
+
+
+	def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):
+		icd_embed = self.disease_encoder.forward_code_lst3(icdcode_lst3)
+		return self.disease2out(icd_embed)
+
+def dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, global_icd):
+	## label_vec: (n,)
+	y = label_vec 
+
+	num_icd = len(global_icd)
+	from HINT.utils import smiles_lst2fp 
+	fp_lst = [smiles_lst2fp(smiles_lst).reshape(1,-1) for smiles_lst in smiles_lst2]
+	fp_mat = np.concatenate(fp_lst, 0)
+	# fp_mat = torch.from_numpy(fp_mat)  ### (n,2048)
+
+	icdcode_lst = []
+	for lst2 in icdcode_lst3:
+		lst = list(reduce(lambda x,y:x+y, lst2))
+		lst = [i.split('.')[0] for i in lst]
+		lst = set(lst)	
+		icd_feature = np.zeros((1,num_icd), np.int32)
+		for ele in lst:
+			if ele in global_icd:
+				idx = global_icd.index(ele)
+				icd_feature[0,idx] = 1 
+		icdcode_lst.append(icd_feature)
+	icdcode_mat = np.concatenate(icdcode_lst, 0)
+	X = np.concatenate([fp_mat, icdcode_mat], 1)
+	X = torch.from_numpy(X)
+	X = X.float()
+	# icdcode_mat = torch.from_numpy(icdcode_mat) 
+
+	# X = torch.cat([fp_mat, icdcode_mat], 1)
+	return X, y 
+
+
+class FFNN(nn.Sequential):
+	def __init__(self, molecule_dim, diseasecode_dim, 
+					global_icd, 
+					protocol_dim = 0,
+					prefix_name = 'FFNN', 
+					epoch = 10,
+					lr = 3e-4, 
+					weight_decay = 0, 
+					):
+		super(FFNN, self).__init__()
+		self.molecule_dim = molecule_dim 
+		self.diseasecode_dim = diseasecode_dim 
+		self.protocol_dim = protocol_dim 
+		self.prefix_name = prefix_name 
+		self.epoch = epoch 
+		self.lr = lr 
+		self.weight_decay = weight_decay 
+		self.global_icd = global_icd 
+		self.num_icd = len(global_icd)
+
+		self.fc_dims = [self.molecule_dim + self.diseasecode_dim + self.protocol_dim, 2000, 1000, 200, 50, 1]
+		self.fc_layers = nn.ModuleList([nn.Linear(v,self.fc_dims[i+1]) for i,v in enumerate(self.fc_dims[:-1])])
+		self.loss = nn.BCEWithLogitsLoss()
+		self.save_name = prefix_name 
+
+	def forward(self, X):
+		for i in range(len(self.fc_layers) - 1):
+			fc_layer = self.fc_layers[i]
+			X = fc_layer(X)
+		last_layer = self.fc_layers[-1]
+		pred = F.sigmoid(last_layer(X))
+		return pred 
+
+	def learn(self, train_loader, valid_loader, test_loader):
+		opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)
+		train_loss_record = [] 
+		valid_loss = self.test(valid_loader, return_loss=True)
+		valid_loss_record = [valid_loss]
+		best_valid_loss = valid_loss
+		best_model = deepcopy(self)
+
+		for ep in tqdm(range(self.epoch)):
+			for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:
+				X, _ = dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, self.global_icd)
+				output = self.forward(X).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,
+				loss = self.loss(output, label_vec.float())
+				train_loss_record.append(loss.item())
+				opt.zero_grad() 
+				loss.backward() 
+				opt.step()
+			valid_loss = self.test(valid_loader, return_loss=True)
+			valid_loss_record.append(valid_loss)
+			if valid_loss < best_valid_loss:
+				best_valid_loss = valid_loss 
+				best_model = deepcopy(self)
+
+		self.plot_learning_curve(train_loss_record, valid_loss_record)
+		self = deepcopy(best_model)
+		auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)
+
+	def evaluation(self, predict_all, label_all, threshold = 0.5):
+		import pickle, os
+		from sklearn.metrics import roc_curve, precision_recall_curve
+		with open("predict_label.txt", 'w') as fout:
+			for i,j in zip(predict_all, label_all):
+				fout.write(str(i)[:4] + '\t' + str(j)[:4]+'\n')
+		auc_score = roc_auc_score(label_all, predict_all)
+		figure_folder = "figure"
+		#### ROC-curve 
+		fpr, tpr, thresholds = roc_curve(label_all, predict_all, pos_label=1)
+		# roc_curve =plt.figure()
+		# plt.plot(fpr,tpr,'-',label=self.save_name + ' ROC Curve ')
+		# plt.legend(fontsize = 15)
+		#plt.savefig(os.path.join(figure_folder,name+"_roc_curve.png"))
+		#### PR-curve
+		precision, recall, thresholds = precision_recall_curve(label_all, predict_all)
+		# plt.plot(recall,precision, label = self.save_name + ' PR Curve')
+		# plt.legend(fontsize = 15)
+		# plt.savefig(os.path.join(figure_folder,self.save_name + "_pr_curve.png"))
+		label_all = [int(i) for i in label_all]
+		float2binary = lambda x:0 if x<threshold else 1
+		predict_all = list(map(float2binary, predict_all))
+		f1score = f1_score(label_all, predict_all)
+		prauc_score = average_precision_score(label_all, predict_all)
+		# print(predict_all)
+		precision = precision_score(label_all, predict_all)
+		recall = recall_score(label_all, predict_all)
+		accuracy = accuracy_score(label_all, predict_all)
+		predict_1_ratio = sum(predict_all) / len(predict_all)
+		label_1_ratio = sum(label_all) / len(label_all)
+		return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 
+
+	def generate_predict(self, dataloader):
+		whole_loss = 0 
+		label_all, predict_all = [], []
+		for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:
+			X, _ = dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, self.global_icd) 
+			output = self.forward(X).view(-1)  
+			loss = self.loss(output, label_vec.float())
+			whole_loss += loss.item()
+			predict_all.extend([i.item() for i in torch.sigmoid(output)])
+			label_all.extend([i.item() for i in label_vec])
+
+		return whole_loss, predict_all, label_all
+
+	def bootstrap_test(self, dataloader, validloader = None, sample_num = 20):
+		best_threshold = 0.5
+		# if validloader is not None:
+		# 	best_threshold = self.select_threshold_for_binary(validloader)
+		self.eval()
+		whole_loss, predict_all, label_all = self.generate_predict(dataloader)
+		from HINT.utils import plot_hist
+		plt.clf()
+		prefix_name = "./figure/" + self.save_name 
+		plot_hist(prefix_name, predict_all, label_all)		
+		def bootstrap(length, sample_num):
+			idx = [i for i in range(length)]
+			from random import choices 
+			bootstrap_idx = [choices(idx, k = length) for i in range(sample_num)]
+			return bootstrap_idx 
+		results_lst = []
+		bootstrap_idx_lst = bootstrap(len(predict_all), sample_num = sample_num)
+		for bootstrap_idx in bootstrap_idx_lst: 
+			bootstrap_label = [label_all[idx] for idx in bootstrap_idx]		
+			bootstrap_predict = [predict_all[idx] for idx in bootstrap_idx]
+			results = self.evaluation(bootstrap_predict, bootstrap_label, threshold = best_threshold)
+			results_lst.append(results)
+		self.train() 
+		auc = [results[0] for results in results_lst]
+		f1score = [results[1] for results in results_lst]
+		prauc_score = [results[2] for results in results_lst]
+		print("PR-AUC   mean: "+str(np.mean(prauc_score))[:6], "std: "+str(np.std(prauc_score))[:6])
+		print("F1       mean: "+str(np.mean(f1score))[:6], "std: "+str(np.std(f1score))[:6])
+		print("ROC-AUC  mean: "+ str(np.mean(auc))[:6], "std: " + str(np.std(auc))[:6])
+
+	def test(self, dataloader, return_loss = True, validloader=None):
+		# if validloader is not None:
+		# 	best_threshold = self.select_threshold_for_binary(validloader)
+		self.eval()
+		best_threshold = 0.5 
+		whole_loss, predict_all, label_all = self.generate_predict(dataloader)
+		# from HINT.utils import plot_hist
+		# plt.clf()
+		# prefix_name = "./figure/" + self.save_name 
+		# plot_hist(prefix_name, predict_all, label_all)
+		self.train()
+		if return_loss:
+			return whole_loss
+		else:
+			print_num = 5
+			auc_score, f1score, prauc_score, precision, recall, accuracy, \
+			predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)
+			print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \
+				 + "\nPR-AUC: " + str(prauc_score)[:print_num] \
+				 + "\nPrecision: " + str(precision)[:print_num] \
+				 + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \
+				 + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \
+				 + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])
+			return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 
+
+	def plot_learning_curve(self, train_loss_record, valid_loss_record):
+		plt.plot(train_loss_record)
+		plt.savefig("./figure/" + self.save_name + '_train_loss.jpg')
+		plt.clf() 
+		plt.plot(valid_loss_record)
+		plt.savefig("./figure/" + self.save_name + '_valid_loss.jpg')
+		plt.clf() 
+
+
+class ADMET(nn.Sequential):
+	def __init__(self, mpnn_model, device):
+		super(ADMET, self).__init__()
+		self.num = 5 
+		self.mpnn_model = mpnn_model 
+		self.device = device 
+		self.mpnn_dim = mpnn_model.mpnn_hidden_size 
+		self.admet_model = []
+		self.global_embed_size = self.mpnn_dim 
+		self.highway_num_layer = 2 
+		for i in range(5):
+			admet_fc = nn.Linear(self.mpnn_model.mpnn_hidden_size, self.global_embed_size).to(device)
+			admet_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)
+			self.admet_model.append(nn.ModuleList([admet_fc, admet_highway]))
+		self.admet_model = nn.ModuleList(self.admet_model)
+
+		self.admet_pred = nn.ModuleList([nn.Linear(self.global_embed_size,1).to(device) for i in range(5)])
+		self.f = F.relu 
+
+		self.device = device 
+		self = self.to(device)
+
+	def feed_lst_of_module(self, input_feature, lst_of_module):
+		x = input_feature
+		for single_module in lst_of_module:
+			x = self.f(single_module(x))
+		return x 
+
+	def forward(self, smiles_lst, idx):
+		assert idx in list(range(5))
+		'''
+			xxxxxxxxxxxx
+		'''
+		embeds = self.mpnn_model.forward_smiles_lst_lst(smiles_lst)
+		embeds = self.feed_lst_of_module(embeds, self.admet_model[idx]) 
+		output = self.admet_pred[idx](embeds)
+		return output 
+
+	def test(self, valid_loader):
+		pass 
+
+	def learn(self, train_loader, valid_loader, idx):
+		opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)
+		train_loss_record = [] 
+		valid_loss = self.test(valid_loader, return_loss=True)
+		valid_loss_record = [valid_loss]
+		best_valid_loss = valid_loss
+		best_model = deepcopy(self)
+
+		for ep in tqdm(range(self.epoch)):
+			for smiles_lst in train_loader:
+				output = self.forward(smiles_lst).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,
+				loss = self.loss(output, label_vec.float())
+				train_loss_record.append(loss.item())
+				opt.zero_grad() 
+				loss.backward() 
+				opt.step()
+			valid_loss = self.test(valid_loader, return_loss=True)
+			valid_loss_record.append(valid_loss)
+			if valid_loss < best_valid_loss:
+				best_valid_loss = valid_loss 
+				best_model = deepcopy(self)
+
+		self = deepcopy(best_model)