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