specific

DiseaseSpecific/KG_extractor.py

@@ -0,0 +1,479 @@
+#%%
+import torch
+import numpy as np
+from torch.autograd import Variable
+from sklearn import metrics
+
+import datetime
+from typing import Dict, Tuple, List
+import logging
+import os
+import utils
+import pickle as pkl
+import json 
+import torch.backends.cudnn as cudnn
+
+from tqdm import tqdm
+
+import sys
+sys.path.append("..")
+import Parameters
+
+parser = utils.get_argument_parser()
+parser = utils.add_attack_parameters(parser)
+parser.add_argument('--mode', type=str, default='sentence', help='sentence, finetune, biogpt, bioBART')
+parser.add_argument('--action', type=str, default='parse', help='parse or extract')
+parser.add_argument('--ratio', type = str, default='', help='ratio of the number of changed words')
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ 
+utils.seed_all(args.seed)
+np.set_printoptions(precision=5)
+cudnn.benchmark = False
+
+data_path = os.path.join('processed_data', args.data)
+target_path = os.path.join(data_path, 'DD_target_{0}_{1}_{2}_{3}_{4}_{5}.txt'.format(args.model, args.data, args.target_split, args.target_size, 'exists:'+str(args.target_existed), args.attack_goal))
+attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate)))
+modified_attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ args.mode))
+attack_data = utils.load_data(attack_path, drop=False)
+#%%
+with open(os.path.join(data_path, 'entities_reverse_dict.json')) as fl:
+ id_to_meshid = json.load(fl)
+with open(os.path.join(data_path, 'entities_dict.json'), 'r') as fl:
+ meshid_to_id = json.load(fl)
+with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+ entity_raw_name = pkl.load(fl)
+with open(Parameters.GNBRfile+'retieve_sentence_through_edgetype', 'rb') as fl:
+ retieve_sentence_through_edgetype = pkl.load(fl)
+with open(Parameters.GNBRfile+'raw_text_of_each_sentence', 'rb') as fl:
+ raw_text_sen = pkl.load(fl)
+with open(Parameters.GNBRfile+'original_entity_raw_name', 'rb') as fl:
+ full_entity_raw_name = pkl.load(fl)
+for k, v in entity_raw_name.items():
+ assert v in full_entity_raw_name[k]
+
+#find unique
+once_set = set()
+twice_set = set()
+
+with open('generate_abstract/valid_entity.json', 'r') as fl:
+ valid_entity = json.load(fl)
+valid_entity = set(valid_entity)
+
+good_name = set()
+for k, v, in full_entity_raw_name.items():
+ names = list(v)
+ for name in names:
+ # if name == 'in a':
+ # print(names)
+ good_name.add(name)
+ # if name not in once_set:
+ # once_set.add(name)
+ # else:
+ # twice_set.add(name)
+# assert 'WNK4' in once_set
+# good_name = set.difference(once_set, twice_set)
+# assert 'in a' not in good_name
+# assert 'STE20' not in good_name
+# assert 'STE20' not in valid_entity
+# assert 'STE20-related proline-alanine-rich kinase' not in good_name
+# assert 'STE20-related proline-alanine-rich kinase' not in valid_entity
+# raise Exception
+
+name_to_type = {}
+name_to_meshid = {}
+
+for k, v, in full_entity_raw_name.items():
+ names = list(v)
+ for name in names:
+ if name in good_name:
+ name_to_type[name] = k.split('_')[0]
+ name_to_meshid[name] = k
+
+import spacy
+import networkx as nx
+import pprint
+
+def check(p, s):
+
+ if p < 1 or p >= len(s):
+ return True
+ return not((s[p]>='a' and s[p]<='z') or (s[p]>='A' and s[p]<='Z') or (s[p]>='0' and s[p]<='9'))
+
+def raw_to_format(sen):
+
+ text = sen
+ l = 0
+ ret = []
+ while(l < len(text)):
+ bo =False
+ if text[l] != ' ':
+ for i in range(len(text), l, -1): # reversing is important !!!
+ cc = text[l:i]
+ if (cc in good_name or cc in valid_entity) and check(l-1, text) and check(i, text):
+ ret.append(cc.replace(' ', '_'))
+ l = i
+ bo = True
+ break
+ if not bo:
+ ret.append(text[l])
+ l += 1
+ return ''.join(ret)
+
+if args.mode == 'sentence':
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_chat.json', 'r') as fl:
+ draft = json.load(fl)
+elif args.mode == 'finetune':
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_sentence_finetune.json', 'r') as fl:
+ draft = json.load(fl)
+elif args.mode == 'bioBART':
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}{args.ratio}_bioBART_finetune.json', 'r') as fl:
+ draft = json.load(fl)
+elif args.mode == 'biogpt':
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_biogpt.json', 'r') as fl:
+ draft = json.load(fl)
+else:
+ raise Exception('No!!!')
+
+nlp = spacy.load("en_core_web_sm")
+
+type_set = set()
+for aa in range(36):
+ dependency_sen_dict = retieve_sentence_through_edgetype[aa]['manual']
+ tmp_dict = retieve_sentence_through_edgetype[aa]['auto']
+ dependencys = list(dependency_sen_dict.keys()) + list(tmp_dict.keys())
+ for dependency in dependencys:
+ dep_list = dependency.split(' ')
+ for sub_dep in dep_list:
+ sub_dep_list = sub_dep.split('|')
+ assert(len(sub_dep_list) == 3)
+ type_set.add(sub_dep_list[1])
+# print('Type:', type_set)
+
+if args.action == 'parse':
+# dp_path, sen_list = list(dependency_sen_dict.items())[0]
+# check
+# paper_id, sen_id = sen_list[0]
+# sen = raw_text_sen[paper_id][sen_id]
+# doc = nlp(sen['text'])
+# print(dp_path, '\n')
+# pprint.pprint(sen)
+# print()
+# for token in doc:
+# print((token.head.text, token.text, token.dep_))
+
+ out = ''
+ for k, v_dict in draft.items():
+ input = v_dict['in']
+ output = v_dict['out']
+ if input == '':
+ continue
+ output = output.replace('\n', ' ')
+ doc = nlp(output)
+ for sen in doc.sents:
+ out += raw_to_format(sen.text) + '\n'
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_{args.mode}_parsein.txt', 'w') as fl:
+ fl.write(out)
+elif args.action == 'extract':
+
+ # dependency_to_type_id = {}
+ # for k, v in Parameters.edge_type_to_id.items():
+ # dependency_to_type_id[k] = {}
+ # for type in v:
+ # LL = list(retieve_sentence_through_edgetype[type]['manual'].keys()) + list(retieve_sentence_through_edgetype[type]['auto'].keys())
+ # for dp in LL:
+ # dependency_to_type_id[k][dp] = type
+ if os.path.exists('generate_abstract/dependency_to_type_id.pickle'):
+ with open('generate_abstract/dependency_to_type_id.pickle', 'rb') as fl:
+ dependency_to_type_id = pkl.load(fl)
+ else:
+ dependency_to_type_id = {}
+ print('Loading path data ...')
+ for k in Parameters.edge_type_to_id.keys():
+ start, end = k.split('-')
+ dependency_to_type_id[k] = {}
+ inner_edge_type_to_id = Parameters.edge_type_to_id[k]
+ inner_edge_type_dict = Parameters.edge_type_dict[k]
+ cal_manual_num = [0] * len(inner_edge_type_to_id)
+ with open('../GNBRdata/part-i-'+start+'-'+end+'-path-theme-distributions.txt', 'r') as fl:
+ for i, line in tqdm(list(enumerate(fl.readlines()))):
+ tmp = line.split('\t')
+ if i == 0:
+ head = [tmp[i] for i in range(1, len(tmp), 2)]
+ assert ' '.join(head) == ' '.join(inner_edge_type_dict[0])
+ continue
+ probability = [float(tmp[i]) for i in range(1, len(tmp), 2)]
+ flag_list = [int(tmp[i]) for i in range(2, len(tmp), 2)]
+ indices = np.where(np.asarray(flag_list) == 1)[0]
+ if len(indices) >= 1:
+ tmp_p = [cal_manual_num[i] for i in indices]
+ p = indices[np.argmin(tmp_p)]
+ cal_manual_num[p] += 1
+ else:
+ p = np.argmax(probability)
+ assert tmp[0].lower() not in dependency_to_type_id.keys()
+ dependency_to_type_id[k][tmp[0].lower()] = inner_edge_type_to_id[p]
+ with open('generate_abstract/dependency_to_type_id.pickle', 'wb') as fl:
+ pkl.dump(dependency_to_type_id, fl)
+ 
+ # record = []
+ # with open(f'generate_abstract/par_parseout.txt', 'r') as fl:
+ # Tmp = []
+ # tmp = []
+ # for i,line in enumerate(fl.readlines()):
+ # # print(len(line), line)
+ # line = line.replace('\n', '')
+ # if len(line) > 1:
+ # tmp.append(line)
+ # else:
+ # Tmp.append(tmp)
+ # tmp = []
+ # if len(Tmp) == 3:
+ # record.append(Tmp)
+ # Tmp = []
+
+ # print(len(record))
+ # record_index = 0 
+ # add = 0
+ # Attack = []
+ # for ii in range(100):
+
+ # # input = v_dict['in']
+ # # output = v_dict['out']
+ # # output = output.replace('\n', ' ')
+ # s, r, o = attack_data[ii]
+ # dependency_sen_dict = retieve_sentence_through_edgetype[int(r)]['manual']
+ 
+ # target_dp = set()
+ # for dp_path, sen_list in dependency_sen_dict.items():
+ # target_dp.add(dp_path)
+ # DP_list = []
+ # for _ in range(1):
+ # dp_dict = {}
+ # data = record[record_index]
+ # record_index += 1
+ # dp_paths = data[2]
+ # nodes_list = []
+ # edges_list = []
+ # for line in dp_paths:
+ # ttp, tmp = line.split('(')
+ # assert tmp[-1] == ')'
+ # tmp = tmp[:-1]
+ # e1, e2 = tmp.split(', ')
+ # if not ttp in type_set and ':' in ttp:
+ # ttp = ttp.split(':')[0]
+ # dp_dict[f'{e1}_x_{e2}'] = [e1, ttp, e2]
+ # dp_dict[f'{e2}_x_{e1}'] = [e1, ttp, e2]
+ # nodes_list.append(e1)
+ # nodes_list.append(e2)
+ # edges_list.append((e1, e2))
+ # nodes_list = list(set(nodes_list))
+ # pure_name = [('-'.join(name.split('-')[:-1])).replace('_', ' ') for name in nodes_list]
+ # graph = nx.Graph(edges_list)
+
+ # type_list = [name_to_type[name] if name in good_name else '' for name in pure_name]
+ # # print(type_list)
+ # # for i in range(len(type_list)):
+ # # print(pure_name[i], type_list[i])
+ # for i in range(len(nodes_list)):
+ # if type_list[i] != '':
+ # for j in range(len(nodes_list)):
+ # if i != j and type_list[j] != '':
+ # if f'{type_list[i]}-{type_list[j]}' in Parameters.edge_type_to_id.keys():
+ # # print(f'{type_list[i]}_{type_list[j]}')
+ # ret_path = []
+ # sp = nx.shortest_path(graph, source=nodes_list[i], target=nodes_list[j])
+ # start = sp[0]
+ # end = sp[-1]
+ # for k in range(len(sp)-1):
+ # e1, ttp, e2 = dp_dict[f'{sp[k]}_x_{sp[k+1]}']
+ # if e1 == start:
+ # e1 = 'start_entity-x'
+ # if e2 == start:
+ # e2 = 'start_entity-x'
+ # if e1 == end:
+ # e1 = 'end_entity-x'
+ # if e2 == end:
+ # e2 = 'end_entity-x'
+ # ret_path.append(f'{"-".join(e1.split("-")[:-1])}|{ttp}|{"-".join(e2.split("-")[:-1])}'.lower())
+ # dependency_P = ' '.join(ret_path)
+ # DP_list.append((f'{type_list[i]}-{type_list[j]}', 
+ # name_to_meshid[pure_name[i]], 
+ # name_to_meshid[pure_name[j]], 
+ # dependency_P))
+ 
+ # boo = False
+ # modified_attack = []
+ # for k, ss, tt, dp in DP_list:
+ # if dp in dependency_to_type_id[k].keys():
+ # tp = str(dependency_to_type_id[k][dp])
+ # id_ss = str(meshid_to_id[ss])
+ # id_tt = str(meshid_to_id[tt])
+ # modified_attack.append(f'{id_ss}*{tp}*{id_tt}')
+ # if int(dependency_to_type_id[k][dp]) == int(r):
+ # # if id_to_meshid[s] == ss and id_to_meshid[o] == tt:
+ # boo = True
+ # modified_attack = list(set(modified_attack))
+ # modified_attack = [k.split('*') for k in modified_attack]
+ # if boo:
+ # add += 1
+ # # else:
+ # # print(ii)
+ 
+ # # for i in range(len(type_list)):
+ # # if type_list[i]:
+ # # print(pure_name[i], type_list[i])
+ # # for k, ss, tt, dp in DP_list:
+ # # print(k, dp)
+ # # print(record[record_index - 1])
+ # # raise Exception('No!!')
+ # Attack.append(modified_attack)
+
+ record = []
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_{args.mode}_parseout.txt', 'r') as fl:
+ Tmp = []
+ tmp = []
+ for i,line in enumerate(fl.readlines()):
+ # print(len(line), line)
+ line = line.replace('\n', '')
+ if len(line) > 1:
+ tmp.append(line)
+ else:
+ if len(Tmp) == 2:
+ if len(tmp) == 1 and '/' in tmp[0].split(' ')[0]:
+ Tmp.append([])
+ record.append(Tmp)
+ Tmp = []
+ Tmp.append(tmp)
+ if len(Tmp) == 2 and tmp[0][:5] != '(ROOT':
+ print(record[-1][2])
+ raise Exception('??')
+ tmp = []
+ if len(Tmp) == 3:
+ record.append(Tmp)
+ Tmp = []
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_{args.mode}_parsein.txt', 'r') as fl:
+ parsin = fl.readlines()
+
+ print('Record len', len(record), 'Parsin len:', len(parsin))
+ record_index = 0 
+ add = 0
+
+ Attack = []
+ for ii, (k, v_dict) in enumerate(tqdm(draft.items())):
+
+ input = v_dict['in']
+ output = v_dict['out']
+ output = output.replace('\n', ' ')
+ s, r, o = attack_data[ii]
+ assert ii == int(k.split('_')[-1])
+ 
+ DP_list = []
+ if input != '':
+
+ dependency_sen_dict = retieve_sentence_through_edgetype[int(r)]['manual']
+ target_dp = set()
+ for dp_path, sen_list in dependency_sen_dict.items():
+ target_dp.add(dp_path)
+ doc = nlp(output)
+ 
+ for sen in doc.sents:
+ dp_dict = {}
+ if record_index >= len(record):
+ break
+ data = record[record_index]
+ record_index += 1
+ dp_paths = data[2]
+ nodes_list = []
+ edges_list = []
+ for line in dp_paths:
+ aa = line.split('(')
+ if len(aa) == 1:
+ print(ii)
+ print(sen)
+ print(data)
+ raise Exception
+ ttp, tmp = aa[0], aa[1]
+ assert tmp[-1] == ')'
+ tmp = tmp[:-1]
+ e1, e2 = tmp.split(', ')
+ if not ttp in type_set and ':' in ttp:
+ ttp = ttp.split(':')[0]
+ dp_dict[f'{e1}_x_{e2}'] = [e1, ttp, e2]
+ dp_dict[f'{e2}_x_{e1}'] = [e1, ttp, e2]
+ nodes_list.append(e1)
+ nodes_list.append(e2)
+ edges_list.append((e1, e2))
+ nodes_list = list(set(nodes_list))
+ pure_name = [('-'.join(name.split('-')[:-1])).replace('_', ' ') for name in nodes_list]
+ graph = nx.Graph(edges_list)
+
+ type_list = [name_to_type[name] if name in good_name else '' for name in pure_name]
+ # print(type_list)
+ for i in range(len(nodes_list)):
+ if type_list[i] != '':
+ for j in range(len(nodes_list)):
+ if i != j and type_list[j] != '':
+ if f'{type_list[i]}-{type_list[j]}' in Parameters.edge_type_to_id.keys():
+ # print(f'{type_list[i]}_{type_list[j]}')
+ ret_path = []
+ sp = nx.shortest_path(graph, source=nodes_list[i], target=nodes_list[j])
+ start = sp[0]
+ end = sp[-1]
+ for k in range(len(sp)-1):
+ e1, ttp, e2 = dp_dict[f'{sp[k]}_x_{sp[k+1]}']
+ if e1 == start:
+ e1 = 'start_entity-x'
+ if e2 == start:
+ e2 = 'start_entity-x'
+ if e1 == end:
+ e1 = 'end_entity-x'
+ if e2 == end:
+ e2 = 'end_entity-x'
+ ret_path.append(f'{"-".join(e1.split("-")[:-1])}|{ttp}|{"-".join(e2.split("-")[:-1])}'.lower())
+ dependency_P = ' '.join(ret_path)
+ DP_list.append((f'{type_list[i]}-{type_list[j]}', 
+ name_to_meshid[pure_name[i]], 
+ name_to_meshid[pure_name[j]], 
+ dependency_P))
+ 
+ boo = False
+ modified_attack = []
+ for k, ss, tt, dp in DP_list:
+ if dp in dependency_to_type_id[k].keys():
+ tp = str(dependency_to_type_id[k][dp])
+ id_ss = str(meshid_to_id[ss])
+ id_tt = str(meshid_to_id[tt])
+ modified_attack.append(f'{id_ss}*{tp}*{id_tt}')
+ if int(dependency_to_type_id[k][dp]) == int(r):
+ if id_to_meshid[s] == ss and id_to_meshid[o] == tt:
+ boo = True
+ modified_attack = list(set(modified_attack))
+ modified_attack = [k.split('*') for k in modified_attack]
+ if boo:
+ # print(DP_list)
+ add += 1
+ Attack.append(modified_attack)
+ print(add)
+ print('End record_index:', record_index)
+ with open(modified_attack_path, 'wb') as fl:
+ pkl.dump(Attack, fl)
+else:
+ raise Exception('Wrong action !!')

DiseaseSpecific/attack.py

@@ -0,0 +1,841 @@
+#%%
+import pickle as pkl
+from typing import Dict, Tuple, List
+import os
+import numpy as np
+import json
+import dill
+import logging
+import argparse 
+import math
+from pprint import pprint
+import pandas as pd
+from collections import defaultdict
+import copy
+import time
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+import torch.autograd as autograd
+from torch.nn import functional as F
+from torch.nn.modules.loss import CrossEntropyLoss
+
+from model import Distmult, Complex, Conve
+import utils
+
+import sys
+
+import dill
+
+sys.path.append("..")
+import Parameters
+
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+
+logger = None
+def generate_nghbrs_single_entity(x, edge_nghbrs, bound):
+
+ ret_S = set(x)
+ ret_L = [x]
+ b = 0 
+ while(b < len(ret_L)):
+ s = ret_L[b]
+ if s in edge_nghbrs.keys():
+ for v in edge_nghbrs[s]:
+ if v not in ret_S:
+ ret_S.add(v)
+ ret_L.append(v)
+ if len(ret_L) == bound:
+ return ret_L
+ b += 1
+ return ret_L
+
+def generate_nghbrs(target_data, edge_nghbrs, args):
+ n_dict = {}
+ for i, (s, r, o) in enumerate(target_data):
+ L_s = generate_nghbrs_single_entity(s, edge_nghbrs, args.neighbor_num)
+ L_o = generate_nghbrs_single_entity(o, edge_nghbrs, args.neighbor_num)
+ n_dict[i] = list(set(L_s + L_o))
+ n_dict[i].sort()
+ return n_dict 
+#%%
+def check_edge(s, r, o, used_trip = None, args = None):
+ """Double check"""
+ if args is None:
+ return True
+ if not args.target_existed:
+ assert (s+'_'+o in used_trip) == args.target_existed
+ else:
+ s = entityid_to_nodetype[s]
+ o = entityid_to_nodetype[o]
+ r_tp = Parameters.edge_id_to_type[int(r)]
+ r_tp = r_tp.split(':')[0]
+ r_tp = r_tp.split('-')
+ assert s == r_tp[0] and o == r_tp[1]
+
+def get_model_loss(batch, model, device, args = None):
+ s,r,o = batch[:,0], batch[:,1], batch[:,2]
+
+ emb_s = model.emb_e(s).squeeze(dim=1)
+ emb_r = model.emb_rel(r).squeeze(dim=1)
+ emb_o = model.emb_e(o).squeeze(dim=1)
+
+ if args.add_reciprocals:
+ r_rev = r + n_rel
+ emb_rrev = model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+
+ pred_sr = model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = model.loss(pred_sr, o) # Cross entropy loss
+
+ pred_or = model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = model.loss(pred_or, s)
+
+ train_loss = loss_sr + loss_or
+ return train_loss
+
+def get_model_loss_without_softmax(batch, model, device=None):
+
+ with torch.no_grad():
+ s,r,o = batch[:,0], batch[:,1], batch[:,2]
+
+ emb_s = model.emb_e(s).squeeze(dim=1)
+ emb_r = model.emb_rel(r).squeeze(dim=1)
+
+ pred = model.forward(emb_s, emb_r)
+ return -pred[range(o.shape[0]), o]
+
+def lp_regularizer(model, weight, p):
+ trainable_params = [model.emb_e.weight, model.emb_rel.weight]
+ norm = 0
+ for i in range(len(trainable_params)):
+ norm += weight * torch.sum( torch.abs(trainable_params[i]) ** p)
+ return norm
+
+def n3_regularizer(factors, weight, p):
+ norm = 0
+ for f in factors:
+ norm += weight * torch.sum(torch.abs(f) ** p)
+ return norm / factors[0].shape[0] 
+
+def get_train_loss(batch, model, device, args):
+ #batch = batch[0].to(device)
+ s,r,o = batch[:,0], batch[:,1], batch[:,2]
+
+ emb_s = model.emb_e(s).squeeze(dim=1)
+ emb_r = model.emb_rel(r).squeeze(dim=1)
+ emb_o = model.emb_e(o).squeeze(dim=1)
+
+ if args.add_reciprocals:
+ r_rev = r + n_rel
+ emb_rrev = model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+
+ pred_sr = model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = model.loss(pred_sr, o) # loss is cross entropy loss
+
+ pred_or = model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = model.loss(pred_or, s)
+
+ train_loss = loss_sr + loss_or
+ 
+ if (args.reg_weight != 0.0 and args.reg_norm == 3):
+ #self.logger.info('Computing regularizer weight')
+ if model == 'complex':
+ emb_dim = args.embedding_dim #int(self.args.embedding_dim/2)
+ lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
+ rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
+ rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
+ rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])
+
+ #print(lhs[0].shape, lhs[1].shape)
+ factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ else:
+ factors_sr = (emb_s, emb_r, emb_o)
+ factors_or = (emb_s, emb_rrev, emb_o)
+
+ train_loss += n3_regularizer(factors_sr, args.reg_weight, p=3)
+ train_loss += n3_regularizer(factors_or, args.reg_weight, p=3)
+
+ if (args.reg_weight != 0.0 and args.reg_norm == 2):
+ train_loss += lp_regularizer(model, args.reg_weight, p=2)
+ 
+ return train_loss
+def hv(loss, model_params, v):
+ grad = autograd.grad(loss, model_params, create_graph=True, retain_graph=True)
+ Hv = autograd.grad(grad, model_params, grad_outputs=v)
+ return Hv
+def gather_flat_grad(grads):
+ views = []
+ for p in grads:
+ if p.data.is_sparse:
+ view = p.data.to_dense().view(-1)
+ else:
+ view = p.data.view(-1)
+ views.append(view)
+ return torch.cat(views, 0)
+
+def get_inverse_hvp_lissa(v, model, device, param_influence, train_data, args):
+
+ damping = args.damping
+ num_samples = args.lissa_repeat
+ scale = args.scale 
+ train_batch_size = args.lissa_batch_size
+ lissa_num_batches = math.ceil(train_data.shape[0]/train_batch_size)
+ recursion_depth = int(lissa_num_batches*args.lissa_depth)
+
+ ihvp = None
+ # print('inversing hvp...')
+ for i in range(num_samples):
+ cur_estimate = v
+ #lissa_data_iterator = iter(train_loader)
+ input_data = torch.from_numpy(train_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+ 
+ b_begin = 0
+ for j in range(recursion_depth):
+ model.zero_grad() # same as optimizer.zero_grad()
+ if b_begin >= actual_examples.shape[0]:
+ b_begin = 0
+ input_data = torch.from_numpy(train_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+ 
+ input_batch = actual_examples[b_begin: b_begin + train_batch_size]
+ input_batch = input_batch.to(device)
+ 
+ train_loss = get_train_loss(input_batch, model, device, args)
+ 
+ hvp = hv(train_loss, param_influence, cur_estimate)
+ cur_estimate = [_a + (1-damping)*_b - _c / scale for _a, _b, _c in zip(v, cur_estimate, hvp)]
+ # if (j%200 == 0) or (j == recursion_depth -1 ):
+ # logger.info("Recursion at depth %s: norm is %f" % (j, np.linalg.norm(gather_flat_grad(cur_estimate).cpu().numpy())))
+ 
+ b_begin += train_batch_size
+ 
+ if ihvp == None:
+ ihvp = [_a / scale for _a in cur_estimate]
+ else:
+ ihvp = [_a + _b / scale for _a, _b in zip(ihvp, cur_estimate)]
+
+ # logger.info("Final ihvp norm is %f" % (np.linalg.norm(gather_flat_grad(ihvp).cpu().numpy())))
+ return_ihvp = gather_flat_grad(ihvp)
+ return_ihvp /= num_samples
+ 
+ return return_ihvp
+
+#%%
+def before_global_attack(device, n_rel, data, target_data, neighbors, model, 
+ filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
+ entityid_to_nodetype, batch_size, args, lissa_path, target_disease):
+
+ if os.path.exists(lissa_path) and not args.update_lissa:
+ with open(lissa_path, 'rb') as fl:
+ ret = dill.load(fl)
+ return ret
+ ret = {}
+
+ test_data = []
+ for i in target_disease:
+ tp = entityid_to_nodetype[str(i)]
+ # r = torch.LongTensor([[10]]).to(device)
+ assert tp == 'disease'
+ if tp == 'disease':
+ for target in target_data:
+ test_data.append([str(target), str(10), str(i)])
+ test_data = np.array(test_data)
+
+ for target_trip in tqdm(test_data):
+
+ target_trip_ori = target_trip
+ trip_name = '_'.join(list(target_trip_ori))
+ target_trip = target_trip[None, :] # add a batch dimension
+ target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)
+ # target_s, target_r, target_o = target_trip[:,0], target_trip[:,1], target_trip[:,2]
+ # target_vec = model.score_triples_vec(target_s, target_r, target_o)
+
+ model.eval()
+ model.zero_grad()
+ target_loss = get_model_loss(target_trip, model, device)
+ target_grads = autograd.grad(target_loss, param_influence)
+
+ model.train()
+ inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
+ param_influence, data, args)
+ model.eval()
+ inverse_hvp = inverse_hvp.detach().cpu().unsqueeze(0)
+ ret[trip_name] = inverse_hvp
+ with open(lissa_path, 'wb') as fl:
+ dill.dump(ret, fl)
+ return ret
+ 
+def global_addtion_attack(device, n_rel, data, target_data, neighbors, model, 
+ filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
+ entityid_to_nodetype, batch_size, args, lissa, target_disease):
+
+ logger.info('------ Generating edits per target triple ------')
+ start_time = time.time()
+ logger.info('Start time: {0}'.format(str(start_time)))
+
+ used_trip = set()
+ print("Processing used triples ...")
+ for s, r, o in tqdm(data):
+ used_trip.add(s+'_'+o)
+ # used_trip.add(o+'_'+s)
+ print('Size of used triples:', len(used_trip))
+ logger.info('Size of used triples: {0}'.format(len(used_trip)))
+
+ ret_trip = []
+ score_record = []
+ real_add_rank_ratio = 0
+
+ with open(score_path, 'rb') as fl:
+ score_record = pkl.load(fl)
+ for i, target in enumerate(target_data):
+
+ print('\n\n------ Attacking target tripid:', i, 'tot:', len(target_data), ' ------')
+ # lissa_hvp = []
+ target_trip = []
+ for disease in target_disease:
+ target_trip.append([target, str(10), disease])
+ # nm = '{}_{}_{}'.format(target, 10, disease)
+ # lissa_hvp.append(lissa[nm])
+ # lissa_hvp = torch.cat(lissa_hvp, dim = 0).to(device)
+
+ target_trip = np.array(target_trip)
+ target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)
+
+ model.eval()
+ model.zero_grad()
+ target_loss = get_model_loss(target_trip, model, device)
+ target_grads = autograd.grad(target_loss, param_influence)
+
+ model.train()
+ inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
+ param_influence, data, args)
+
+ model.eval()
+
+ nghbr_trip = []
+ s = str(target)
+ tp = entityid_to_nodetype[s]
+ for nghbr in tqdm(neighbors):
+ o = str(nghbr)
+ if s!=o and s+'_'+o not in used_trip:
+ for r in range(n_rel):
+ if (tp, r) in filters["rhs"].keys() and filters["rhs"][(tp, r)][int(o)] == True:
+ nghbr_trip.append([s, str(r), o])
+
+ nghbr_trip = np.asarray(nghbr_trip)
+ influences = [] 
+ edge_losses = []
+ 
+ # nghbr_cos_log_prob, nghbr_LM_log_prob = score_record[i]
+ # assert nghbr_cos_log_prob.shape[0] == nghbr_trip.shape[0]
+ 
+ for train_trip in tqdm(nghbr_trip):
+ #model.train() #batch norm cannot be used here
+ train_trip = train_trip[None, :] # add batch dim
+ train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+ #### L-train gradient ####
+ edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+ edge_losses.append(edge_loss.unsqueeze(0).detach())
+ model.zero_grad()
+ train_loss = get_model_loss(train_trip, model, device, args)
+ train_grads = autograd.grad(train_loss, param_influence)
+ train_grads = gather_flat_grad(train_grads)
+ influence = torch.dot(inverse_hvp, train_grads) #default dim=1
+ influences.append(influence.unsqueeze(0).detach())
+
+ edge_losses = torch.cat(edge_losses, dim = -1)
+ influences = torch.cat(influences, dim = -1)
+ edge_losses_log_prob = torch.log(F.softmax(-edge_losses, dim = -1))
+ influences_log_prob = torch.log(F.softmax(influences, dim = -1))
+
+ inf_score_sorted, influences_sort = torch.sort(influences_log_prob, -1, descending=True)
+ edge_score_sorted, edge_sort = torch.sort(edge_losses_log_prob, -1, descending=True)
+ influences_sort = influences_sort.cpu().numpy()
+ edge_sort = edge_sort.cpu().numpy()
+ inf_score_sorted = inf_score_sorted.cpu().numpy()
+ edge_score_sorted = edge_score_sorted.cpu().numpy()
+
+ logger.info('')
+ logger.info('Top 8 inf_score: {}'.format(" ".join(map(str, list(inf_score_sorted[:8])))))
+ logger.info('Top 8 edge_score: {}'.format(" ".join(map(str, list(edge_score_sorted[:8])))))
+
+ nghbr_cos_log_prob = influences_log_prob.detach().cpu().numpy()
+ nghbr_LM_log_prob = edge_losses_log_prob.detach().cpu().numpy()
+ max_sim = np.max(nghbr_cos_log_prob)
+ min_sim = np.min(nghbr_cos_log_prob)
+ max_LM = np.max(nghbr_LM_log_prob)
+ min_LM = np.min(nghbr_LM_log_prob)
+
+ # final_score = nghbr_cos_log_prob + nghbr_LM_log_prob
+ final_score = nghbr_cos_log_prob
+
+ index = np.argmax(final_score[:-1])
+ # p = np.where(index == edge_sort)[0][0]
+ # logger.info('Added edge\'s edge rank ratio: {}'.format(p / edge_sort.shape[0]))
+ real_add_rank_ratio += p
+ add_trip = nghbr_trip[index]
+ logger.info('max_inf: {0:.8}, min_inf: {1:.8}, max_edge: {2:.8}, min_edge: {3:.8}'.format(max_sim, min_sim, max_LM, min_LM))
+ logger.info('Attack trip: {0}_{1}_{2}.\n Influnce score: {3:.8}. Edge score: {4:.8}.'.format(add_trip[0], add_trip[1], add_trip[2],
+ nghbr_cos_log_prob[index], nghbr_LM_log_prob[index]))
+ ret_trip.append(add_trip) 
+ score_record.append((nghbr_cos_log_prob, nghbr_LM_log_prob)) 
+ real_add_rank_ratio = real_add_rank_ratio / target_data.shape[0]
+ logger.info('Mean real ratio: {}.'.format(real_add_rank_ratio))
+ return ret_trip, score_record
+
+def addition_attack(param_influence, device, n_rel, data, target_data, neighbors, model, 
+ filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
+ entityid_to_nodetype, batch_size, args, load_Record = False, divide_bound = None, data_mean = None, data_std = None, cache_intermidiate = True):
+
+ if logger:
+ logger.info('------ Generating edits per target triple ------')
+ start_time = time.time()
+ if logger:
+ logger.info('Start time: {0}'.format(str(start_time)))
+
+ used_trip = set()
+ print("Processing used triples ...")
+ for s, r, o in tqdm(data):
+ used_trip.add(s+'_'+o)
+ # used_trip.add(o+'_'+s)
+ print('Size of used triples:', len(used_trip))
+ if logger:
+ logger.info('Size of used triples: {0}'.format(len(used_trip)))
+
+ nghbr_trip_len = []
+ ret_trip = []
+ score_record = []
+ direct_add_rank_ratio = 0
+ real_add_rank_ratio = 0
+ bad_ratio = 0
+
+ RRcord = []
+ print('****'*10)
+ if load_Record:
+ print('Load intermidiate file')
+ with open(intermidiate_path, 'rb') as fl:
+ RRcord = dill.load(fl)
+ else:
+ print('Donnot load intermidiate file')
+
+ for i, target_trip in enumerate(target_data):
+
+ print('\n\n------ Attacking target tripid:', i, ' ------')
+ target_nghbrs = neighbors[i]
+ for a in target_nghbrs:
+ if str(a) == '-1':
+ raise Exception('pppp')
+
+ target_trip_ori = target_trip
+ check_edge(target_trip[0], target_trip[1], target_trip[2], used_trip)
+ target_trip = target_trip[None, :] # add a batch dimension
+ target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)
+ # target_s, target_r, target_o = target_trip[:,0], target_trip[:,1], target_trip[:,2]
+ # target_vec = model.score_triples_vec(target_s, target_r, target_o)
+
+ model.eval()
+
+ if load_Record:
+ o_target_trip, nghbr_trip, edge_losses, influences, edge_losses_log_prob, influences_log_prob = RRcord[i]
+ assert (o_target_trip.cpu() == target_trip.cpu()).sum().item() == 3
+ else:
+ model.zero_grad()
+ target_loss = get_model_loss(target_trip, model, device, args)
+ target_grads = autograd.grad(target_loss, param_influence)
+
+ model.train()
+ inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
+ param_influence, data, args)
+
+ model.eval()
+ nghbr_trip = []
+ valid_trip = 0
+ if args.candidate_mode == 'quadratic':
+ s_o_list = [(i, j) for i in target_nghbrs for j in target_nghbrs]
+ elif args.candidate_mode == 'linear':
+ s_o_list = [(j, i) for i in target_nghbrs for j in [target_trip_ori[0], target_trip_ori[2]]] \
+ + [(i, j) for i in target_nghbrs for j in [target_trip_ori[0], target_trip_ori[2]]]
+ else:
+ raise Exception('Wrong candidate_mode: '+args.candidate_mode)
+ for s, o in tqdm(s_o_list): 
+ tp = entityid_to_nodetype[s]
+ if s!=o and s+'_'+o not in used_trip:
+ for r in range(n_rel):
+ if (tp, r) in filters["rhs"].keys() and filters["rhs"][(tp, r)][int(o)] == True:
+ # check_edge(s, r, o)
+ valid_trip += 1
+ nghbr_trip.append([s, str(r), o])
+ # logger.info('{0}_{1}_{2}'.format(s, str(r), o))
+ nghbr_trip_len.append(len(nghbr_trip))
+ print('Valid trip:', valid_trip)
+
+ if target_trip_ori[0]+'_'+target_trip_ori[2] not in used_trip:
+ nghbr_trip.append(target_trip_ori)
+ nghbr_trip = np.asarray(nghbr_trip)
+ print("Edge scoring ...")
+
+ influences = []
+ edge_losses = []
+
+ for train_trip in tqdm(nghbr_trip):
+ #model.train() #batch norm cannot be used here
+ train_trip = train_trip[None, :] # add batch dim
+ train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+ #### L-train gradient ####
+ edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+ edge_losses.append(edge_loss.unsqueeze(0).detach())
+ model.zero_grad()
+ train_loss = get_model_loss(train_trip, model, device, args)
+ train_grads = autograd.grad(train_loss, param_influence)
+ train_grads = gather_flat_grad(train_grads)
+ influence = torch.dot(inverse_hvp, train_grads) #default dim=1
+ influences.append(influence.unsqueeze(0).detach()) 
+
+ edge_losses = torch.cat(edge_losses, dim = -1)
+ influences = torch.cat(influences, dim = -1)
+ edge_losses_log_prob = torch.log(F.softmax(-edge_losses, dim = -1))
+ influences_log_prob = torch.log(F.softmax(influences, dim = -1))
+ std_scale = torch.std(edge_losses_log_prob) / torch.std(influences_log_prob)
+ influences_log_prob = (influences_log_prob - influences_log_prob.mean()) * std_scale + edge_losses_log_prob.mean()
+ 
+ RRcord.append([target_trip.detach(), nghbr_trip, edge_losses, influences, edge_losses_log_prob, influences_log_prob])
+
+ inf_score_sorted, influences_sort = torch.sort(influences_log_prob, -1, descending=True)
+ edge_score_sorted, edge_sort = torch.sort(edge_losses_log_prob, -1, descending=True)
+
+ influences_sort = influences_sort.cpu().numpy()
+ edge_sort = edge_sort.cpu().numpy()
+ inf_score_sorted = inf_score_sorted.cpu().numpy()
+ edge_score_sorted = edge_score_sorted.cpu().numpy()
+ edge_losses = edge_losses.cpu().numpy()
+
+ p = np.where(influences_sort[0] == edge_sort)[0][0]
+ direct_add_rank_ratio += p / edge_sort.shape[0]
+ if logger:
+ logger.info('Top 8 inf_score: {}'.format(" ".join(map(str, list(inf_score_sorted[:8])))))
+ logger.info('Top 8 edge_score: {}'.format(" ".join(map(str, list(edge_score_sorted[:8])))))
+
+ nghbr_cos_log_prob = influences_log_prob.detach().cpu().numpy()
+ nghbr_LM_log_prob = edge_losses_log_prob.detach().cpu().numpy()
+ max_sim = nghbr_cos_log_prob[influences_sort[0]]
+ min_sim = nghbr_cos_log_prob[influences_sort[-1]]
+ max_LM = nghbr_LM_log_prob[edge_sort[0]]
+ min_LM = nghbr_LM_log_prob[edge_sort[-1]]
+ direct_score_0 = 0
+ direct_score_1 = 0
+ if target_trip_ori[0]+'_'+target_trip_ori[2] not in used_trip:
+ direct_score_0 = nghbr_cos_log_prob[-1]
+ direct_score_1 = nghbr_LM_log_prob[-1]
+ 
+ # bound = math.log(1 / nghbr_LM_log_prob.shape[0])
+ bound = 1 - args.reasonable_rate
+ edge_losses = (edge_losses - data_mean) / data_std
+ edge_losses_prob = 1 / ( 1 + np.exp(edge_losses - divide_bound) )
+ nghbr_LM_log_prob[edge_losses_prob < bound] = -(1e20)
+
+ final_score = nghbr_cos_log_prob + nghbr_LM_log_prob
+
+ index = np.argmax(final_score[:-1])
+ sort_index = [(i, final_score[i])for i in range(len(final_score) - 1)]
+ sort_index = sorted(sort_index, key=lambda x: x[1], reverse=True)
+ assert sort_index[0][0] == index
+
+ p = np.where(index == edge_sort)[0][0]
+ if logger:
+ logger.info('Bad edge ratio: {}'.format((edge_losses_prob < bound).mean()))
+ logger.info('Bounded edge\'s edge rank ratio: {}'.format(p / edge_sort.shape[0]))
+ real_add_rank_ratio += p / edge_sort.shape[0]
+ bad_ratio += (edge_losses_prob < bound).mean()
+
+ add_trip = nghbr_trip[index]
+
+ if (int(add_trip[0]) == int(-1)):
+ add_trip[0], add_trip[1], add_trip[2] = -1, -1, -1
+ print(final_score.shape, index, edge_losses_prob[index], bound)
+ raise Exception('??')
+
+ if logger:
+ logger.info('max_inf: {0:.8}, min_inf: {1:.8}, max_edge: {2:.8}, min_edge: {3:.8}'.format(max_sim, min_sim, max_LM, min_LM))
+ logger.info('Target trip: {0}_{1}_{2}. Attack trip: {3}_{4}_{5}.\n Influnce score: {6:.8}. Edge score: {7:.8}. Direct score: {8:.8} + {9:.8}'.format(target_trip_ori[0],target_trip_ori[1], target_trip_ori[2], 
+ add_trip[0], add_trip[1], add_trip[2],
+ nghbr_cos_log_prob[index], nghbr_LM_log_prob[index],
+ direct_score_0, direct_score_1))
+ if (args.added_edge_num == '' or int(args.added_edge_num) == 1):
+ ret_trip.append(add_trip)
+ else:
+ edge_num = int(args.added_edge_num)
+ for i in range(edge_num):
+ ret_trip.append(nghbr_trip[sort_index[i][0]])
+ score_record.append((nghbr_cos_log_prob, nghbr_LM_log_prob)) 
+ 
+ if not load_Record and cache_intermidiate:
+ with open(intermidiate_path, 'wb') as fl:
+ dill.dump(RRcord, fl)
+ direct_add_rank_ratio = direct_add_rank_ratio / target_data.shape[0]
+ real_add_rank_ratio = real_add_rank_ratio / target_data.shape[0]
+ bad_ratio = bad_ratio / target_data.shape[0]
+ if logger:
+ logger.info('Mean direct ratio: {}. Mean real ratio: {}. Mean bad ratio: {}'.format(direct_add_rank_ratio, real_add_rank_ratio, bad_ratio))
+ return ret_trip, score_record
+
+def calculate_edge_bound(data, model, device, n_ent):
+
+ tmp = np.random.choice(a = data.shape[0], size = data.shape[0] // 10, replace=False)
+ existed_data= data[tmp, :]
+
+ print('calculating edge bound ...')
+ print(existed_data.shape)
+
+ existed_edge = set()
+ for src_trip in existed_data:
+ existed_edge.add('_'.join(list(src_trip)))
+ 
+ not_existed = []
+ for s, r, o in existed_data:
+
+ if np.random.randint(0, n_ent) % 2 == 0:
+ while True:
+ oo = np.random.randint(0, n_ent)
+ if '_'.join([s, r, str(oo)]) not in existed_edge:
+ not_existed.append([s, r, str(oo)])
+ break
+ else:
+ while True:
+ ss = np.random.randint(0, n_ent)
+ if '_'.join([str(ss), r, o]) not in existed_edge:
+ not_existed.append([str(ss), r, o])
+ break 
+ existed_data = np.array(existed_data)
+ not_existed = np.array(not_existed)
+ existed_data = torch.from_numpy(existed_data.astype('int64')).to(device)
+ not_existed = torch.from_numpy(not_existed.astype('int64')).to(device)
+ loss_existed = get_model_loss_without_softmax(existed_data, model).cpu().numpy()
+ loss_not_existed = get_model_loss_without_softmax(not_existed, model).cpu().numpy()
+ tot_loss = np.hstack((loss_existed, loss_not_existed))
+ tot_mean, tot_std = np.mean(tot_loss), np.std(tot_loss)
+
+ loss_existed = (loss_existed - tot_mean) / tot_std
+ loss_not_existed = (loss_not_existed - tot_mean) / tot_std
+
+ print('Tot mean: {}, Tot std: {}'.format(tot_mean, tot_std))
+
+ # print(np.mean(loss_existed), np.std(loss_existed), np.max(loss_existed))
+ # print(np.mean(loss_not_existed), np.std(loss_not_existed), np.min(loss_not_existed))
+ l_mean, l_std = np.mean(loss_existed), np.std(loss_existed)
+ r_mean, r_std = np.mean(loss_not_existed), np.std(loss_not_existed)
+
+ A = -1/(l_std**2) + 1/(r_std**2)
+ B = 2 * (-r_mean/(r_std**2) + l_mean/(l_std**2))
+ C = (r_mean**2)/(r_std**2)-(l_mean**2)/(l_std**2) + np.log((r_std**2)/(l_std**2))
+
+ delta = B**2 - 4*A*C
+
+ x_1 = ( -B + math.sqrt(delta) ) / (2*A)
+ x_2 = ( -B - math.sqrt(delta) ) / (2*A)
+
+ x = None
+ if (x_1 > l_mean and x_1 < r_mean):
+ x = x_1
+ if (x_2 > l_mean and x_2 < r_mean):
+ x = x_2
+ if not x:
+ raise Exception('Bad model!!!!')
+ TP = (loss_existed < x).mean()
+ TN = (loss_not_existed > x).mean()
+ FP = (loss_not_existed < x).mean()
+ FN = (loss_existed > x).mean()
+ print('X:{}, TP:{}, TN:{}, FP:{}, FN{}'.format(x, TP, TN, FP, FN))
+
+ sig_existed = 1 / ( 1 + np.exp(loss_existed- x) ) # negtive important
+ sig_not_existed = 1 / ( 1 + np.exp(loss_not_existed - x) )
+
+ print('Positive mean score:', sig_existed.mean(),'Negetive mean score:', sig_not_existed.mean())
+
+ return x, tot_mean, tot_std
+
+
+#%%
+if __name__ == '__main__':
+ parser = utils.get_argument_parser()
+ parser = utils.add_attack_parameters(parser)
+ args = parser.parse_args()
+ args = utils.set_hyperparams(args)
+
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ args.device = device
+ args.device1 = device
+ if torch.cuda.device_count() >= 2:
+ args.device = "cuda:0"
+ args.device1 = "cuda:1"
+
+ utils.seed_all(args.seed)
+ np.set_printoptions(precision=5)
+ cudnn.benchmark = False
+
+ model_name = '{0}_{1}_{2}_{3}_{4}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop)
+ model_path = 'saved_models/{0}_{1}.model'.format(args.data, model_name)
+ data_path = os.path.join('processed_data', args.data)
+ target_path = os.path.join(data_path, 'DD_target_{0}_{1}_{2}_{3}_{4}_{5}.txt'.format(args.model, args.data, args.target_split, args.target_size, 'exists:'+str(args.target_existed), args.attack_goal))
+ lissa_path = 'lissa/{0}_{1}_{2}'.format(args.model, 
+ args.data, 
+ args.target_size)
+ intermidiate_path = 'intermidiate/{0}_{1}_{2}_{3}_{4}_{5}_{6}'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal)
+ log_path = 'logs/attack_logs/cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate))
+ print(log_path)
+ attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ str(args.added_edge_num)))
+
+ logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+ datefmt = '%m/%d/%Y %H:%M:%S',
+ level = logging.INFO,
+ filename = log_path
+ )
+ logger = logging.getLogger(__name__)
+ logger.info(vars(args))
+ #%%
+ n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+ data = utils.load_data(os.path.join(data_path, 'all.txt'))
+ with open(os.path.join(data_path, 'filter.pickle'), 'rb') as fl:
+ filters = pkl.load(fl)
+ with open(os.path.join(data_path, 'entityid_to_nodetype.json'), 'r') as fl:
+ entityid_to_nodetype = json.load(fl)
+ with open(os.path.join(data_path, 'edge_nghbrs.pickle'), 'rb') as fl:
+ edge_nghbrs = pkl.load(fl)
+ with open(os.path.join(data_path, 'disease_meshid.pickle'), 'rb') as fl:
+ disease_meshid = pkl.load(fl)
+ with open(os.path.join(data_path, 'entities_dict.json'), 'r') as fl:
+ entity_to_id = json.load(fl)
+ with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+ entity_raw_name = pkl.load(fl)
+ #%%
+ init_mask = np.asarray([0] * n_ent).astype('int64')
+ init_mask = (init_mask == 1)
+ for k, v in filters.items():
+ for kk, vv in v.items():
+ tmp = init_mask.copy()
+ tmp[np.asarray(vv)] = True
+ t = torch.ByteTensor(tmp).to(args.device)
+ filters[k][kk] = t
+ #%%
+ model = utils.load_model(model_path, args, n_ent, n_rel, args.device)
+ divide_bound, data_mean, data_std = calculate_edge_bound(data, model, args.device, n_ent)
+ # index = torch.LongTensor([0, 1]).to(device)
+ # print(model.emb_rel(index)[:, :32])
+ # print(model.emb_e(index)[:, :32])
+ # raise Exception
+ #%%
+ target_data = utils.load_data(target_path)
+ if args.attack_goal == 'single':
+ neighbors = generate_nghbrs(target_data, edge_nghbrs, args)
+ elif args.attack_goal == 'global':
+ s_set = set()
+ for s, r, o in target_data:
+ s_set.add(s)
+ target_data = list(s_set)
+ target_data.sort()
+ target_data = np.array(target_data, dtype=str)
+ neighbors = []
+ for i in list(range(n_ent)):
+ tp = entityid_to_nodetype[str(i)]
+ # r = torch.LongTensor([[10]]).to(device)
+ if tp == 'gene':
+ neighbors.append(str(i))
+ target_disease = []
+ tid = 1
+ bound = 50
+ while True:
+ meshid = disease_meshid[tid][0]
+ fre = disease_meshid[tid][1]
+ if len(entity_raw_name[meshid]) > 4:
+ target_disease.append(entity_to_id[meshid])
+ bound -= 1
+ if bound == 0:
+ break
+ tid += 1
+ else:
+ raise Exception('Wrong attack_goal: '+args.attack_goal)
+
+ param_optimizer = list(model.named_parameters())
+ param_influence = []
+ for n,p in param_optimizer:
+ param_influence.append(p)
+ if args.attack_goal == 'single':
+ len_list = []
+ for v in neighbors.values():
+ len_list.append(len(v))
+ mean_len = np.mean(len_list)
+ else:
+ mean_len = len(neighbors)
+ print('Mean length of neighbors:', mean_len)
+ logger.info("Mean length of neighbors: {0}".format(mean_len))
+
+ # GPT_LM = LMscore_calculator(data_path, args)
+ lissa_num_batches = math.ceil(data.shape[0]/args.lissa_batch_size)
+ logger.info('-------- Lissa Params for IHVP --------')
+ logger.info('Damping: {0}'.format(args.damping))
+ logger.info('Lissa_repeat: {0}'.format(args.lissa_repeat))
+ logger.info('Lissa_depth: {0}'.format(args.lissa_depth))
+ logger.info('Scale: {0}'.format(args.scale))
+ logger.info('Lissa batch size: {0}'.format(args.lissa_batch_size))
+ logger.info('Lissa num bacthes: {0}'.format(lissa_num_batches))
+
+ score_path = os.path.join('attack_results', args.data, 'score_cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ str(args.added_edge_num)))
+
+ if args.attack_goal == 'single':
+ attack_trip, score_record = addition_attack(param_influence, args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, load_Record = args.load_existed, divide_bound = divide_bound, data_mean = data_mean, data_std = data_std)
+ else:
+ # lissa = before_global_attack(args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, lissa_path, target_disease)
+
+ attack_trip, score_record = global_addtion_attack(args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, None, target_disease)
+
+ utils.save_data(attack_path, attack_trip)
+
+ logger.info("Attack triples are saved in " + attack_path)
+ with open(score_path, 'wb') as fl:
+ pkl.dump(score_record, fl)

DiseaseSpecific/edge_to_abstract.py

@@ -0,0 +1,530 @@
+#%%
+import torch
+import numpy as np
+from torch.autograd import Variable
+# from sklearn import metrics
+
+import datetime
+from typing import Dict, Tuple, List
+import logging
+import os
+import utils
+import pickle as pkl
+import json 
+import torch.backends.cudnn as cudnn
+
+from tqdm import tqdm
+
+import sys
+sys.path.append("..")
+import Parameters
+
+parser = utils.get_argument_parser()
+parser = utils.add_attack_parameters(parser)
+parser.add_argument('--mode', type=str, default='sentence', help='sentence, biogpt or finetune')
+parser.add_argument('--ratio', type = str, default='', help='ratio of the number of changed words')
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ 
+utils.seed_all(args.seed)
+np.set_printoptions(precision=5)
+cudnn.benchmark = False
+
+data_path = os.path.join('processed_data', args.data)
+target_path = os.path.join(data_path, 'DD_target_{0}_{1}_{2}_{3}_{4}_{5}.txt'.format(args.model, args.data, args.target_split, args.target_size, 'exists:'+str(args.target_existed), args.attack_goal))
+attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate)))
+# target_data = utils.load_data(target_path)
+attack_data = utils.load_data(attack_path, drop=False)
+# assert target_data.shape == attack_data.shape
+#%%
+
+with open(os.path.join(data_path, 'entities_reverse_dict.json')) as fl:
+ id_to_meshid = json.load(fl)
+with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+ entity_raw_name = pkl.load(fl)
+with open(Parameters.GNBRfile+'retieve_sentence_through_edgetype', 'rb') as fl:
+ retieve_sentence_through_edgetype = pkl.load(fl)
+with open(Parameters.GNBRfile+'raw_text_of_each_sentence', 'rb') as fl:
+ raw_text_sen = pkl.load(fl)
+
+if not os.path.exists('generate_abstract/valid_entity.json'):
+ valid_entity = set()
+ for paper_id, paper in raw_text_sen.items():
+ for sen_id, sen in paper.items():
+ text = sen['text'].split(' ')
+ for a in text:
+ if '_' in a:
+ valid_entity.add(a.replace('_', ' '))
+ with open('valid_entity.json', 'w') as fl:
+ json.dump(list(valid_entity), fl, indent=4)
+ print('Valid entity saved!!')
+
+if args.mode == 'sentence':
+ import torch 
+ from torch.nn.modules.loss import CrossEntropyLoss
+ from transformers import AutoTokenizer
+ from transformers import BioGptForCausalLM
+ criterion = CrossEntropyLoss(reduction="none")
+
+ print('Generating GPT input ...')
+
+ tokenizer = AutoTokenizer.from_pretrained('microsoft/biogpt')
+ tokenizer.pad_token = tokenizer.eos_token
+ model = BioGptForCausalLM.from_pretrained('microsoft/biogpt', pad_token_id=tokenizer.eos_token_id)
+ model.to(device)
+ model.eval()
+ GPT_batch_size = 32
+ single_sentence = {}
+ test_text = []
+ test_dp = []
+ test_parse = []
+ for i, (s, r, o) in enumerate(tqdm(attack_data)):
+
+ if int(s) != -1:
+
+ dependency_sen_dict = retieve_sentence_through_edgetype[int(r)]['manual']
+ candidate_sen = []
+ Dp_path = []
+ L = len(dependency_sen_dict.keys())
+ bound = 500 // L
+ if bound == 0:
+ bound = 1
+ for dp_path, sen_list in dependency_sen_dict.items():
+ if len(sen_list) > bound:
+ index = np.random.choice(np.array(range(len(sen_list))), bound, replace=False)
+ sen_list = [sen_list[aa] for aa in index]
+ candidate_sen += sen_list
+ Dp_path += [dp_path] * len(sen_list)
+
+ text_s = entity_raw_name[id_to_meshid[s]]
+ text_o = entity_raw_name[id_to_meshid[o]]
+ candidate_text_sen = []
+ candidate_ori_sen = []
+ candidate_parse_sen = []
+
+ for paper_id, sen_id in candidate_sen:
+ sen = raw_text_sen[paper_id][sen_id]
+ text = sen['text']
+ candidate_ori_sen.append(text)
+ ss = sen['start_formatted']
+ oo = sen['end_formatted']
+ text = text.replace('-LRB-', '(')
+ text = text.replace('-RRB-', ')')
+ text = text.replace('-LSB-', '[')
+ text = text.replace('-RSB-', ']')
+ text = text.replace('-LCB-', '{')
+ text = text.replace('-RCB-', '}')
+ parse_text = text
+ parse_text = parse_text.replace(ss, text_s.replace(' ', '_'))
+ parse_text = parse_text.replace(oo, text_o.replace(' ', '_'))
+ text = text.replace(ss, text_s)
+ text = text.replace(oo, text_o)
+ text = text.replace('_', ' ')
+ candidate_text_sen.append(text)
+ candidate_parse_sen.append(parse_text)
+ tokens = tokenizer( candidate_text_sen,
+ truncation = True,
+ padding = True,
+ max_length = 300,
+ return_tensors="pt")
+ target_ids = tokens['input_ids'].to(device)
+ attention_mask = tokens['attention_mask'].to(device)
+
+ L = len(candidate_text_sen)
+ assert L > 0
+ ret_log_L = []
+ for l in range(0, L, GPT_batch_size):
+ R = min(L, l + GPT_batch_size)
+ target = target_ids[l:R, :]
+ attention = attention_mask[l:R, :]
+ outputs = model(input_ids = target,
+ attention_mask = attention,
+ labels = target)
+ logits = outputs.logits
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = target[..., 1:].contiguous()
+ Loss = criterion(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1))
+ Loss = Loss.view(-1, shift_logits.shape[1])
+ attention = attention[..., 1:].contiguous()
+ log_Loss = (torch.mean(Loss * attention.float(), dim = 1) / torch.mean(attention.float(), dim = 1))
+ ret_log_L.append(log_Loss.detach())
+ 
+
+ ret_log_L = list(torch.cat(ret_log_L, -1).cpu().numpy())
+ sen_score = list(zip(candidate_text_sen, ret_log_L, candidate_ori_sen, Dp_path, candidate_parse_sen))
+ sen_score.sort(key = lambda x: x[1])
+ test_text.append(sen_score[0][2])
+ test_dp.append(sen_score[0][3])
+ test_parse.append(sen_score[0][4])
+ single_sentence.update({f'{s}_{r}_{o}_{i}': sen_score[0][0]})
+
+ else:
+ single_sentence.update({f'{s}_{r}_{o}_{i}': ''})
+
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_sentence.json', 'w') as fl:
+ json.dump(single_sentence, fl, indent=4)
+ # with open('generate_abstract/test.txt', 'w') as fl:
+ # fl.write('\n'.join(test_text))
+ # with open('generate_abstract/dp.txt', 'w') as fl:
+ # fl.write('\n'.join(test_dp))
+ with open (f'generate_abstract/path/{args.target_split}_{args.reasonable_rate}_path.json', 'w') as fl:
+ fl.write('\n'.join(test_dp))
+ with open (f'generate_abstract/path/{args.target_split}_{args.reasonable_rate}_temp.json', 'w') as fl:
+ fl.write('\n'.join(test_text))
+
+elif args.mode == 'finetune':
+
+ import spacy
+ import pprint
+ from transformers import AutoModel, AutoTokenizer,BartForConditionalGeneration
+
+ print('Finetuning ...')
+
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_chat.json', 'r') as fl:
+ draft = json.load(fl)
+ with open (f'generate_abstract/path/{args.target_split}_{args.reasonable_rate}_path.json', 'r') as fl:
+ dpath = fl.readlines()
+ 
+ nlp = spacy.load("en_core_web_sm")
+ if os.path.exists(f'generate_abstract/bioBART/{args.target_split}_{args.reasonable_rate}{args.ratio}_candidates.json'):
+ with open(f'generate_abstract/bioBART/{args.target_split}_{args.reasonable_rate}{args.ratio}_candidates.json', 'r') as fl:
+ ret_candidates = json.load(fl)
+ # if False:
+ # pass
+ else:
+
+ def find_mini_span(vec, words, check_set):
+ 
+
+ def cal(text, sset):
+ add = 0
+ for tt in sset:
+ if tt in text:
+ add += 1
+ return add
+ text = ' '.join(words)
+ max_add = cal(text, check_set)
+
+ minn = 10000000
+ span = ''
+ rc = None
+ for i in range(len(vec)):
+ if vec[i] == True:
+ p = -1
+ for j in range(i+1, len(vec)+1):
+ if vec[j-1] == True:
+ text = ' '.join(words[i:j])
+ if cal(text, check_set) == max_add:
+ p = j
+ break
+ if p > 0:
+ if (p-i) < minn:
+ minn = p-i
+ span = ' '.join(words[i:p])
+ rc = (i, p)
+ if rc:
+ for i in range(rc[0], rc[1]):
+ vec[i] = True
+ return vec, span
+ 
+ def mask_func(tokenized_sen):
+
+ if len(tokenized_sen) == 0:
+ return []
+ token_list = []
+ # for sen in tokenized_sen:
+ # for token in sen:
+ # token_list.append(token)
+ for sen in tokenized_sen:
+ token_list += sen.text.split(' ')
+ if args.ratio == '':
+ P = 0.3
+ else:
+ P = float(args.ratio)
+
+ ret_list = []
+ i = 0
+ mask_num = 0
+ while i < len(token_list):
+ t = token_list[i]
+ if '.' in t or '(' in t or ')' in t or '[' in t or ']' in t:
+ ret_list.append(t)
+ i += 1
+ mask_num = 0
+ else:
+ length = np.random.poisson(3)
+ if np.random.rand() < P and length > 0:
+ if mask_num < 8:
+ ret_list.append('<mask>')
+ mask_num += 1
+ i += length
+ else:
+ ret_list.append(t)
+ i += 1
+ mask_num = 0
+ return [' '.join(ret_list)]
+ 
+ model = BartForConditionalGeneration.from_pretrained('GanjinZero/biobart-large')
+ model.eval()
+ model.to(device)
+ tokenizer = AutoTokenizer.from_pretrained('GanjinZero/biobart-large')
+
+ ret_candidates = {}
+ dpath_i = 0
+
+ for i,(k, v) in enumerate(tqdm(draft.items())):
+
+ input = v['in'].replace('\n', '')
+ output = v['out'].replace('\n', '')
+ s, r, o = attack_data[i]
+
+ if int(s) == -1:
+ ret_candidates[str(i)] = {'span': '', 'prompt' : '', 'out' : [], 'in': [], 'assist': []}
+ continue
+
+ path_text = dpath[dpath_i].replace('\n', '')
+ dpath_i += 1
+ text_s = entity_raw_name[id_to_meshid[s]]
+ text_o = entity_raw_name[id_to_meshid[o]]
+
+ doc = nlp(output)
+ words= input.split(' ')
+ tokenized_sens = [sen for sen in doc.sents]
+ sens = np.array([sen.text for sen in doc.sents])
+
+ checkset = set([text_s, text_o])
+ e_entity = set(['start_entity', 'end_entity'])
+ for path in path_text.split(' '):
+ a, b, c = path.split('|')
+ if a not in e_entity:
+ checkset.add(a)
+ if c not in e_entity:
+ checkset.add(c)
+ vec = []
+ l = 0
+ while(l < len(words)):
+ bo =False
+ for j in range(len(words), l, -1): # reversing is important !!!
+ cc = ' '.join(words[l:j])
+ if (cc in checkset):
+ vec += [True] * (j-l)
+ l = j
+ bo = True
+ break
+ if not bo:
+ vec.append(False)
+ l += 1
+ vec, span = find_mini_span(vec, words, checkset)
+ # vec = np.vectorize(lambda x: x in checkset)(words)
+ vec[-1] = True
+ prompt = []
+ mask_num = 0
+ for j, bo in enumerate(vec):
+ if not bo:
+ mask_num += 1
+ else:
+ if mask_num > 0:
+ # mask_num = mask_num // 3 # span length ~ poisson distribution (lambda = 3)
+ mask_num = max(mask_num, 1)
+ mask_num= min(8, mask_num)
+ prompt += ['<mask>'] * mask_num
+ prompt.append(words[j])
+ mask_num = 0
+ prompt = ' '.join(prompt)
+ Text = []
+ Assist = []
+
+ for j in range(len(sens)):
+ Bart_input = list(sens[:j]) + [prompt] +list(sens[j+1:])
+ assist = list(sens[:j]) + [input] +list(sens[j+1:])
+ Text.append(' '.join(Bart_input))
+ Assist.append(' '.join(assist))
+ 
+ for j in range(len(sens)):
+ Bart_input = mask_func(tokenized_sens[:j]) + [input] + mask_func(tokenized_sens[j+1:])
+ assist = list(sens[:j]) + [input] +list(sens[j+1:])
+ Text.append(' '.join(Bart_input))
+ Assist.append(' '.join(assist))
+
+ batch_size = len(Text) // 2
+ Outs = []
+ for l in range(2):
+ A = tokenizer(Text[batch_size * l:batch_size * (l+1)],
+ truncation = True,
+ padding = True,
+ max_length = 1024,
+ return_tensors="pt")
+ input_ids = A['input_ids'].to(device)
+ attention_mask = A['attention_mask'].to(device)
+ aaid = model.generate(input_ids, attention_mask = attention_mask, num_beams = 5, max_length = 1024)
+ outs = tokenizer.batch_decode(aaid, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+ Outs += outs
+ ret_candidates[str(i)] = {'span': span, 'prompt' : prompt, 'out' : Outs, 'in': Text, 'assist': Assist}
+ with open(f'generate_abstract/bioBART/{args.target_split}_{args.reasonable_rate}{args.ratio}_candidates.json', 'w') as fl:
+ json.dump(ret_candidates, fl, indent = 4)
+ 
+ from torch.nn.modules.loss import CrossEntropyLoss
+ from transformers import BioGptForCausalLM
+ criterion = CrossEntropyLoss(reduction="none")
+
+ tokenizer = AutoTokenizer.from_pretrained('microsoft/biogpt')
+ tokenizer.pad_token = tokenizer.eos_token
+ model = BioGptForCausalLM.from_pretrained('microsoft/biogpt', pad_token_id=tokenizer.eos_token_id)
+ model.to(device)
+ model.eval()
+
+ scored = {}
+ ret = {}
+ dpath_i = 0
+ for i,(k, v) in enumerate(tqdm(draft.items())):
+
+ span = ret_candidates[str(i)]['span']
+ prompt = ret_candidates[str(i)]['prompt']
+ sen_list = ret_candidates[str(i)]['out']
+ BART_in = ret_candidates[str(i)]['in']
+ Assist = ret_candidates[str(i)]['assist']
+
+ s, r, o = attack_data[i]
+
+ if int(s) == -1:
+ ret[k] = {'prompt': '', 'in':'', 'out': ''}
+ continue
+
+ text_s = entity_raw_name[id_to_meshid[s]]
+ text_o = entity_raw_name[id_to_meshid[o]]
+
+ def process(text):
+
+ for i in range(ord('A'), ord('Z')+1):
+ text = text.replace(f'.{chr(i)}', f'. {chr(i)}') 
+ return text
+
+ sen_list = [process(text) for text in sen_list]
+ path_text = dpath[dpath_i].replace('\n', '')
+ dpath_i += 1
+
+ checkset = set([text_s, text_o])
+ e_entity = set(['start_entity', 'end_entity'])
+ for path in path_text.split(' '):
+ a, b, c = path.split('|')
+ if a not in e_entity:
+ checkset.add(a)
+ if c not in e_entity:
+ checkset.add(c)
+
+ input = v['in'].replace('\n', '')
+ output = v['out'].replace('\n', '')
+
+ doc = nlp(output)
+ gpt_sens = [sen.text for sen in doc.sents]
+ assert len(gpt_sens) == len(sen_list) // 2
+
+ word_sets = []
+ for sen in gpt_sens:
+ word_sets.append(set(sen.split(' ')))
+
+ def sen_align(word_sets, modified_word_sets):
+ 
+ l = 0
+ while(l < len(modified_word_sets)):
+ if len(word_sets[l].intersection(modified_word_sets[l])) > len(word_sets[l]) * 0.8:
+ l += 1
+ else:
+ break
+ if l == len(modified_word_sets):
+ return -1, -1, -1, -1
+ r = l + 1
+ r1 = None
+ r2 = None
+ for pos1 in range(r, len(word_sets)):
+ for pos2 in range(r, len(modified_word_sets)):
+ if len(word_sets[pos1].intersection(modified_word_sets[pos2])) > len(word_sets[pos1]) * 0.8:
+ r1 = pos1
+ r2 = pos2
+ break
+ if r1 is not None:
+ break
+ if r1 is None:
+ r1 = len(word_sets)
+ r2 = len(modified_word_sets)
+ return l, r1, l, r2
+
+ replace_sen_list = []
+ boundary = []
+ assert len(sen_list) % 2 == 0
+ for j in range(len(sen_list) // 2):
+ doc = nlp(sen_list[j])
+ sens = [sen.text for sen in doc.sents]
+ modified_word_sets = [set(sen.split(' ')) for sen in sens]
+ l1, r1, l2, r2 = sen_align(word_sets, modified_word_sets)
+ boundary.append((l1, r1, l2, r2))
+ if l1 == -1:
+ replace_sen_list.append(sen_list[j])
+ continue
+ check_text = ' '.join(sens[l2: r2])
+ replace_sen_list.append(' '.join(gpt_sens[:l1] + [check_text] + gpt_sens[r1:]))
+ sen_list = replace_sen_list + sen_list[len(sen_list) // 2:]
+
+ old_L = len(sen_list)
+ sen_list.append(output)
+ sen_list += Assist
+ tokens = tokenizer( sen_list,
+ truncation = True,
+ padding = True,
+ max_length = 1024,
+ return_tensors="pt")
+ target_ids = tokens['input_ids'].to(device)
+ attention_mask = tokens['attention_mask'].to(device)
+ L = len(sen_list)
+ ret_log_L = []
+ for l in range(0, L, 5):
+ R = min(L, l + 5)
+ target = target_ids[l:R, :]
+ attention = attention_mask[l:R, :]
+ outputs = model(input_ids = target,
+ attention_mask = attention,
+ labels = target)
+ logits = outputs.logits
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = target[..., 1:].contiguous()
+ Loss = criterion(shift_logits.view(-1, shift_logits.shape[-1]), shift_labels.view(-1))
+ Loss = Loss.view(-1, shift_logits.shape[1])
+ attention = attention[..., 1:].contiguous()
+ log_Loss = (torch.mean(Loss * attention.float(), dim = 1) / torch.mean(attention.float(), dim = 1))
+ ret_log_L.append(log_Loss.detach())
+ log_Loss = torch.cat(ret_log_L, -1).cpu().numpy()
+
+ real_log_Loss = log_Loss.copy()
+
+ log_Loss = log_Loss[:old_L]
+
+ p = np.argmin(log_Loss)
+ content = []
+ for i in range(len(real_log_Loss)):
+ content.append([sen_list[i], str(real_log_Loss[i])])
+ scored[k] = {'path':path_text, 'prompt': prompt, 'in':input, 's':text_s, 'o':text_o, 'out': content, 'bound': boundary}
+ p_p = p
+ # print('Old_L:', old_L)
+
+ if real_log_Loss[p] > real_log_Loss[p+1+old_L]:
+ p_p = p+1+old_L
+
+ if real_log_Loss[p] > real_log_Loss[old_L]:
+ if real_log_Loss[p] > real_log_Loss[p+1+old_L]:
+ p = p+1+old_L
+ ret[k] = {'prompt': prompt, 'in':input, 'out': sen_list[p]}
+ with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}{args.ratio}_bioBART_finetune.json', 'w') as fl:
+ json.dump(ret, fl, indent=4)
+ with open(f'generate_abstract/bioBART/{args.target_split}_{args.reasonable_rate}{args.ratio}_scored.json', 'w') as fl:
+ json.dump(scored, fl, indent=4)
+else:
+ raise Exception('Wrong mode !!')

DiseaseSpecific/evaluation.py

@@ -0,0 +1,499 @@
+import torch
+import numpy as np
+from torch.autograd import Variable
+from sklearn import metrics
+
+import datetime
+from typing import Dict, Tuple, List
+import logging
+import os
+import utils
+import pickle as pkl
+import json 
+from tqdm import tqdm
+import torch.backends.cudnn as cudnn
+
+import sys
+sys.path.append("..")
+import Parameters
+
+logger = logging.getLogger(__name__)
+
+def get_model_loss_without_softmax(batch, model, device=None):
+
+ with torch.no_grad():
+ s,r,o = batch[:,0], batch[:,1], batch[:,2]
+
+ emb_s = model.emb_e(s).squeeze(dim=1)
+ emb_r = model.emb_rel(r).squeeze(dim=1)
+
+ pred = model.forward(emb_s, emb_r)
+ return -pred[range(o.shape[0]), o]
+
+def check(trip, model, reasonable_rate, device, data_mean = -4.008113861083984, data_std = 5.153779983520508, divide_bound = 0.05440050354114886):
+
+ if args.model == 'distmult':
+ pass
+ elif args.model == 'conve':
+ data_mean = 13.890259742
+ data_std = 12.396190643
+ divide_bound = -0.1986345871
+ else:
+ raise Exception('Wrong model!!')
+ trip = np.array(trip)
+ train_trip = trip[None, :] 
+ train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+ edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze().item()
+
+ bound = 1 - reasonable_rate
+ edge_loss = (edge_loss - data_mean) / data_std
+ edge_loss_prob = 1 / ( 1 + np.exp(edge_loss - divide_bound))
+ return edge_loss_prob > bound
+
+
+def get_ranking(model, queries,
+ valid_filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
+ device, batch_size, entityid_to_nodetype, exists_edge):
+ """
+ Ranking for target generation.
+ """
+ ranks = []
+ total_nums = []
+ b_begin = 0
+
+ for b_begin in range(0, len(queries), 1):
+ b_queries = queries[b_begin : b_begin+1]
+ s,r,o = b_queries[:,0], b_queries[:,1], b_queries[:,2]
+ r_rev = r
+ lhs_score = model.score_or(o, r_rev, sigmoid=False) #this gives scores not probabilities
+ # print(b_queries.shape)
+ for i, query in enumerate(b_queries):
+
+ if not args.target_existed:
+ tp1 = entityid_to_nodetype[str(query[0].item())]
+ tp2 = entityid_to_nodetype[str(query[2].item())]
+ filter = valid_filters['lhs'][(tp2, query[1].item())].clone()
+ filter[exists_edge['lhs'][str(query[2].item())]] = False
+ filter = (filter == False)
+ else:
+ tp1 = entityid_to_nodetype[str(query[0].item())]
+ tp2 = entityid_to_nodetype[str(query[2].item())]
+ filter = valid_filters['lhs'][(tp2, query[1].item())]
+ filter = (filter == False)
+ 
+ # if (str(query[2].item())) == '16566':
+ # print('16566', filter.sum(), valid_filters['lhs'][(tp2, query[1].item())].sum(), tp2, query[1].item())
+ # raise Exception('??')
+
+ score = lhs_score
+ # target_value = rhs_score[i, query[0].item()].item()
+ # zero all known cases (this are not interesting)
+ # this corresponds to the filtered setting
+ score[i][filter] = 1e6
+ total_nums.append(n_ent - filter.sum().item())
+ # write base the saved values
+ # if b_begin < len(queries) // 2:
+ # score[i][query[2].item()] = target_value
+ # else:
+ # score[i][query[0].item()] = target_value
+
+ # sort and rank
+ min_values, sort_v = torch.sort(score, dim=1, descending=False) #low scores get low number ranks
+
+ sort_v = sort_v.cpu().numpy()
+ 
+ for i, query in enumerate(b_queries):
+ # find the rank of the target entities
+ rank = np.where(sort_v[i]==query[0].item())[0][0]
+
+ # rank+1, since the lowest rank is rank 1 not rank 0
+ ranks.append(rank)
+
+ #logger.info('Ranking done for all queries')
+ return ranks, total_nums
+ 
+ 
+def evaluation(model, queries,
+ valid_filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
+ device, batch_size, entityid_to_nodetype, exists_edge, eval_type = '', attack_data = None, ori_ranks = None, ori_totals = None):
+ 
+ #get ranking
+ ranks, total_nums = get_ranking(model, queries, valid_filters, device, batch_size, entityid_to_nodetype, exists_edge)
+ ranks, total_nums = np.array(ranks), np.array(total_nums)
+ # print(ranks)
+ # print(total_nums)
+ # print(ranks)
+ # print(total_nums)
+
+ ranks = total_nums - ranks
+
+ if (attack_data is not None):
+ for i, tri in enumerate(attack_data):
+ if args.mode == '':
+ if args.added_edge_num == '' or int(args.added_edge_num) == 1:
+ if int(tri[0]) == -1:
+ ranks[i] = ori_ranks[i]
+ total_nums[i] = ori_totals[i]
+ else:
+ if int(tri[0][0]) == -1:
+ ranks[i] = ori_ranks[i]
+ total_nums[i] = ori_totals[i]
+ else:
+ if len(tri) == 0:
+ ranks[i] = ori_ranks[i]
+ total_nums[i] = ori_totals[i]
+
+ mean = (ranks / total_nums).mean()
+ std = (ranks / total_nums).std()
+ #final logging
+ hits_at = np.arange(1,11)
+ hits_at_both = list(map(lambda x: np.mean((ranks <= x), dtype=np.float64).item(), 
+ hits_at))
+ mr = np.mean(ranks, dtype=np.float64).item()
+ 
+ mrr = np.mean(1. / ranks, dtype=np.float64).item()
+ 
+ logger.info('')
+ logger.info('-'*50)
+ # logger.info(split+'_'+save_name)
+ logger.info('')
+ if eval_type:
+ logger.info(eval_type)
+ else:
+ logger.info('after attck')
+
+ for i in hits_at:
+ logger.info('Hits @{0}: {1}'.format(i, hits_at_both[i-1]))
+ logger.info('Mean rank: {0}'.format( mr))
+ logger.info('Mean reciprocal rank lhs: {0}'.format(mrr))
+ logger.info('Mean proportion: {0}'.format(mean))
+ logger.info('Std proportion: {0}'.format(std))
+ logger.info('Mean candidate num: {0}'.format(np.mean(total_nums)))
+ 
+# with open(os.path.join('results', split + '_' + save_name + '.txt'), 'a') as text_file:
+# text_file.write('Epoch: {0}\n'.format(epoch))
+# text_file.write('Lhs denotes ranking by subject corruptions \n')
+# text_file.write('Rhs denotes ranking by object corruptions \n')
+# for i in hits_at:
+# text_file.write('Hits left @{0}: {1}\n'.format(i, hits_at_lhs[i-1]))
+# text_file.write('Hits right @{0}: {1}\n'.format(i, hits_at_rhs[i-1]))
+# text_file.write('Hits @{0}: {1}\n'.format(i, np.mean([hits_at_lhs[i-1],hits_at_rhs[i-1]]).item()))
+# text_file.write('Mean rank lhs: {0}\n'.format( mr_lhs))
+# text_file.write('Mean rank rhs: {0}\n'.format(mr_rhs))
+# text_file.write('Mean rank: {0}\n'.format( np.mean([mr_lhs, mr_rhs])))
+# text_file.write('MRR lhs: {0}\n'.format( mrr_lhs))
+# text_file.write('MRR rhs: {0}\n'.format(mrr_rhs))
+# text_file.write('MRR: {0}\n'.format(np.mean([mrr_rhs, mrr_lhs])))
+# text_file.write('-------------------------------------------------\n')
+ 
+ 
+ results = {}
+ for i in hits_at:
+ results['hits @{}'.format(i)] = hits_at_both[i-1]
+ results['mrr'] = mrr
+ results['mr'] = mr
+ results['proportion'] = mean
+ results['std'] = std
+ 
+ return results, list(ranks), list(total_nums)
+
+
+parser = utils.get_argument_parser()
+parser = utils.add_attack_parameters(parser)
+parser = utils.add_eval_parameters(parser)
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ 
+utils.seed_all(args.seed)
+np.set_printoptions(precision=5)
+cudnn.benchmark = False
+ 
+data_path = os.path.join('processed_data', args.data)
+target_path = os.path.join(data_path, 'DD_target_{0}_{1}_{2}_{3}_{4}_{5}.txt'.format(args.model, args.data, args.target_split, args.target_size, 'exists:'+str(args.target_existed), args.attack_goal))
+
+log_path = 'logs/evaluation_logs/cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ args.mode)
+record_path = 'eval_record/{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}{9}{10}'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ args.mode,
+ str(args.added_edge_num),
+ args.mask_ratio)
+init_record_path = 'eval_record/{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ 'init')
+
+if args.seperate:
+ record_path += '_seperate'
+ log_path += '_seperate'
+else:
+ record_path += '_batch'
+
+if args.direct:
+ log_path += '_direct'
+ record_path += '_direct'
+else:
+ log_path += '_nodirect'
+ record_path += '_nodirect'
+
+dis_turbrbed_path_pre = os.path.join(data_path, 'evaluation')
+logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+ datefmt = '%m/%d/%Y %H:%M:%S',
+ level = logging.INFO,
+ filename = log_path
+ )
+logger = logging.getLogger(__name__)
+logger.info(vars(args))
+
+n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+model_name = '{0}_{1}_{2}_{3}_{4}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop)
+model_path = 'saved_models/{0}_{1}.model'.format(args.data, model_name)
+model = utils.load_model(model_path, args, n_ent, n_rel, device)
+
+ori_data = utils.load_data(os.path.join(data_path, 'all.txt'))
+target_data = utils.load_data(target_path)
+
+index = range(len(target_data))
+index = np.random.permutation(index)
+target_data = target_data[index]
+
+if args.direct:
+ assert args.attack_goal == 'single'
+ raise Exception('This option is abandoned in this version .')
+ # disturbed_data = list(ori_data) + list(target_data)
+else:
+ 
+ attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}{9}{10}.txt'.format(args.model, 
+ args.target_split, 
+ args.target_size, 
+ 'exists:'+str(args.target_existed),
+ args.neighbor_num,
+ args.candidate_mode,
+ args.attack_goal,
+ str(args.reasonable_rate),
+ args.mode,
+ str(args.added_edge_num),
+ args.mask_ratio))
+ if args.mode == '':
+ attack_data = utils.load_data(attack_path, drop=False)
+ if not(args.added_edge_num == '' or int(args.added_edge_num) == 1):
+ assert int(args.added_edge_num) * len(target_data) == len(attack_data)
+ attack_data = attack_data.reshape((len(target_data), int(args.added_edge_num), 3))
+ attack_data = attack_data[index]
+ else:
+ assert len(target_data) == len(attack_data)
+ attack_data = attack_data[index]
+ # if not args.seperate:
+ # disturbed_data = list(ori_data) + list(attack_data)
+ else:
+ with open(attack_path, 'rb') as fl:
+ attack_data = pkl.load(fl)
+
+ tmp_attack_data = []
+ for vv in attack_data:
+ a_attack = []
+ for v in vv:
+ if check(v, model, args.reasonable_rate, device):
+ a_attack.append(v)
+ tmp_attack_data.append(a_attack)
+ attack_data = tmp_attack_data
+ attack_data = [attack_data[i] for i in index]
+
+ # if not args.seperate:
+ # disturbed_data = list(ori_data)
+ # if args.mode == '':
+ # for aa in list(attack_data):
+ # if int(aa[0]) != -1:
+ # disturbed_data.append(aa)
+ # else:
+ # for vv in attack_data:
+ # for v in vv:
+ # disturbed_data.append(v)
+
+with open(os.path.join(data_path, 'filter.pickle'), 'rb') as fl:
+ valid_filters = pkl.load(fl)
+with open(os.path.join(data_path, 'entityid_to_nodetype.json'), 'r') as fl:
+ entityid_to_nodetype = json.load(fl)
+with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+ entity_raw_name = pkl.load(fl)
+with open(os.path.join(data_path, 'disease_meshid.pickle'), 'rb') as fl:
+ disease_meshid = pkl.load(fl)
+with open(os.path.join(data_path, 'entities_dict.json'), 'r') as fl:
+ entity_to_id = json.load(fl)
+
+if args.attack_goal == 'global':
+ raise Exception('Please refer to pagerank method in global setting.')
+ # target_disease = []
+ # tid = 1
+ # bound = 50
+ # while True:
+ # meshid = disease_meshid[tid][0]
+ # fre = disease_meshid[tid][1]
+ # if len(entity_raw_name[meshid]) > 4:
+ # target_disease.append(entity_to_id[meshid])
+ # bound -= 1
+ # if bound == 0:
+ # break
+ # tid += 1
+ # s_set = set()
+ # for s, r, o in target_data:
+ # s_set.add(s)
+ # target_data = list(s_set)
+ # target_data.sort()
+
+ # target_list = []
+ # for s in target_data:
+ # for o in target_disease:
+ # target_list.append([str(s), str(10), str(o)])
+ # target_data = np.array(target_list, dtype = str)
+
+init_mask = np.asarray([0] * n_ent).astype('int64')
+init_mask = (init_mask == 1)
+for k, v in valid_filters.items():
+ for kk, vv in v.items():
+ tmp = init_mask.copy()
+ tmp[np.asarray(vv)] = True
+ t = torch.ByteTensor(tmp).to(device)
+ valid_filters[k][kk] = t
+# print('what??', valid_filters['lhs'][('disease', 10)].sum())
+
+exists_edge = {'lhs':{}, 'rhs':{}}
+for s, r, o in ori_data:
+ if s not in exists_edge['rhs'].keys():
+ exists_edge['rhs'][s] = []
+ if o not in exists_edge['lhs'].keys():
+ exists_edge['lhs'][o] = []
+ exists_edge['rhs'][s].append(int(o))
+ exists_edge['lhs'][o].append(int(s))
+target_data = torch.from_numpy(target_data.astype('int64')).to(device)
+# print(target_data[:5, :])
+ori_results, ori_ranks, ori_totals = evaluation(model, target_data, valid_filters, device, args.test_batch_size, entityid_to_nodetype, exists_edge, 'original')
+print('Original:', ori_results)
+with open(init_record_path, 'wb') as fl:
+ pkl.dump([ori_results, ori_ranks, ori_totals], fl)
+
+# raise Exception('Check Original Rank!!!')
+
+thread_name = args.model+'_'+args.target_split+'_'+args.attack_goal+'_'+str(args.reasonable_rate)+str(args.added_edge_num)+str(args.mask_ratio)
+if args.direct:
+ thread_name += '_direct'
+else:
+ thread_name += '_nodirect'
+if args.seperate:
+ thread_name += '_seperate'
+else:
+ thread_name += '_batch'
+thread_name += args.mode
+
+disturbed_data_path = os.path.join(dis_turbrbed_path_pre, 'all_{}.txt'.format(thread_name))
+
+if args.seperate:
+ # assert len(attack_data) * len(target_disease) == len(target_data)
+ assert len(attack_data) == len(target_data)
+ # final_result = None
+ Ranks = []
+ Totals = []
+ print('Training model {}...'.format(thread_name))
+ for i in tqdm(range(len(attack_data))):
+ attack_trip = attack_data[i]
+ if args.mode == '':
+ attack_trip = [attack_trip]
+ # target = target_data[i*len(target_disease) : (i+1)*len(target_disease)]
+ target = target_data[i: i+1, :]
+ if len(attack_trip) > 0 and int(attack_trip[0][0]) != -1:
+ disturbed_data = list(ori_data) + attack_trip
+ disturbed_data = np.array(disturbed_data)
+ utils.save_data(disturbed_data_path, disturbed_data)
+
+ cmd = 'CUDA_VISIBLE_DEVICES={} python main_multiprocess.py --data {} --model {} --thread-name {}'.format(args.cuda_name,args.data, args.model, thread_name)
+ os.system(cmd)
+ model_name = '{0}_{1}_{2}_{3}_{4}_{5}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop, thread_name)
+ model_path = 'saved_models/evaluation/{0}_{1}.model'.format(args.data, model_name)
+ model = utils.load_model(model_path, args, n_ent, n_rel, device)
+ a_results, a_ranks, a_total_nums = evaluation(model, target, valid_filters, device, args.test_batch_size, entityid_to_nodetype, exists_edge)
+ assert len(a_ranks) == 1
+ if not final_result:
+ final_result = a_results
+ else:
+ for k in final_result.keys():
+ final_result[k] += a_results[k]
+ Ranks += a_ranks
+ Totals += a_total_nums
+ else:
+ Ranks += [ori_ranks[i]]
+ Totals += [ori_totals[i]]
+ final_result['proportion'] += ori_ranks[i] / ori_totals[i]
+ for k in final_result.keys():
+ final_result[k] /= attack_data.shape[0]
+ print('Final !!!')
+ print(final_result)
+ logger.info('Final !!!!')
+ for k, v in final_result.items():
+ logger.info('{} : {}'.format(k, v))
+ tmp = np.array(Ranks) / np.array(Totals)
+ print('Std:', np.std(tmp))
+ with open(record_path, 'wb') as fl:
+ pkl.dump([final_result, Ranks, Totals], fl)
+
+else:
+ assert len(target_data) == len(attack_data)
+ print('Attack shape:' , len(attack_data))
+ Results = []
+ Ranks = []
+ Totals = []
+ for l in range(0, len(target_data), 50):
+ r = min(l+50, len(target_data))
+ t_target_data = target_data[l:r]
+ t_attack_data = attack_data[l:r]
+ t_ori_ranks = ori_ranks[l:r]
+ t_ori_totals = ori_totals[l:r]
+ if args.mode == '':
+ if not(args.added_edge_num == '' or int(args.added_edge_num) == 1):
+ tt_attack_data = []
+ for vv in t_attack_data:
+ tt_attack_data += list(vv)
+ t_attack_data = tt_attack_data
+ else:
+ assert args.mode == 'sentence' or args.mode == 'bioBART'
+ tt_attack_data = []
+ for vv in t_attack_data:
+ tt_attack_data += vv
+ t_attack_data = tt_attack_data
+ disturbed_data = list(ori_data) + list(t_attack_data)
+
+
+ utils.save_data(disturbed_data_path, disturbed_data)
+ cmd = 'CUDA_VISIBLE_DEVICES={} python main_multiprocess.py --data {} --model {} --thread-name {}'.format(args.cuda_name,args.data, args.model, thread_name)
+ print('Training model {}...'.format(thread_name))
+ os.system(cmd)
+ model_name = '{0}_{1}_{2}_{3}_{4}_{5}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop, thread_name)
+ model_path = 'saved_models/evaluation/{0}_{1}.model'.format(args.data, model_name)
+ model = utils.load_model(model_path, args, n_ent, n_rel, device)
+ a_results, a_ranks, a_totals = evaluation(model, t_target_data, valid_filters, device, args.test_batch_size, entityid_to_nodetype, exists_edge, attack_data = attack_data[l:r], ori_ranks = t_ori_ranks, ori_totals = t_ori_totals)
+ print(f'************Current l: {l}\n', a_results)
+ assert len(a_ranks) == t_target_data.shape[0]
+ Results += [a_results]
+ Ranks += list(a_ranks)
+ Totals += list(a_totals)
+ with open(record_path, 'wb') as fl:
+ pkl.dump([Results, Ranks, Totals, index], fl)

DiseaseSpecific/main.py

@@ -0,0 +1,377 @@
+#%%
+import pickle as pkl
+from typing import Dict, Tuple, List
+import os
+import numpy as np
+import json
+import logging
+import argparse 
+import math
+from pprint import pprint
+import pandas as pd
+from collections import defaultdict
+import copy
+import time
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+import torch.autograd as autograd
+
+from model import Distmult, Complex, Conve
+import utils
+
+# from evaluation import evaluation
+
+#%%
+class Main(object):
+ def __init__(self, args):
+ self.args = args 
+ 
+ self.model_name = '{0}_{1}_{2}_{3}_{4}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop)
+ #leaving batches from the model_name since they do not depend on model_architecture 
+ # also leaving kernel size and filters, siinice don't intend to change those
+ self.model_path = 'saved_models/{0}_{1}.model'.format(args.data, self.model_name)
+ 
+ self.log_path = 'logs/{0}_{1}_{2}_{3}.log'.format(args.data, self.model_name, args.epochs, args.train_batch_size)
+ self.loss_path = 'losses/{0}_{1}_{2}_{3}.pickle'.format(args.data, self.model_name, args.epochs, args.train_batch_size)
+ 
+ logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+ datefmt = '%m/%d/%Y %H:%M:%S',
+ level = logging.INFO,
+ filename = self.log_path)
+ self.logger = logging.getLogger(__name__)
+ self.logger.info(vars(self.args))
+ self.logger.info('\n')
+ 
+ 
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ 
+ self.load_data()
+ self.model = self.add_model()
+ self.optimizer = self.add_optimizer(self.model.parameters())
+ 
+ if self.args.save_influence_map:
+ self.logger.info('-------- Argument save_influence_map is set. Will use GR to compute and save influence maps ----------\n')
+ # when we want to save influence during training
+ self.args.add_reciprocals = False # to keep things simple
+ # init an empty influence map
+ self.influence_map = defaultdict(float)
+ #self.influence_path = 'influence_maps/{0}_{1}.json'.format(args.data, self.model_name)
+ self.influence_path = 'influence_maps/{0}_{1}.pickle'.format(args.data, self.model_name)
+ # Initialize a copy of the model prams to track previous weights in an epoch
+ self.previous_weights = [copy.deepcopy(param) for param in self.model.parameters()]
+ self.logger.info('Shape for previous weights: {}, {}'.format(self.previous_weights[0].shape, self.previous_weights[1].shape))
+ 
+ def load_data(self):
+ ''' 
+ Load the train, valid datasets
+ '''
+ data_path = os.path.join('processed_data', self.args.data)
+ n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+ self.n_ent = n_ent
+ self.n_rel = n_rel
+ 
+ self.train_data = utils.load_data(os.path.join(data_path, 'all.txt'))
+ # print(type(self.train_data), self.train_data.shape) #(1996432, 3)
+ tmp = np.random.choice(a = self.train_data.shape[0], size = int(self.train_data.shape[0] * self.args.KG_valid_rate), replace=False)
+ self.valid_data= self.train_data[tmp, :]
+
+ 
+ def add_model(self):
+ 
+ if self.args.model is None:
+ model = Distmult(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'distmult':
+ model = Distmult(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'complex':
+ model = Complex(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'conve':
+ model = Conve(self.args, self.n_ent, self.n_rel)
+ else:
+ self.logger.info('Unknown model: {0}', self.args.model)
+ raise Exception("Unknown model!")
+ model.to(self.device)
+ return model
+ 
+ def add_optimizer(self, parameters):
+ return torch.optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.lr_decay)
+ 
+ def save_model(self):
+ state = {
+ 'state_dict': self.model.state_dict(),
+ 'optimizer': self.optimizer.state_dict(),
+ 'args': vars(self.args)
+ }
+ torch.save(state, self.model_path)
+ self.logger.info('Saving model to {0}'.format(self.model_path))
+ 
+ def load_model(self):
+ self.logger.info('Loading saved model from {0}'.format(self.model_path))
+ state = torch.load(self.model_path)
+ model_params = state['state_dict']
+ params = [(key, value.size(), value.numel()) for key, value in model_params.items()]
+ for key, size, count in params:
+ self.logger.info(key, size, count)
+ self.model.load_state_dict(model_params)
+ self.optimizer.load_state_dict(state['optimizer'])
+ 
+ def lp_regularizer(self):
+ # Apply p-norm regularization; assign weights to each param
+ weight = self.args.reg_weight
+ p = self.args.reg_norm
+ 
+ trainable_params = [self.model.emb_e.weight, self.model.emb_rel.weight]
+ norm = 0
+ for i in range(len(trainable_params)):
+ #norm += weight * trainable_params[i].norm(p = p)**p
+ norm += weight * torch.sum( torch.abs(trainable_params[i]) ** p)
+ 
+ return norm
+ 
+ def n3_regularizer(self, factors):
+ # factors are the embeddings for lhs, rel, rhs for triples in a batch
+ weight = self.args.reg_weight
+ p = self.args.reg_norm
+ 
+ norm = 0
+ for f in factors:
+ norm += weight * torch.sum(torch.abs(f) ** p)
+ 
+ return norm / factors[0].shape[0] # scale by number of triples in batch
+ 
+ def get_influence_map(self):
+ """
+ Turns the influence map into a list, ready to be written to disc. (before: numpy)
+ :return: the influence map with lists as values
+ """
+ assert self.args.save_influence_map == True
+ 
+ for key in self.influence_map:
+ self.influence_map[key] = self.influence_map[key].tolist()
+ #self.logger.info('get_influence_map passed')
+ return self.influence_map
+
+ def evaluate(self, split, batch_size, epoch):
+ """
+ The same as self.run_epoch()
+ """
+
+ self.model.eval()
+ losses = []
+
+ with torch.no_grad():
+ input_data = torch.from_numpy(self.valid_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+
+ batch_size = self.args.valid_batch_size
+ for b_begin in tqdm(range(0, actual_examples.shape[0], batch_size)):
+
+ input_batch = actual_examples[b_begin: b_begin + batch_size]
+ input_batch = input_batch.to(self.device)
+ 
+ s,r,o = input_batch[:,0], input_batch[:,1], input_batch[:,2]
+ 
+ emb_s = self.model.emb_e(s).squeeze(dim=1)
+ emb_r = self.model.emb_rel(r).squeeze(dim=1)
+ emb_o = self.model.emb_e(o).squeeze(dim=1)
+ 
+ if self.args.add_reciprocals:
+ r_rev = r + self.n_rel
+ emb_rrev = self.model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+ 
+ pred_sr = self.model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = self.model.loss(pred_sr, o) # cross entropy loss
+ 
+ pred_or = self.model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = self.model.loss(pred_or, s)
+ 
+ total_loss = loss_sr + loss_or
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 3):
+ #self.logger.info('Computing regularizer weight')
+ if self.args.model == 'complex':
+ emb_dim = self.args.embedding_dim #int(self.args.embedding_dim/2)
+ lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
+ rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
+ rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
+ rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])
+ 
+ #print(lhs[0].shape, lhs[1].shape)
+ factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ else:
+ factors_sr = (emb_s, emb_r, emb_o)
+ factors_or = (emb_s, emb_rrev, emb_o)
+ 
+ total_loss += self.n3_regularizer(factors_sr)
+ total_loss += self.n3_regularizer(factors_or)
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 2):
+ total_loss += self.lp_regularizer()
+
+ losses.append(total_loss.item())
+
+ loss = np.mean(losses)
+ self.logger.info('[Epoch:{}]: Validating Loss:{:.6}\n'.format(epoch, loss))
+ return loss
+
+
+ def run_epoch(self, epoch):
+ self.model.train()
+ losses = []
+ 
+ #shuffle the train dataset
+ input_data = torch.from_numpy(self.train_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+ 
+ batch_size = self.args.train_batch_size
+ 
+ for b_begin in tqdm(range(0, actual_examples.shape[0], batch_size)):
+ self.optimizer.zero_grad()
+ input_batch = actual_examples[b_begin: b_begin + batch_size]
+ input_batch = input_batch.to(self.device)
+ 
+ s,r,o = input_batch[:,0], input_batch[:,1], input_batch[:,2]
+ 
+ emb_s = self.model.emb_e(s).squeeze(dim=1)
+ emb_r = self.model.emb_rel(r).squeeze(dim=1)
+ emb_o = self.model.emb_e(o).squeeze(dim=1)
+ 
+ if self.args.add_reciprocals:
+ r_rev = r + self.n_rel
+ emb_rrev = self.model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+ 
+ pred_sr = self.model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = self.model.loss(pred_sr, o) # loss is cross entropy loss
+ 
+ pred_or = self.model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = self.model.loss(pred_or, s)
+ 
+ total_loss = loss_sr + loss_or
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 3):
+ #self.logger.info('Computing regularizer weight')
+ if self.args.model == 'complex':
+ emb_dim = self.args.embedding_dim #int(self.args.embedding_dim/2)
+ lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
+ rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
+ rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
+ rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])
+ 
+ #print(lhs[0].shape, lhs[1].shape)
+ factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2))
+ factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2))
+ else:
+ factors_sr = (emb_s, emb_r, emb_o)
+ factors_or = (emb_s, emb_rrev, emb_o)
+ 
+ total_loss += self.n3_regularizer(factors_sr)
+ total_loss += self.n3_regularizer(factors_or)
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 2):
+ total_loss += self.lp_regularizer()
+ 
+ 
+ total_loss.backward()
+ self.optimizer.step()
+ losses.append(total_loss.item())
+ 
+ if self.args.save_influence_map: #for gradient rollback
+ with torch.no_grad():
+ prev_emb_e = self.previous_weights[0]
+ prev_emb_rel = self.previous_weights[1]
+ # need to compute the influence value per-triple
+ for idx in range(input_batch.shape[0]):
+ head, rel, tail = s[idx], r[idx], o[idx]
+ inf_head = (emb_s[idx] - prev_emb_e[head]).cpu().detach().numpy()
+ inf_tail = (emb_o[idx] - prev_emb_e[tail]).cpu().detach().numpy()
+ inf_rel = (emb_r[idx] - prev_emb_rel[rel]).cpu().detach().numpy()
+ #print(inf_head.shape, inf_tail.shape, inf_rel.shape)
+
+ #write the influences to dictionary
+ key_trip = '{0}_{1}_{2}'.format(head.item(), rel.item(), tail.item())
+ key = '{0}_s'.format(key_trip)
+ self.influence_map[key] += inf_head
+ #self.logger.info('Written to influence map. Key: {0}, Value shape: {1}'.format(key, inf_head.shape))
+ key = '{0}_r'.format(key_trip)
+ self.influence_map[key] += inf_rel
+ key = '{0}_o'.format(key_trip)
+ self.influence_map[key] += inf_tail
+
+ # update the previous weights to be tracked
+ self.previous_weights = [copy.deepcopy(param) for param in self.model.parameters()]
+ 
+ if (b_begin%5000 == 0) or (b_begin== (actual_examples.shape[0]-1)):
+ self.logger.info('[E:{} | {}]: Train Loss:{:.6}'.format(epoch, b_begin, np.mean(losses)))
+ 
+ loss = np.mean(losses)
+ self.logger.info('[Epoch:{}]: Training Loss:{:.6}\n'.format(epoch, loss))
+ return loss
+ 
+ def fit(self):
+ self.model.init()
+ self.logger.info(self.model)
+ 
+ self.logger.info('------ Start the model training ------')
+ start_time = time.time()
+ self.logger.info('Start time: {0}'.format(str(start_time)))
+ 
+ 
+ train_losses = []
+ valid_losses = []
+ best_val = 10000000000.
+ for epoch in range(self.args.epochs):
+
+ print("="*15,'epoch:',epoch,'='*15)
+ train_loss = self.run_epoch(epoch)
+ train_losses.append(train_loss)
+
+ if train_loss < best_val:
+ best_val = train_loss
+ self.save_model()
+ print("Train loss: {0}, Best loss: {1}\n\n".format(train_loss, best_val))
+
+ 
+ with open(self.loss_path, "wb") as fl: 
+ pkl.dump({"train loss":train_losses, "valid loss":valid_losses}, fl)
+ self.logger.info('Time taken to train the model: {0}'.format(str(time.time() - start_time)))
+ start_time = time.time()
+ 
+ if self.args.save_influence_map: #save the influence map
+ with open(self.influence_path, "wb") as fl: #Pickling
+ pkl.dump(self.get_influence_map(), fl)
+ self.logger.info('Finished saving influence map')
+ self.logger.info('Time taken to save the influence map: {0}'.format(str(time.time() - start_time)))
+ 
+#%%
+parser = utils.get_argument_parser()
+
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+
+utils.seed_all(args.seed)
+np.set_printoptions(precision=5)
+cudnn.benchmark = False
+
+model = Main(args)
+model.fit()

DiseaseSpecific/main_multiprocess.py

@@ -0,0 +1,391 @@
+"""Multiprocess for main.py"""
+#%%
+import pickle as pkl
+from typing import Dict, Tuple, List
+import os
+import numpy as np
+import json
+import logging
+import argparse 
+import math
+from pprint import pprint
+import pandas as pd
+from collections import defaultdict
+import copy
+import time
+
+import torch
+from torch.utils.data import DataLoader
+import torch.backends.cudnn as cudnn
+import torch.autograd as autograd
+
+from model import Distmult, Complex, Conve
+import utils
+
+# from evaluation import evaluation
+
+#%%
+class Main(object):
+ def __init__(self, args):
+ self.args = args 
+ 
+ self.model_name = '{0}_{1}_{2}_{3}_{4}_{5}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop, args.thread_name)
+ #leaving batches from the model_name since they do not depend on model_architecture 
+ # also leaving kernel size and filters, siinice don't intend to change those
+ self.model_path = 'saved_models/evaluation/{0}_{1}.model'.format(args.data, self.model_name)
+ 
+ self.log_path = 'logs/evaluation_logs/{0}_{1}_{2}_{3}_{4}.log'.format(args.data, self.model_name, args.epochs, args.train_batch_size, args.thread_name)
+ self.loss_path = 'losses/evaluation_losses/{0}_{1}_{2}_{3}_{4}.pickle'.format(args.data, self.model_name, args.epochs, args.train_batch_size, args.thread_name)
+ 
+ logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+ datefmt = '%m/%d/%Y %H:%M:%S',
+ level = logging.INFO,
+ filename = self.log_path)
+ self.logger = logging.getLogger(__name__)
+ self.logger.info(vars(self.args))
+ self.logger.info('\n') 
+ 
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ 
+ self.load_data()
+ self.model = self.add_model()
+ self.optimizer = self.add_optimizer(self.model.parameters())
+ 
+ if self.args.save_influence_map:
+ self.logger.info('-------- Argument save_influence_map is set. Will use GR to compute and save influence maps ----------\n')
+ # when we want to save influence during training
+ self.args.add_reciprocals = False # to keep things simple
+ # init an empty influence map
+ self.influence_map = defaultdict(float)
+ #self.influence_path = 'influence_maps/{0}_{1}.json'.format(args.data, self.model_name)
+ self.influence_path = 'influence_maps/{0}_{1}.pickle'.format(args.data, self.model_name)
+ # Initialize a copy of the model prams to track previous weights in an epoch
+ self.previous_weights = [copy.deepcopy(param) for param in self.model.parameters()]
+ self.logger.info('Shape for previous weights: {}, {}'.format(self.previous_weights[0].shape, self.previous_weights[1].shape))
+ 
+ def load_data(self):
+ ''' 
+ Load the train, valid datasets
+ '''
+ data_path = os.path.join('processed_data', self.args.data)
+
+ n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+ self.n_ent = n_ent
+ self.n_rel = n_rel
+ 
+ self.train_data = utils.load_data(os.path.join(data_path, 'evaluation', 'all_{}.txt'.format(self.args.thread_name)))
+ 
+ self.valid_data= self.train_data[-100:, :].copy()
+ # self.train_data = utils.load_data()
+ 
+ def add_model(self):
+ 
+ if self.args.model is None:
+ model = Distmult(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'distmult':
+ model = Distmult(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'complex':
+ model = Complex(self.args, self.n_ent, self.n_rel)
+ elif self.args.model == 'conve':
+ model = Conve(self.args, self.n_ent, self.n_rel)
+ else:
+ self.logger.info('Unknown model: {0}', self.args.model)
+ raise Exception("Unknown model!")
+ model.to(self.device)
+ return model
+ 
+ def add_optimizer(self, parameters):
+ #if self.args.optimizer == 'adam' : return torch.optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.lr_decay)
+ #else : return torch.optim.SGD(parameters, lr=self.args.lr, weight_decay=self.args.lr_decay)
+ return torch.optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.lr_decay)
+ 
+ def save_model(self):
+ state = {
+ 'state_dict': self.model.state_dict(),
+ 'optimizer': self.optimizer.state_dict(),
+ 'args': vars(self.args)
+ }
+ torch.save(state, self.model_path)
+ self.logger.info('Saving model to {0}'.format(self.model_path))
+ 
+ def load_model(self):
+ self.logger.info('Loading saved model from {0}'.format(self.model_path))
+ state = torch.load(self.model_path)
+ model_params = state['state_dict']
+ params = [(key, value.size(), value.numel()) for key, value in model_params.items()]
+ for key, size, count in params:
+ self.logger.info(key, size, count)
+ self.model.load_state_dict(model_params)
+ self.optimizer.load_state_dict(state['optimizer'])
+ 
+ def lp_regularizer(self):
+ # Apply p-norm regularization; assign weights to each param
+ weight = self.args.reg_weight
+ p = self.args.reg_norm
+ 
+ trainable_params = [self.model.emb_e.weight, self.model.emb_rel.weight]
+ norm = 0
+ for i in range(len(trainable_params)):
+ #norm += weight * trainable_params[i].norm(p = p)**p
+ norm += weight * torch.sum( torch.abs(trainable_params[i]) ** p)
+ 
+ return norm
+ 
+ def n3_regularizer(self, factors):
+ # factors are the embeddings for lhs, rel, rhs for triples in a batch
+ weight = self.args.reg_weight
+ p = self.args.reg_norm
+ 
+ norm = 0
+ for f in factors:
+ norm += weight * torch.sum(torch.abs(f) ** p)
+ 
+ return norm / factors[0].shape[0] # scale by number of triples in batch
+ 
+ def get_influence_map(self):
+ """
+ Turns the influence map into a list, ready to be written to disc. (before: numpy)
+ :return: the influence map with lists as values
+ """
+ assert self.args.save_influence_map == True
+ 
+ for key in self.influence_map:
+ self.influence_map[key] = self.influence_map[key].tolist()
+ #self.logger.info('get_influence_map passed')
+ return self.influence_map
+
+ def evaluate(self, split, batch_size, epoch):
+ """
+ The same as self.run_epoch()
+ """
+
+ self.model.eval()
+ losses = []
+
+ with torch.no_grad():
+ input_data = torch.from_numpy(self.valid_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+
+ batch_size = self.args.valid_batch_size
+ for b_begin in range(0, actual_examples.shape[0], batch_size):
+
+ input_batch = actual_examples[b_begin: b_begin + batch_size]
+ input_batch = input_batch.to(self.device)
+ 
+ s,r,o = input_batch[:,0], input_batch[:,1], input_batch[:,2]
+ 
+ emb_s = self.model.emb_e(s).squeeze(dim=1)
+ emb_r = self.model.emb_rel(r).squeeze(dim=1)
+ emb_o = self.model.emb_e(o).squeeze(dim=1)
+ 
+ if self.args.add_reciprocals:
+ r_rev = r + self.n_rel
+ emb_rrev = self.model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+ 
+ pred_sr = self.model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = self.model.loss(pred_sr, o) # cross entropy loss
+ 
+ pred_or = self.model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = self.model.loss(pred_or, s)
+ 
+ total_loss = loss_sr + loss_or
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 3):
+ #self.logger.info('Computing regularizer weight')
+ if self.args.model == 'complex':
+ emb_dim = self.args.embedding_dim #int(self.args.embedding_dim/2)
+ lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
+ rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
+ rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
+ rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])
+ 
+ #print(lhs[0].shape, lhs[1].shape)
+ factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
+ )
+ else:
+ factors_sr = (emb_s, emb_r, emb_o)
+ factors_or = (emb_s, emb_rrev, emb_o)
+ 
+ total_loss += self.n3_regularizer(factors_sr)
+ total_loss += self.n3_regularizer(factors_or)
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 2):
+ total_loss += self.lp_regularizer()
+
+ losses.append(total_loss.item())
+
+ loss = np.mean(losses)
+ self.logger.info('[Epoch:{}]: Validating Loss:{:.6}\n'.format(epoch, loss))
+ return loss
+
+
+ def run_epoch(self, epoch):
+ self.model.train()
+ losses = []
+ 
+ #shuffle the train dataset
+ input_data = torch.from_numpy(self.train_data.astype('int64'))
+ actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
+ del input_data
+ 
+ batch_size = self.args.train_batch_size
+ 
+ for b_begin in range(0, actual_examples.shape[0], batch_size):
+ self.optimizer.zero_grad()
+ input_batch = actual_examples[b_begin: b_begin + batch_size]
+ input_batch = input_batch.to(self.device)
+ 
+ s,r,o = input_batch[:,0], input_batch[:,1], input_batch[:,2]
+ 
+ emb_s = self.model.emb_e(s).squeeze(dim=1)
+ emb_r = self.model.emb_rel(r).squeeze(dim=1)
+ emb_o = self.model.emb_e(o).squeeze(dim=1)
+ 
+ if self.args.add_reciprocals:
+ r_rev = r + self.n_rel
+ emb_rrev = self.model.emb_rel(r_rev).squeeze(dim=1)
+ else:
+ r_rev = r
+ emb_rrev = emb_r
+ 
+ pred_sr = self.model.forward(emb_s, emb_r, mode='rhs')
+ loss_sr = self.model.loss(pred_sr, o) # loss is cross entropy loss
+ 
+ pred_or = self.model.forward(emb_o, emb_rrev, mode='lhs')
+ loss_or = self.model.loss(pred_or, s)
+ 
+ total_loss = loss_sr + loss_or
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 3):
+ #self.logger.info('Computing regularizer weight')
+ if self.args.model == 'complex':
+ emb_dim = self.args.embedding_dim #int(self.args.embedding_dim/2)
+ lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
+ rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
+ rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
+ rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])
+ 
+ #print(lhs[0].shape, lhs[1].shape)
+ factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2))
+ factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
+ torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
+ torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2))
+ else:
+ factors_sr = (emb_s, emb_r, emb_o)
+ factors_or = (emb_s, emb_rrev, emb_o)
+ 
+ total_loss += self.n3_regularizer(factors_sr)
+ total_loss += self.n3_regularizer(factors_or)
+ 
+ if (self.args.reg_weight != 0.0 and self.args.reg_norm == 2):
+ total_loss += self.lp_regularizer()
+ 
+ 
+ total_loss.backward()
+ self.optimizer.step()
+ losses.append(total_loss.item())
+ 
+ if self.args.save_influence_map: #for gradient rollback
+ with torch.no_grad():
+ prev_emb_e = self.previous_weights[0]
+ prev_emb_rel = self.previous_weights[1]
+ # need to compute the influence value per-triple
+ for idx in range(input_batch.shape[0]):
+ head, rel, tail = s[idx], r[idx], o[idx]
+ inf_head = (emb_s[idx] - prev_emb_e[head]).cpu().detach().numpy()
+ inf_tail = (emb_o[idx] - prev_emb_e[tail]).cpu().detach().numpy()
+ inf_rel = (emb_r[idx] - prev_emb_rel[rel]).cpu().detach().numpy()
+ #print(inf_head.shape, inf_tail.shape, inf_rel.shape)
+
+ #write the influences to dictionary
+ key_trip = '{0}_{1}_{2}'.format(head.item(), rel.item(), tail.item())
+ key = '{0}_s'.format(key_trip)
+ self.influence_map[key] += inf_head
+ #self.logger.info('Written to influence map. Key: {0}, Value shape: {1}'.format(key, inf_head.shape))
+ key = '{0}_r'.format(key_trip)
+ self.influence_map[key] += inf_rel
+ key = '{0}_o'.format(key_trip)
+ self.influence_map[key] += inf_tail
+
+ # update the previous weights to be tracked
+ self.previous_weights = [copy.deepcopy(param) for param in self.model.parameters()]
+ 
+ if (b_begin%5000 == 0) or (b_begin== (actual_examples.shape[0]-1)):
+ self.logger.info('[E:{} | {}]: Train Loss:{:.6}'.format(epoch, b_begin, np.mean(losses)))
+ 
+ loss = np.mean(losses)
+ self.logger.info('[Epoch:{}]: Training Loss:{:.6}\n'.format(epoch, loss))
+ return loss
+ 
+ def fit(self):
+ # if self.args.resume:
+ # self.load_model()
+ # results = self.evaluate(split=self.args.resume_split, batch_size = self.args.test_batch_size, epoch = -1)
+ # pprint(results)
+ 
+ # else:
+ self.model.init()
+ self.logger.info(self.model)
+ 
+ self.logger.info('------ Start the model training ------')
+ start_time = time.time()
+ self.logger.info('Start time: {0}'.format(str(start_time)))
+ 
+ 
+ train_losses = []
+ valid_losses = []
+ best_val = 10000000000.
+ for epoch in range(self.args.epochs):
+
+ train_loss = self.run_epoch(epoch)
+ train_losses.append(train_loss)
+ 
+ # Don't use valid_data here !!!!!!!!!
+
+ # valid_loss = self.evaluate(split='valid', batch_size = self.args.valid_batch_size, epoch = epoch)
+ # valid_losses.append(valid_loss)
+ # results_test = self.evaluate(split='test', batch_size = self.args.test_batch_size, epoch = epoch)
+ if train_loss < best_val:
+ best_val = train_loss
+ self.save_model()
+ self.logger.info("Train loss: {0}, Best loss: {1}\n\n".format(train_loss, best_val))
+ # print("Valid loss: {0}, Best loss: {1}\n\n".format(valid_loss, best_val))
+ 
+ with open(self.loss_path, "wb") as fl: 
+ pkl.dump({"train loss":train_losses, "valid loss":valid_losses}, fl)
+ self.logger.info('Time taken to train the model: {0}'.format(str(time.time() - start_time)))
+ start_time = time.time()
+ 
+ if self.args.save_influence_map: #save the influence map
+ with open(self.influence_path, "wb") as fl: #Pickling
+ pkl.dump(self.get_influence_map(), fl)
+ self.logger.info('Finished saving influence map')
+ self.logger.info('Time taken to save the influence map: {0}'.format(str(time.time() - start_time)))
+ 
+#%%
+parser = utils.get_argument_parser()
+parser.add_argument('--thread-name', type = str, required=True, help = "This parameter will be automatically determined.")
+
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+# if args.reproduce_results:
+# args = utils.set_hyperparams(args)
+
+utils.seed_all(args.seed)
+np.set_printoptions(precision=5)
+cudnn.benchmark = False
+
+model = Main(args)
+model.fit()

DiseaseSpecific/model.py

@@ -0,0 +1,504 @@
+import torch
+from torch.nn import functional as F, Parameter
+from torch.autograd import Variable
+from torch.nn.init import xavier_normal_, xavier_uniform_
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+
+class Distmult(torch.nn.Module):
+ def __init__(self, args, num_entities, num_relations):
+ super(Distmult, self).__init__()
+ 
+ if args.max_norm:
+ self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, max_norm=1.0)
+ self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim)
+ else:
+ self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, padding_idx=None)
+ self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim, padding_idx=None)
+ 
+ self.inp_drop = torch.nn.Dropout(args.input_drop)
+ self.loss = torch.nn.CrossEntropyLoss()
+ 
+ self.init()
+ 
+ def init(self):
+ xavier_normal_(self.emb_e.weight)
+ xavier_normal_(self.emb_rel.weight)
+ 
+ def score_sr(self, sub, rel, sigmoid = False):
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ 
+ #sub_emb = self.inp_drop(sub_emb)
+ #rel_emb = self.inp_drop(rel_emb) 
+ 
+ pred = torch.mm(sub_emb*rel_emb, self.emb_e.weight.transpose(1,0))
+ if sigmoid:
+ pred = torch.sigmoid(pred) 
+ return pred
+ 
+ def score_or(self, obj, rel, sigmoid = False):
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ 
+ #obj_emb = self.inp_drop(obj_emb)
+ #rel_emb = self.inp_drop(rel_emb) 
+ 
+ pred = torch.mm(obj_emb*rel_emb, self.emb_e.weight.transpose(1,0))
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+ return pred
+ 
+ 
+ def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False):
+ '''
+ When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r)
+ For distmult, computations for both modes are equivalent, so we do not need if-else block
+ '''
+ sub_emb = self.inp_drop(sub_emb)
+ rel_emb = self.inp_drop(rel_emb) 
+ 
+ pred = torch.mm(sub_emb*rel_emb, self.emb_e.weight.transpose(1,0))
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - score
+ '''
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ 
+ pred = torch.sum(sub_emb*rel_emb*obj_emb, dim=-1)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False):
+ '''
+ Inputs - embeddings of subject, relation, object
+ Return - score
+ '''
+ pred = torch.sum(emb_s*emb_r*emb_o, dim=-1)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples_vec(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - a vector score for the triple instead of reducing over the embedding dimension
+ '''
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ 
+ pred = sub_emb*rel_emb*obj_emb
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+class Complex(torch.nn.Module):
+ def __init__(self, args, num_entities, num_relations):
+ super(Complex, self).__init__()
+ 
+ if args.max_norm:
+ self.emb_e = torch.nn.Embedding(num_entities, 2*args.embedding_dim, max_norm=1.0)
+ self.emb_rel = torch.nn.Embedding(num_relations, 2*args.embedding_dim)
+ else:
+ self.emb_e = torch.nn.Embedding(num_entities, 2*args.embedding_dim, padding_idx=None)
+ self.emb_rel = torch.nn.Embedding(num_relations, 2*args.embedding_dim, padding_idx=None)
+ 
+ self.inp_drop = torch.nn.Dropout(args.input_drop)
+ self.loss = torch.nn.CrossEntropyLoss()
+ 
+ self.init()
+ 
+ def init(self):
+ xavier_normal_(self.emb_e.weight)
+ xavier_normal_(self.emb_rel.weight)
+ 
+ def score_sr(self, sub, rel, sigmoid = False):
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ 
+ s_real, s_img = torch.chunk(rel_emb, 2, dim=-1)
+ rel_real, rel_img = torch.chunk(sub_emb, 2, dim=-1)
+ emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1)
+ 
+ realo_realreal = s_real*rel_real
+ realo_imgimg = s_img*rel_img
+ realo = realo_realreal - realo_imgimg
+ real = torch.mm(realo, emb_e_real.transpose(1,0))
+ 
+ imgo_realimg = s_real*rel_img
+ imgo_imgreal = s_img*rel_real
+ imgo = imgo_realimg + imgo_imgreal
+ img = torch.mm(imgo, emb_e_img.transpose(1,0))
+ 
+ pred = real + img
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred) 
+ return pred
+ 
+ 
+ def score_or(self, obj, rel, sigmoid = False):
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ 
+ rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1)
+ o_real, o_img = torch.chunk(obj_emb, 2, dim=-1)
+ emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1)
+
+ #rel_real = self.inp_drop(rel_real)
+ #rel_img = self.inp_drop(rel_img)
+ #o_real = self.inp_drop(o_real)
+ #o_img = self.inp_drop(o_img)
+
+ # complex space bilinear product (equivalent to HolE)
+# realrealreal = torch.mm(rel_real*o_real, emb_e_real.transpose(1,0))
+# realimgimg = torch.mm(rel_img*o_img, emb_e_real.transpose(1,0))
+# imgrealimg = torch.mm(rel_real*o_img, emb_e_img.transpose(1,0))
+# imgimgreal = torch.mm(rel_img*o_real, emb_e_img.transpose(1,0))
+# pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+ 
+ reals_realreal = rel_real*o_real
+ reals_imgimg = rel_img*o_img
+ reals = reals_realreal + reals_imgimg
+ real = torch.mm(reals, emb_e_real.transpose(1,0))
+ 
+ imgs_realimg = rel_real*o_img
+ imgs_imgreal = rel_img*o_real
+ imgs = imgs_realimg - imgs_imgreal
+ img = torch.mm(imgs, emb_e_img.transpose(1,0))
+ 
+ pred = real + img
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+ return pred
+ 
+ 
+ def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False):
+ '''
+ When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r)
+ 
+ '''
+ if mode == 'lhs':
+ rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1)
+ o_real, o_img = torch.chunk(sub_emb, 2, dim=-1)
+ emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1)
+ 
+ rel_real = self.inp_drop(rel_real)
+ rel_img = self.inp_drop(rel_img)
+ o_real = self.inp_drop(o_real)
+ o_img = self.inp_drop(o_img)
+ 
+ reals_realreal = rel_real*o_real
+ reals_imgimg = rel_img*o_img
+ reals = reals_realreal + reals_imgimg
+ real = torch.mm(reals, emb_e_real.transpose(1,0))
+
+ imgs_realimg = rel_real*o_img
+ imgs_imgreal = rel_img*o_real
+ imgs = imgs_realimg - imgs_imgreal
+ img = torch.mm(imgs, emb_e_img.transpose(1,0))
+
+ pred = real + img
+ 
+ else:
+ s_real, s_img = torch.chunk(rel_emb, 2, dim=-1)
+ rel_real, rel_img = torch.chunk(sub_emb, 2, dim=-1)
+ emb_e_real, emb_e_img = torch.chunk(self.emb_e.weight, 2, dim=-1)
+ 
+ s_real = self.inp_drop(s_real)
+ s_img = self.inp_drop(s_img)
+ rel_real = self.inp_drop(rel_real)
+ rel_img = self.inp_drop(rel_img)
+ 
+ realo_realreal = s_real*rel_real
+ realo_imgimg = s_img*rel_img
+ realo = realo_realreal - realo_imgimg
+ real = torch.mm(realo, emb_e_real.transpose(1,0))
+
+ imgo_realimg = s_real*rel_img
+ imgo_imgreal = s_img*rel_real
+ imgo = imgo_realimg + imgo_imgreal
+ img = torch.mm(imgo, emb_e_img.transpose(1,0))
+
+ pred = real + img
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - score
+ '''
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ 
+ s_real, s_img = torch.chunk(sub_emb, 2, dim=-1)
+ rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1)
+ o_real, o_img = torch.chunk(obj_emb, 2, dim=-1)
+ 
+ realrealreal = torch.sum(s_real*rel_real*o_real, dim=-1)
+ realimgimg = torch.sum(s_real*rel_img*o_img, axis=-1)
+ imgrealimg = torch.sum(s_img*rel_real*o_img, axis=-1)
+ imgimgreal = torch.sum(s_img*rel_img*o_real, axis=-1)
+ 
+ pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False):
+ '''
+ Inputs - embeddings of subject, relation, object
+ Return - score
+ '''
+ 
+ s_real, s_img = torch.chunk(emb_s, 2, dim=-1)
+ rel_real, rel_img = torch.chunk(emb_r, 2, dim=-1)
+ o_real, o_img = torch.chunk(emb_o, 2, dim=-1)
+ 
+ realrealreal = torch.sum(s_real*rel_real*o_real, dim=-1)
+ realimgimg = torch.sum(s_real*rel_img*o_img, axis=-1)
+ imgrealimg = torch.sum(s_img*rel_real*o_img, axis=-1)
+ imgimgreal = torch.sum(s_img*rel_img*o_real, axis=-1)
+ 
+ pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples_vec(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - a vector score for the triple instead of reducing over the embedding dimension
+ '''
+ sub_emb = self.emb_e(sub).squeeze(dim=1)
+ rel_emb = self.emb_rel(rel).squeeze(dim=1)
+ obj_emb = self.emb_e(obj).squeeze(dim=1)
+ 
+ s_real, s_img = torch.chunk(sub_emb, 2, dim=-1)
+ rel_real, rel_img = torch.chunk(rel_emb, 2, dim=-1)
+ o_real, o_img = torch.chunk(obj_emb, 2, dim=-1)
+ 
+ realrealreal = s_real*rel_real*o_real
+ realimgimg = s_real*rel_img*o_img
+ imgrealimg = s_img*rel_real*o_img
+ imgimgreal = s_img*rel_img*o_real
+ 
+ pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+class Conve(torch.nn.Module):
+
+ #Too slow !!!!
+
+ def __init__(self, args, num_entities, num_relations):
+ super(Conve, self).__init__()
+ 
+ if args.max_norm:
+ self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, max_norm=1.0)
+ self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim)
+ else:
+ self.emb_e = torch.nn.Embedding(num_entities, args.embedding_dim, padding_idx=None)
+ self.emb_rel = torch.nn.Embedding(num_relations, args.embedding_dim, padding_idx=None)
+ 
+ self.inp_drop = torch.nn.Dropout(args.input_drop)
+ self.hidden_drop = torch.nn.Dropout(args.hidden_drop)
+ self.feature_drop = torch.nn.Dropout2d(args.feat_drop)
+ 
+ self.embedding_dim = args.embedding_dim #default is 200
+ self.num_filters = args.num_filters # default is 32
+ self.kernel_size = args.kernel_size # default is 3
+ self.stack_width = args.stack_width # default is 20
+ self.stack_height = args.embedding_dim // self.stack_width
+ 
+ self.bn0 = torch.nn.BatchNorm2d(1)
+ self.bn1 = torch.nn.BatchNorm2d(self.num_filters)
+ self.bn2 = torch.nn.BatchNorm1d(args.embedding_dim)
+
+ self.conv1 = torch.nn.Conv2d(1, out_channels=self.num_filters, 
+ kernel_size=(self.kernel_size, self.kernel_size), 
+ stride=1, padding=0, bias=args.use_bias)
+ #self.conv1 = torch.nn.Conv2d(1, 32, (3, 3), 1, 0, bias=args.use_bias) # <-- default
+ 
+ flat_sz_h = int(2*self.stack_width) - self.kernel_size + 1
+ flat_sz_w = self.stack_height - self.kernel_size + 1
+ self.flat_sz = flat_sz_h*flat_sz_w*self.num_filters
+ self.fc = torch.nn.Linear(self.flat_sz, args.embedding_dim)
+ 
+ self.register_parameter('b', Parameter(torch.zeros(num_entities)))
+ self.loss = torch.nn.CrossEntropyLoss()
+
+ self.init()
+ 
+ def init(self):
+ xavier_normal_(self.emb_e.weight)
+ xavier_normal_(self.emb_rel.weight)
+ 
+ def concat(self, e1_embed, rel_embed, form='plain'):
+ if form == 'plain':
+ e1_embed = e1_embed. view(-1, 1, self.stack_width, self.stack_height)
+ rel_embed = rel_embed.view(-1, 1, self.stack_width, self.stack_height)
+ stack_inp = torch.cat([e1_embed, rel_embed], 2)
+
+ elif form == 'alternate':
+ e1_embed = e1_embed. view(-1, 1, self.embedding_dim)
+ rel_embed = rel_embed.view(-1, 1, self.embedding_dim)
+ stack_inp = torch.cat([e1_embed, rel_embed], 1)
+ stack_inp = torch.transpose(stack_inp, 2, 1).reshape((-1, 1, 2*self.stack_width, self.stack_height))
+
+ else: raise NotImplementedError
+ return stack_inp
+ 
+ def conve_architecture(self, sub_emb, rel_emb):
+ stacked_inputs = self.concat(sub_emb, rel_emb)
+ stacked_inputs = self.bn0(stacked_inputs)
+ x = self.inp_drop(stacked_inputs)
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = F.relu(x)
+ x = self.feature_drop(x)
+ #x = x.view(x.shape[0], -1)
+ x = x.view(-1, self.flat_sz)
+ x = self.fc(x)
+ x = self.hidden_drop(x)
+ x = self.bn2(x)
+ x = F.relu(x)
+ 
+ return x
+ 
+ def score_sr(self, sub, rel, sigmoid = False):
+ sub_emb = self.emb_e(sub)
+ rel_emb = self.emb_rel(rel)
+ 
+ x = self.conve_architecture(sub_emb, rel_emb)
+ 
+ pred = torch.mm(x, self.emb_e.weight.transpose(1,0))
+ pred += self.b.expand_as(pred) 
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred) 
+ return pred
+ 
+ def score_or(self, obj, rel, sigmoid = False):
+ obj_emb = self.emb_e(obj)
+ rel_emb = self.emb_rel(rel)
+ 
+ x = self.conve_architecture(obj_emb, rel_emb)
+ pred = torch.mm(x, self.emb_e.weight.transpose(1,0))
+ pred += self.b.expand_as(pred)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+ return pred
+ 
+ 
+ def forward(self, sub_emb, rel_emb, mode='rhs', sigmoid=False):
+ '''
+ When mode is 'rhs' we expect (s,r); for 'lhs', we expect (o,r)
+ For conve, computations for both modes are equivalent, so we do not need if-else block
+ '''
+ x = self.conve_architecture(sub_emb, rel_emb)
+ 
+ pred = torch.mm(x, self.emb_e.weight.transpose(1,0))
+ pred += self.b.expand_as(pred)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - score
+ '''
+ sub_emb = self.emb_e(sub)
+ rel_emb = self.emb_rel(rel)
+ obj_emb = self.emb_e(obj)
+ x = self.conve_architecture(sub_emb, rel_emb)
+ 
+ pred = torch.mm(x, obj_emb.transpose(1,0))
+ #print(pred.shape)
+ pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding
+ # above works fine for single input triples; 
+ # but if input is batch of triples, then this is a matrix of (num_trip x num_trip) where diagonal is scores
+ # so use torch.diagonal() after calling this function
+ pred = torch.diagonal(pred)
+ # or could have used : pred= torch.sum(x*obj_emb, dim=-1)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_emb(self, emb_s, emb_r, emb_o, sigmoid=False):
+ '''
+ Inputs - embeddings of subject, relation, object
+ Return - score
+ '''
+ x = self.conve_architecture(emb_s, emb_r)
+ 
+ pred = torch.mm(x, emb_o.transpose(1,0))
+ #pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding - don't know which obj
+ # above works fine for single input triples; 
+ # but if input is batch of triples, then this is a matrix of (num_trip x num_trip) where diagonal is scores
+ # so use torch.diagonal() after calling this function
+ pred = torch.diagonal(pred)
+ # or could have used : pred= torch.sum(x*obj_emb, dim=-1)
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred
+ 
+ def score_triples_vec(self, sub, rel, obj, sigmoid=False):
+ '''
+ Inputs - subject, relation, object
+ Return - a vector score for the triple instead of reducing over the embedding dimension
+ '''
+ sub_emb = self.emb_e(sub)
+ rel_emb = self.emb_rel(rel)
+ obj_emb = self.emb_e(obj)
+ 
+ x = self.conve_architecture(sub_emb, rel_emb)
+ 
+ #pred = torch.mm(x, obj_emb.transpose(1,0))
+ pred = x*obj_emb
+ #print(pred.shape, self.b[obj].shape) #shapes are [7,200] and [7]
+ #pred += self.b[obj].expand_as(pred) #taking the bias value for object embedding - can't add scalar to vector
+ 
+ #pred = sub_emb*rel_emb*obj_emb
+ 
+ if sigmoid:
+ pred = torch.sigmoid(pred)
+
+ return pred

DiseaseSpecific/utils.py

@@ -0,0 +1,195 @@
+'''
+A file modified on https://github.com/PeruBhardwaj/AttributionAttack/blob/main/KGEAttack/ConvE/utils.py
+'''
+#%%
+import logging
+import time
+from tqdm import tqdm
+import io
+import pandas as pd
+import numpy as np
+import os
+import json
+
+import argparse
+import torch
+import random
+
+from yaml import parse
+
+from model import Conve, Distmult, Complex
+
+logger = logging.getLogger(__name__)
+#%%
+def generate_dicts(data_path):
+ with open (os.path.join(data_path, 'entities_dict.json'), 'r') as f:
+ ent_to_id = json.load(f)
+ with open (os.path.join(data_path, 'relations_dict.json'), 'r') as f:
+ rel_to_id = json.load(f)
+ n_ent = len(list(ent_to_id.keys()))
+ n_rel = len(list(rel_to_id.keys()))
+ 
+ return n_ent, n_rel, ent_to_id, rel_to_id
+
+def save_data(file_name, data):
+ with open(file_name, 'w') as fl:
+ for item in data:
+ fl.write("%s\n" % "\t".join(map(str, item)))
+
+def load_data(file_name, drop = True):
+ df = pd.read_csv(file_name, sep='\t', header=None, names=None, dtype=str)
+ if drop:
+ df = df.drop_duplicates()
+ else:
+ pass
+ return df.values
+
+def seed_all(seed=1):
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ os.environ['PYTHONHASHSEED'] = str(seed)
+ torch.backends.cudnn.deterministic = True
+
+def add_model(args, n_ent, n_rel):
+ if args.model is None:
+ model = Distmult(args, n_ent, n_rel)
+ elif args.model == 'distmult':
+ model = Distmult(args, n_ent, n_rel)
+ elif args.model == 'complex':
+ model = Complex(args, n_ent, n_rel)
+ elif args.model == 'conve':
+ model = Conve(args, n_ent, n_rel)
+ else:
+ raise Exception("Unknown model!")
+
+ return model
+
+def load_model(model_path, args, n_ent, n_rel, device):
+ # add a model and load the pre-trained params
+ model = add_model(args, n_ent, n_rel)
+ model.to(device)
+ logger.info('Loading saved model from {0}'.format(model_path))
+ state = torch.load(model_path)
+ model_params = state['state_dict']
+ params = [(key, value.size(), value.numel()) for key, value in model_params.items()]
+ for key, size, count in params:
+ logger.info('Key:{0}, Size:{1}, Count:{2}'.format(key, size, count))
+ 
+ model.load_state_dict(model_params)
+ model.eval()
+ logger.info(model)
+ 
+ return model
+
+def add_eval_parameters(parser):
+
+ # parser.add_argument('--eval-mode', type = str, default = 'all', help = 'Method to evaluate the attack performance. Default: all. (all or single)')
+ parser.add_argument('--cuda-name', type = str, required = True, help = 'Start a main thread on each cuda.')
+ parser.add_argument('--direct', action='store_true', help = 'Directly add edge or not.')
+ parser.add_argument('--seperate', action='store_true', help = 'Evaluate seperatly or not')
+ parser.add_argument('--mode', type = str, default = '', help = ' '' or '' ')
+ parser.add_argument('--mask-ratio', type=str, default='', help='Mask ratio for Fig4b')
+ return parser
+
+def add_attack_parameters(parser):
+
+ # parser.add_argument('--target-split', type=str, default='0_100_1', help='Ranks to use for target set. Values are 0 for ranks==1; 1 for ranks <=10; 2 for ranks>10 and ranks<=100. Default: 1')
+ parser.add_argument('--target-split', type=str, default='min', help='Methods for target triple selection. Default: min. (min or top_?, top means top_0.1)')
+ parser.add_argument('--target-size', type=int, default=50, help='Number of target triples. Default: 50')
+ parser.add_argument('--target-existed', action='store_true', help='Whether the targeted s_?_o already exists.')
+
+ # parser.add_argument('--budget', type=int, default=1, help='Budget for each target triple for each corruption side')
+
+ parser.add_argument('--attack-goal', type = str, default='single', help='Attack goal. Default: single. (single or global)')
+ parser.add_argument('--neighbor-num', type = int, default=20, help='Max neighbor num for each side. Default: 20')
+ parser.add_argument('--candidate-mode', type = str, default='quadratic', help = 'The method to generate candidate edge. Default: quadratic. (quadratic or linear)')
+ parser.add_argument('--reasonable-rate', type = float, default=0.7, help = 'The added edge\'s existance rank prob greater than this rate')
+ parser.add_argument('--added-edge-num', type = str, default='', help = 'How many edges to add for each target edge. Default: '' means 1.')
+ # parser.add_argument('--neighbor-num', type = int, default=200, help='Max neighbor num for each side. Default: 200')
+ # parser.add_argument('--candidate-mode', type = str, default='linear', help = 'The method to generate candidate edge. Default: quadratic. (quadratic or linear)')
+ parser.add_argument('--attack-batch-size', type=int, default=256, help='Batch size for processing neighbours of target')
+ parser.add_argument('--template-mode', type=str, default = 'manual', help = 'Template mode for transforming edge to single sentense. Default: manual. (manual or auto)')
+
+ parser.add_argument('--update-lissa', action='store_true', help = 'Update lissa cache or not.')
+
+ parser.add_argument('--GPT-batch-size', type=int, default = 64, help = 'Batch size for GPT2 when calculating LM score. Default: 64')
+ parser.add_argument('--LM-softmax', action='store_true', help = 'Use a softmax head on LM prob or not.')
+ parser.add_argument('--LMprob-mode', type=str, default='relative', help = 'Use the absolute LM score or calculate the destruction score when target word is replaced. Default: absolute. (absolute or relative)')
+
+ parser.add_argument('--load-existed', action='store_true', help = 'Use cached intermidiate results or not, when only --reasonable-rate changed, set this param to True')
+
+ return parser
+
+def get_argument_parser():
+ '''Generate an argument parser'''
+ parser = argparse.ArgumentParser(description='Graph embedding')
+
+ parser.add_argument('--seed', type=int, default=1, metavar='S', help='Random seed (default: 1)')
+ 
+ parser.add_argument('--data', type=str, default='GNBR', help='Dataset to use: { GNBR }')
+ parser.add_argument('--model', type=str, default='distmult', help='Choose from: {distmult, conve, complex}')
+ 
+ parser.add_argument('--transe-margin', type=float, default=0.0, help='Margin value for TransE scoring function. Default:0.0')
+ parser.add_argument('--transe-norm', type=int, default=2, help='P-norm value for TransE scoring function. Default:2')
+ 
+ parser.add_argument('--epochs', type=int, default=100, help='Number of epochs to train (default: 100)')
+ parser.add_argument('--lr', type=float, default=0.001, help='Learning rate (default: 0.001)')
+ parser.add_argument('--lr-decay', type=float, default=0.0, help='Weight decay value to use in the optimizer. Default: 0.0')
+ parser.add_argument('--max-norm', action='store_true', help='Option to add unit max norm constraint to entity embeddings')
+ 
+ parser.add_argument('--train-batch-size', type=int, default=64, help='Batch size for train split (default: 128)')
+ parser.add_argument('--test-batch-size', type=int, default=128, help='Batch size for test split (default: 128)')
+ parser.add_argument('--valid-batch-size', type=int, default=128, help='Batch size for valid split (default: 128)')
+ parser.add_argument('--KG-valid-rate', type = float, default=0.1, help='Validation rate during KG embedding training. (default: 0.1)')
+ 
+ parser.add_argument('--save-influence-map', action='store_true', help='Save the influence map during training for gradient rollback.')
+ parser.add_argument('--add-reciprocals', action='store_true')
+
+ parser.add_argument('--embedding-dim', type=int, default=128, help='The embedding dimension (1D). Default: 128')
+ parser.add_argument('--stack-width', type=int, default=16, help='The first dimension of the reshaped/stacked 2D embedding. Second dimension is inferred. Default: 20')
+ #parser.add_argument('--stack_height', type=int, default=10, help='The second dimension of the reshaped/stacked 2D embedding. Default: 10')
+ parser.add_argument('--hidden-drop', type=float, default=0.3, help='Dropout for the hidden layer. Default: 0.3.')
+ parser.add_argument('--input-drop', type=float, default=0.2, help='Dropout for the input embeddings. Default: 0.2.')
+ parser.add_argument('--feat-drop', type=float, default=0.3, help='Dropout for the convolutional features. Default: 0.2.')
+ parser.add_argument('-num-filters', default=32, type=int, help='Number of filters for convolution')
+ parser.add_argument('-kernel-size', default=3, type=int, help='Kernel Size for convolution')
+ 
+ parser.add_argument('--use-bias', action='store_true', help='Use a bias in the convolutional layer. Default: True')
+ 
+ parser.add_argument('--reg-weight', type=float, default=5e-2, help='Weight for regularization. Default: 5e-2')
+ parser.add_argument('--reg-norm', type=int, default=3, help='Norm for regularization. Default: 2')
+ # parser.add_argument('--resume', action='store_true', help='Restore a saved model.')
+ # parser.add_argument('--resume-split', type=str, default='test', help='Split to evaluate a restored model')
+ # parser.add_argument('--reproduce-results', action='store_true', help='Use the hyperparameters to reproduce the results.')
+ # parser.add_argument('--original-data', type=str, default='FB15k-237', help='Dataset to use; this option is needed to set the hyperparams to reproduce the results for training after attack, default: FB15k-237')
+ return parser
+
+def set_hyperparams(args):
+ if args.model == 'distmult':
+ args.lr = 0.005
+ args.train_batch_size = 1024
+ args.reg_norm = 3
+ elif args.model == 'complex':
+ args.lr = 0.005
+ args.reg_norm = 3
+ args.input_drop = 0.4
+ args.train_batch_size = 1024
+ elif args.model == 'conve':
+ args.lr = 0.005
+ args.train_batch_size = 1024
+ args.reg_weight = 0.0
+ 
+ # args.damping = 0.01 
+ # args.lissa_repeat = 1 
+ # args.lissa_depth = 1
+ # args.scale = 500
+ # args.lissa_batch_size = 100
+
+ args.damping = 0.01 
+ args.lissa_repeat = 1 
+ args.lissa_depth = 1
+ args.scale = 400
+ args.lissa_batch_size = 300
+ return args