Agnostic

DiseaseAgnostic/KG_extractor.py

@@ -0,0 +1,473 @@
+#%%
+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.add_argument('--reasonable-rate', type = float, default=0.7, help = 'The added edge\'s existance rank prob greater than this rate')
+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('--init-mode', type = str, default='random', help = 'How to select target nodes')
+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 = '../DiseaseSpecific/processed_data/GNBR'
+target_path = f'processed_data/target_{args.reasonable_rate}{args.init_mode}.pkl'
+attack_path = f'processed_data/attack_edge_{args.model}_{args.reasonable_rate}{args.init_mode}.pkl'
+modified_attack_path = f'processed_data/attack_edge_{args.model}_{args.reasonable_rate}{args.init_mode}{args.mode}.pkl'
+
+with open(attack_path, 'rb') as fl:
+    Attack_edge_list = pkl.load(fl)
+attack_data = np.array(Attack_edge_list).reshape(-1, 3)
+#%%
+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('../DiseaseSpecific/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.init_mode}{args.reasonable_rate}_chat.json', 'r') as fl:
+        draft = json.load(fl)
+elif args.mode == 'finetune':
+    with open(f'generate_abstract/{args.init_mode}{args.reasonable_rate}_sentence_finetune.json', 'r') as fl:
+        draft = json.load(fl)
+elif args.mode == 'bioBART':
+    with open(f'generate_abstract/{args.init_mode}{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.init_mode}{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.init_mode}{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.init_mode}{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.init_mode}{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]
+        s = str(s)
+        r = str(r)
+        o = str(o)
+        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)
+    final_Attack = Attack
+    print('Len of Attack:', len(Attack))
+    with open(modified_attack_path, 'wb') as fl:
+        pkl.dump(final_Attack, fl)
+else:
+    raise Exception('Wrong action !!')

DiseaseAgnostic/edge_to_abstract.py

@@ -0,0 +1,652 @@
+#%%
+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.add_argument('--reasonable-rate', type = float, default=0.7, help = 'The added edge\'s existance rank prob greater than this rate')
+parser.add_argument('--mode', type=str, default='sentence', help='sentence, biogpt or finetune')
+parser.add_argument('--init-mode', type = str, default='random', help = 'How to select target nodes')
+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 = '../DiseaseSpecific/processed_data/GNBR'
+target_path = f'processed_data/target_{args.reasonable_rate}{args.init_mode}.pkl'
+attack_path = f'processed_data/attack_edge_{args.model}_{args.reasonable_rate}{args.init_mode}.pkl'
+
+# target_data = utils.load_data(target_path)
+with open(target_path, 'rb') as fl:
+    Target_node_list = pkl.load(fl)
+with open(attack_path, 'rb') as fl:
+    Attack_edge_list = pkl.load(fl)
+attack_data = np.array(Attack_edge_list).reshape(-1, 3)
+# assert target_data.shape == attack_data.shape
+#%%
+
+with open('../DiseaseSpecific/processed_data/GNBR/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 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 = 24
+    single_sentence = {}
+    test_text = []
+    test_dp = []
+    test_parse = []
+    for i, (s, r, o) in enumerate(tqdm(attack_data)):
+        
+        s = str(s)
+        r = str(r)
+        o = str(o)
+        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.init_mode}{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.init_mode}{args.reasonable_rate}_path.json', 'w') as fl:
+        fl.write('\n'.join(test_dp))
+    with open (f'generate_abstract/path/{args.init_mode}{args.reasonable_rate}_temp.json', 'w') as fl:
+        fl.write('\n'.join(test_text))
+
+elif args.mode == 'biogpt':
+    pass
+    # from biogpt_generate import GPT_eval
+    # import spacy
+
+    # model = GPT_eval(args.seed)
+
+    # nlp = spacy.load("en_core_web_sm")
+    # with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_sentence.json', 'r') as fl:
+    #     data = json.load(fl)
+
+    # KK = []
+    # input = []
+    # for i,(k, v) in enumerate(data.items()):
+    #     KK.append(k)
+    #     input.append(v)
+    # output = model.eval(input)
+
+    # ret = {}
+    # for i, o in enumerate(output):
+
+    #     o = o.replace('<|abstract|>', '')
+    #     doc = nlp(o)
+    #     sen_list = []
+    #     sen_set = set()
+    #     for sen in doc.sents:
+    #         txt = sen.text
+    #         if not (txt.lower() in sen_set):
+    #             sen_set.add(txt.lower())
+    #             sen_list.append(txt)
+    #     O = ' '.join(sen_list)
+    #     ret[KK[i]] = {'in' : input[i], 'out' : O}
+    
+    # with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_biogpt.json', 'w') as fl:
+    #     json.dump(ret, fl, indent=4)
+
+elif args.mode == 'finetune':
+
+    import spacy
+    import pprint
+    from transformers import AutoModel, AutoTokenizer,BartForConditionalGeneration
+
+    print('Finetuning ...')
+
+    with open(f'generate_abstract/{args.init_mode}{args.reasonable_rate}_chat.json', 'r') as fl:
+        draft = json.load(fl)
+    with open (f'generate_abstract/path/{args.init_mode}{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.init_mode}{args.reasonable_rate}{args.ratio}_candidates.json'):
+        with open(f'generate_abstract/bioBART/{args.init_mode}{args.reasonable_rate}{args.ratio}_candidates.json', 'r') as fl:
+            ret_candidates = json.load(fl)
+    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, position):
+
+        #     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(' ')
+        #     l_p = 0
+        #     r_p = 1
+        #     assert position == 'front' or position == 'back'
+        #     if position == 'back':
+        #         l_p, r_p = r_p, l_p
+        #     P = np.linspace(start = l_p, stop = r_p, num = len(token_list))
+        #     P = (P ** 3) * 0.4
+
+        #     ret_list = []
+        #     for t, p in zip(token_list, list(P)):
+        #         if '.' in t or '(' in t or ')' in t or '[' in t or ']' in t:
+        #             ret_list.append(t)
+        #         else:
+        #             if np.random.rand() < p:
+        #                 ret_list.append('<mask>')
+        #             else:
+        #                 ret_list.append(t)
+        #     return [' '.join(ret_list)]
+        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]
+            s = str(s)
+            o = str(o)
+            r = str(r)
+
+            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, 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.init_mode}{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 = {}
+    case_study = {}
+    p_ret = {}
+    add = 0
+    dpath_i = 0
+    inner_better = 0
+    outter_better = 0
+    better_than_gpt = 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]
+        s = str(s)
+        r = str(r)
+        o = str(o)
+
+        if int(s) == -1:
+            ret[k] = {'prompt': '', 'in':'', 'out': ''}
+            p_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]
+        # sen_list = sen_list[:old_L]
+
+        # mini_span should be preserved
+        # for j in range(len(log_Loss)):
+        #     doc = nlp(sen_list[j])
+        #     sens = [sen.text for sen in doc.sents]
+        #     Len = len(sen_list)
+        #     check_text = ' '.join(sens[j : max(0,len(sens) - Len) + j + 1])
+        #     if span not in check_text:
+        #         log_Loss[j] += 1
+
+        p = np.argmin(log_Loss)
+        if p < old_L // 2:
+            inner_better += 1
+        else:
+            outter_better += 1
+        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[p+1+old_L]:
+            add += 1
+
+        if real_log_Loss[p] < real_log_Loss[old_L]:
+            better_than_gpt += 1
+        else:
+            if real_log_Loss[p] > real_log_Loss[p+1+old_L]:
+                p = p+1+old_L
+        # case_study[k] = {'path':path_text, 'entity_0': text_s, 'entity_1': text_o, 'GPT_in': input, 'Prompt': prompt, 'GPT_out': {'text': output, 'perplexity': str(np.exp(real_log_Loss[old_L]))}, 'BART_in': BART_in[p], 'BART_out': {'text': sen_list[p], 'perplexity': str(np.exp(real_log_Loss[p]))}, 'Assist': {'text': Assist[p], 'perplexity': str(np.exp(real_log_Loss[p+1+old_L]))}}
+        ret[k] = {'prompt': prompt, 'in':input, 'out': sen_list[p]}
+        p_ret[k] = {'prompt': prompt, 'in':input, 'out': sen_list[p_p]}
+    print(add)
+    print('inner_better:', inner_better)
+    print('outter_better:', outter_better)
+    print('better_than_gpt:', better_than_gpt)
+    print('better_than_replace', add)
+    with open(f'generate_abstract/{args.init_mode}{args.reasonable_rate}{args.ratio}_bioBART_finetune.json', 'w') as fl:
+        json.dump(ret, fl, indent=4)
+    # with open(f'generate_abstract/bioBART/case_{args.target_split}_{args.reasonable_rate}_bioBART_finetune.json', 'w') as fl:
+    #     json.dump(case_study, fl, indent=4)
+    with open(f'generate_abstract/bioBART/{args.init_mode}{args.reasonable_rate}{args.ratio}_scored.json', 'w') as fl:
+        json.dump(scored, fl, indent=4)
+    with open(f'generate_abstract/bioBART/{args.init_mode}{args.reasonable_rate}{args.ratio}_perplexity.json', 'w') as fl:
+        json.dump(p_ret, fl, indent=4)
+
+    # 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]
+
+    # 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])
+
+    # fine_dict = {}
+    # for k, v_dict in draft.items():
+
+    #     input = v_dict['in']
+    #     output = v_dict['out']
+    #     fine_dict[k] = {'in':input, 'out': input + ' ' + output}
+
+    # with open(f'generate_abstract/{args.target_split}_{args.reasonable_rate}_sentence_finetune.json', 'w') as fl:
+    #     json.dump(fine_dict, fl, indent=4)
+else:
+    raise Exception('Wrong mode !!')

DiseaseAgnostic/evaluation.py

@@ -0,0 +1,219 @@
+#%%
+import logging
+from symbol import parameters
+from textwrap import indent
+import os
+import tempfile
+import sys
+from matplotlib import collections
+import pandas as pd
+import json
+from glob import glob
+from tqdm import tqdm
+import numpy as np
+from pprint import pprint
+import torch
+import pickle as pkl
+from collections import Counter 
+# print(dir(collections))
+import networkx as nx
+from collections import Counter
+import utils
+from torch.nn import functional as F
+sys.path.append("..")
+import Parameters
+from DiseaseSpecific.attack import calculate_edge_bound, get_model_loss_without_softmax
+
+#%%
+def load_data(file_name):
+    df = pd.read_csv(file_name, sep='\t', header=None, names=None, dtype=str)
+    df = df.drop_duplicates()
+    return df.values
+
+parser = utils.get_argument_parser()
+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('--mode', type = str, default='', help = ' "" or chat or bioBART')
+parser.add_argument('--init-mode', type = str, default='random', help = 'How to select target nodes') # 'single' for case study 
+parser.add_argument('--added-edge-num', type = str, default = '', help = 'Added edge num')
+
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+utils.seed_all(args.seed)
+graph_edge_path = '../DiseaseSpecific/processed_data/GNBR/all.txt'
+idtomeshid_path = '../DiseaseSpecific/processed_data/GNBR/entities_reverse_dict.json'
+model_path = f'../DiseaseSpecific/saved_models/GNBR_{args.model}_128_0.2_0.3_0.3.model'
+data_path = '../DiseaseSpecific/processed_data/GNBR'
+target_path = f'processed_data/target_{args.reasonable_rate}{args.init_mode}.pkl'
+attack_path = f'processed_data/attack_edge_{args.model}_{args.reasonable_rate}{args.init_mode}{args.added_edge_num}{args.mode}.pkl'
+
+with open(Parameters.GNBRfile+'original_entity_raw_name', 'rb') as fl:
+    full_entity_raw_name = pkl.load(fl)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# device = torch.device("cpu")
+
+args.device = device
+n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+model = utils.load_model(model_path, args, n_ent, n_rel, args.device)
+
+graph_edge = utils.load_data(graph_edge_path)
+with open(idtomeshid_path, 'r') as fl:
+    idtomeshid = json.load(fl)
+print(graph_edge.shape, len(idtomeshid))
+
+divide_bound, data_mean, data_std = calculate_edge_bound(graph_edge, model, args.device, n_ent)
+print('Defender ...')
+print(divide_bound, data_mean, data_std)
+
+meshids = list(idtomeshid.values())
+cal = {
+    'chemical' : 0,
+    'disease' : 0,
+    'gene' : 0
+}
+for meshid in meshids:
+    cal[meshid.split('_')[0]] += 1
+# pprint(cal)
+
+def check_reasonable(s, r, o):
+
+    train_trip = np.asarray([[s, r, o]])
+    train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+    edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+    # edge_losse_log_prob = torch.log(F.softmax(-edge_loss, dim = -1))
+
+    edge_loss = edge_loss.item() 
+    edge_loss = (edge_loss - data_mean) / data_std
+    edge_losses_prob =  1 / ( 1 + np.exp(edge_loss - divide_bound) )
+    bound = 1 - args.reasonable_rate
+
+    return (edge_losses_prob > bound),  edge_losses_prob
+
+edgeid_to_edgetype = {}
+edgeid_to_reversemask = {}
+for k, id_list in Parameters.edge_type_to_id.items():
+    for iid, mask in zip(id_list, Parameters.reverse_mask[k]):
+        edgeid_to_edgetype[str(iid)] = k
+        edgeid_to_reversemask[str(iid)] = mask
+
+with open(target_path, 'rb') as fl:
+    Target_node_list = pkl.load(fl)
+with open(attack_path, 'rb') as fl:
+    Attack_edge_list = pkl.load(fl)
+with open(Parameters.UMLSfile+'drug_term', 'rb') as fl:
+    drug_term = pkl.load(fl)
+with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+    entity_raw_name = pkl.load(fl)
+drug_meshid = []
+for meshid, nm in entity_raw_name.items():
+    if nm.lower() in drug_term and meshid.split('_')[0] == 'chemical':
+        drug_meshid.append(meshid)
+drug_meshid = set(drug_meshid)
+
+if args.init_mode == 'single':
+    name_list = []
+    for target in Target_node_list:
+        name = entity_raw_name[idtomeshid[str(target)]]
+        name_list.append(name)
+    with open(f'results/name_list_{args.reasonable_rate}{args.init_mode}.txt', 'w') as fl:
+        fl.write('\n'.join(name_list))
+    # print(Target_node_list)
+    # # print(Attack_edge_list)
+    # addset = set()
+    # if args.added_edge_num == 1:
+    #     for edge in Attack_edge_list:
+    #         addset.add(edge[2])
+    # else:
+    #     for edge_list in Attack_edge_list:
+    #         for edge in edge_list:
+    #             addset.add(edge[2])
+    # print(addset)
+    # print(len(addset))
+    # typeset = set()
+    # for iid in addset:
+    #     typeset.add(idtomeshid[str(iid)].split('_')[0])
+    # print(typeset)
+    # raise Exception('done')
+
+if args.init_mode == 'single':
+    Target_node_list = [[Target_node_list[i]] for i in range(len(Target_node_list))]
+    Attack_edge_list = [[Attack_edge_list[i]] for i in range(len(Attack_edge_list))]
+else:
+    print(len(Attack_edge_list), len(Target_node_list))
+    tmp_target_node_list = []
+    tmp_attack_edge_list = []
+    for l in range(0,len(Target_node_list), 50):
+        r = min(l+50, len(Target_node_list))
+        tmp_target_node_list.append(Target_node_list[l:r])
+        tmp_attack_edge_list.append(Attack_edge_list[l:r])
+    Target_node_list = tmp_target_node_list
+    Attack_edge_list = tmp_attack_edge_list
+
+# for i, init_p in enumerate([0.1, 0.3, 0.5, 0.7, 0.9]):
+
+#     target_node_list = Target_node_list[i]
+#     attack_edge_list = Attack_edge_list[i]
+Init = []
+After = []
+# final_init = []
+# final_after = []
+for i, (target_node_list, attack_edge_list) in enumerate(zip(Target_node_list, Attack_edge_list)):
+
+    G = nx.DiGraph()
+    for s, r, o in graph_edge:
+        assert idtomeshid[s].split('_')[0] == edgeid_to_edgetype[r].split('-')[0]
+        if edgeid_to_reversemask[r] == 1:
+            G.add_edge(int(o), int(s))
+        else:
+            G.add_edge(int(s), int(o))
+
+    pagerank_value_1 = nx.pagerank(G, max_iter = 200, tol=1.0e-7) 
+    
+    for target, attack_list in tqdm(list(zip(target_node_list, attack_edge_list))):
+        pr = list(pagerank_value_1.items())
+        pr.sort(key = lambda x: x[1])
+        list_iid = []
+        for iid, score in pr:
+            tp = idtomeshid[str(iid)].split('_')[0]
+            if tp == 'chemical':
+                # if idtomeshid[str(iid)] in drug_meshid:
+                list_iid.append(iid)
+        init_rank = len(list_iid) - list_iid.index(target)
+        # init_rank = 1 - list_iid.index(target) / len(list_iid)
+        Init.append(init_rank)
+
+    for target, attack_list in tqdm(list(zip(target_node_list, attack_edge_list))):
+        
+        if args.mode == '' and (args.added_edge_num == '' or int(args.added_edge_num) == 1):
+            if int(attack_list[0]) == -1:
+                attack_list = []
+            else:
+                attack_list = [attack_list]
+        if len(attack_list) > 0:
+            for s, r, o in attack_list:
+                bo, prob = check_reasonable(s, r, o)
+                if bo:
+                    if edgeid_to_reversemask[str(r)] == 1:
+                        G.add_edge(int(o), int(s))
+                    else:
+                        G.add_edge(int(s), int(o))
+    pagerank_value_1 = nx.pagerank(G, max_iter = 200, tol=1.0e-7) 
+    for target, attack_list in tqdm(list(zip(target_node_list, attack_edge_list))):
+        pr = list(pagerank_value_1.items())
+        pr.sort(key = lambda x: x[1])
+        list_iid = []
+        for iid, score in pr:
+            tp = idtomeshid[str(iid)].split('_')[0]
+            if tp == 'chemical':
+                # if idtomeshid[str(iid)] in drug_meshid:
+                list_iid.append(iid)
+        after_rank = len(list_iid) - list_iid.index(target)
+        # after_rank = 1 - list_iid.index(target) / len(list_iid)
+        After.append(after_rank)
+    with open(f'results/Init_{args.reasonable_rate}{args.init_mode}.pkl', 'wb') as fl:
+        pkl.dump(Init, fl)
+    with open(f'results/After_{args.model}_{args.reasonable_rate}{args.init_mode}{args.added_edge_num}{args.mode}.pkl', 'wb') as fl:
+        pkl.dump(After, fl)
+    print(np.mean(Init), np.std(Init))
+    print(np.mean(After), np.std(After))

DiseaseAgnostic/generate_target_and_attack.py

@@ -0,0 +1,371 @@
+#%%
+import logging
+from symbol import parameters
+from textwrap import indent
+import os
+import tempfile
+import sys
+from matplotlib import collections
+import pandas as pd
+import json
+from glob import glob
+from tqdm import tqdm
+import numpy as np
+from pprint import pprint
+import torch
+import pickle as pkl
+from collections import Counter 
+# print(dir(collections))
+import networkx as nx
+from collections import Counter
+import utils
+from torch.nn import functional as F
+sys.path.append("..")
+import Parameters
+from DiseaseSpecific.attack import calculate_edge_bound, get_model_loss_without_softmax
+
+#%%
+def load_data(file_name):
+    df = pd.read_csv(file_name, sep='\t', header=None, names=None, dtype=str)
+    df = df.drop_duplicates()
+    return df.values
+
+parser = utils.get_argument_parser()
+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('--init-mode', type = str, default='single', help = 'How to select target nodes') # 'single' for case study 
+parser.add_argument('--added-edge-num', type = str, default = '', help = 'Added edge num')
+
+args = parser.parse_args()
+args = utils.set_hyperparams(args)
+utils.seed_all(args.seed)
+graph_edge_path = '../DiseaseSpecific/processed_data/GNBR/all.txt'
+idtomeshid_path = '../DiseaseSpecific/processed_data/GNBR/entities_reverse_dict.json'
+model_path = f'../DiseaseSpecific/saved_models/GNBR_{args.model}_128_0.2_0.3_0.3.model'
+data_path = '../DiseaseSpecific/processed_data/GNBR'
+with open(Parameters.GNBRfile+'original_entity_raw_name', 'rb') as fl:
+    full_entity_raw_name = pkl.load(fl)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+args.device = device
+n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
+model = utils.load_model(model_path, args, n_ent, n_rel, args.device)
+print(device)
+
+graph_edge = utils.load_data(graph_edge_path)
+with open(idtomeshid_path, 'r') as fl:
+    idtomeshid = json.load(fl)
+print(graph_edge.shape, len(idtomeshid))
+
+divide_bound, data_mean, data_std = calculate_edge_bound(graph_edge, model, args.device, n_ent)
+print('Defender ...')
+print(divide_bound, data_mean, data_std)
+
+meshids = list(idtomeshid.values())
+cal = {
+    'chemical' : 0,
+    'disease' : 0,
+    'gene' : 0
+}
+for meshid in meshids:
+    cal[meshid.split('_')[0]] += 1
+# pprint(cal)
+
+def check_reasonable(s, r, o):
+
+    train_trip = np.asarray([[s, r, o]])
+    train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+    edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+    # edge_losse_log_prob = torch.log(F.softmax(-edge_loss, dim = -1))
+
+    edge_loss = edge_loss.item() 
+    edge_loss = (edge_loss - data_mean) / data_std
+    edge_losses_prob =  1 / ( 1 + np.exp(edge_loss - divide_bound) )
+    bound = 1 - args.reasonable_rate
+
+    return (edge_losses_prob > bound),  edge_losses_prob
+
+edgeid_to_edgetype = {}
+edgeid_to_reversemask = {}
+for k, id_list in Parameters.edge_type_to_id.items():
+    for iid, mask in zip(id_list, Parameters.reverse_mask[k]):
+        edgeid_to_edgetype[str(iid)] = k
+        edgeid_to_reversemask[str(iid)] = mask
+reverse_tot = 0
+G = nx.DiGraph()
+for s, r, o in graph_edge:
+    assert idtomeshid[s].split('_')[0] == edgeid_to_edgetype[r].split('-')[0]
+    if edgeid_to_reversemask[r] == 1:
+        reverse_tot += 1
+        G.add_edge(int(o), int(s))
+    else:
+        G.add_edge(int(s), int(o))
+# print(reverse_tot)
+print('Edge num:', G.number_of_edges(), 'Node num:', G.number_of_nodes())
+pagerank_value_1 = nx.pagerank(G, max_iter = 200, tol=1.0e-7) 
+
+#%%
+with open(Parameters.UMLSfile+'drug_term', 'rb') as fl:
+    drug_term = pkl.load(fl)
+with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
+    entity_raw_name = pkl.load(fl)
+drug_meshid = []
+for meshid, nm in entity_raw_name.items():
+    if nm.lower() in drug_term and meshid.split('_')[0] == 'chemical':
+        drug_meshid.append(meshid)
+drug_meshid = set(drug_meshid)
+pr = list(pagerank_value_1.items())
+pr.sort(key = lambda x: x[1])
+sorted_rank = { 'chemical' : [],
+                'gene' : [],
+                'disease': [],
+                'merged' : []}
+for iid, score in pr:
+    tp = idtomeshid[str(iid)].split('_')[0]
+    if tp == 'chemical':
+        if idtomeshid[str(iid)] in drug_meshid:
+            sorted_rank[tp].append((iid, score))
+    else:
+        sorted_rank[tp].append((iid, score))
+        sorted_rank['merged'].append((iid, score))
+llen = len(sorted_rank['merged']) 
+sorted_rank['merged'] = sorted_rank['merged'][llen * 3 // 4 : ]
+print(len(sorted_rank['chemical']))
+print(len(sorted_rank['gene']), len(sorted_rank['disease']), len(sorted_rank['merged']))
+
+#%%
+Target_node_list = []
+Attack_edge_list = []
+if args.init_mode == '':
+
+    if args.added_edge_num != '' and args.added_edge_num != '1':
+        raise Exception('added_edge_num must be 1 when init_mode=='' ')
+    for init_p in [0.1, 0.3, 0.5, 0.7, 0.9]:
+
+        p  = len(sorted_rank['chemical']) * init_p
+        print('Init p:', init_p)
+        target_node_list = []
+        attack_edge_list = []
+        num_max_eq = 0
+        mean_rank_of_total_max = 0
+        for pp in tqdm(range(int(p)-10, int(p)+10)):
+            target = sorted_rank['chemical'][pp][0]
+            target_node_list.append(target)
+
+            candidate_list = []
+            score_list = []
+            loss_list = []
+            for iid, score in sorted_rank['merged']:
+                a = G.number_of_edges(iid, target) + 1
+                if a != 1:
+                    continue
+                b = G.out_degree(iid) + 1
+                tp = idtomeshid[str(iid)].split('_')[0]
+                edge_losses = []
+                r_list = []
+                for r in range(len(edgeid_to_edgetype)):
+                    r_tp = edgeid_to_edgetype[str(r)]
+                    if (edgeid_to_reversemask[str(r)] == 0 and r_tp.split('-')[0] == tp and r_tp.split('-')[1] == 'chemical'):
+                        train_trip = np.array([[iid, r, target]])
+                        train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+                        edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+                        edge_losses.append(edge_loss.unsqueeze(0).detach())
+                        r_list.append(r)
+                    elif(edgeid_to_reversemask[str(r)] == 1 and r_tp.split('-')[0] == 'chemical' and r_tp.split('-')[1] == tp):
+                        train_trip = np.array([[iid, r, target]]) # add batch dim
+                        train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+                        edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+                        edge_losses.append(edge_loss.unsqueeze(0).detach())
+                        r_list.append(r)
+                if len(edge_losses)==0:
+                    continue
+                min_index = torch.argmin(torch.cat(edge_losses, dim = 0))
+                r = r_list[min_index]
+                r_tp = edgeid_to_edgetype[str(r)]
+                
+                if (edgeid_to_reversemask[str(r)] == 0):
+                    bo, prob = check_reasonable(iid, r, target)
+                    if bo:
+                        candidate_list.append((iid, r, target))
+                        score_list.append(score * a / b)
+                        loss_list.append(edge_losses[min_index].item())
+                if (edgeid_to_reversemask[str(r)] == 1):
+                    bo, prob = check_reasonable(target, r, iid)
+                    if bo:
+                        candidate_list.append((target, r, iid))
+                        score_list.append(score * a / b)
+                        loss_list.append(edge_losses[min_index].item())
+            
+            if len(candidate_list) == 0:
+                attack_edge_list.append((-1, -1, -1))
+                continue
+            norm_score = np.array(score_list) / np.sum(score_list)
+            norm_loss = np.exp(-np.array(loss_list)) / np.sum(np.exp(-np.array(loss_list)))
+
+            total_score = norm_score * norm_loss
+            max_index = np.argmax(total_score)
+            attack_edge_list.append(candidate_list[max_index])
+
+            score_max_index = np.argmax(norm_score)
+            if score_max_index == max_index:
+                num_max_eq += 1
+
+            score_index_list = list(zip(list(range(len(norm_score))), norm_score))
+            score_index_list.sort(key = lambda x: x[1], reverse = True)
+            max_index_in_score = score_index_list.index((max_index, norm_score[max_index]))
+            mean_rank_of_total_max += max_index_in_score / len(norm_score)
+        print('num_max_eq:', num_max_eq)
+        print('mean_rank_of_total_max:', mean_rank_of_total_max / 20)
+        Target_node_list.append(target_node_list)
+        Attack_edge_list.append(attack_edge_list)
+else:
+    assert args.init_mode == 'random' or args.init_mode == 'single'
+    print(f'Init mode : {args.init_mode}')
+    utils.seed_all(args.seed)
+
+    if args.init_mode == 'random':
+        index = np.random.choice(len(sorted_rank['chemical']), 400, replace = False)
+    else:
+        # index = [5807, 6314, 5799, 5831, 3954, 5654, 5649, 5624, 2412, 2407]
+        
+        index = np.random.choice(len(sorted_rank['chemical']), 400, replace = False)
+        with open(f'../pagerank/results/After_distmult_0.7random10.pkl', 'rb') as fl:
+            edge = pkl.load(fl)
+        with open('../pagerank/results/Init_0.7random.pkl', 'rb') as fl:
+            init = pkl.load(fl)
+        increase = (np.array(init) - np.array(edge)) / np.array(init)
+        increase = increase.reshape(-1)
+        selected_index = np.argsort(increase)[::-1][:10]
+        # print(selected_index)
+        # print(increase[selected_index])
+        # print(np.array(init)[selected_index])
+        # print(np.array(edge)[selected_index])
+        index = [index[i] for i in selected_index]
+        # llen = len(sorted_rank['chemical'])
+        # index = np.random.choice(range(llen//4, llen), 4, replace = False)
+        # index = selected_index + list(index)
+        # for i in index:
+        #     ii = str(sorted_rank['chemical'][i][0])
+        #     nm = entity_raw_name[idtomeshid[ii]]
+        #     nmset = full_entity_raw_name[idtomeshid[ii]]
+        #     print('**'*10)
+        #     print(i)
+        #     print(nm)
+        #     print(nmset)
+        # raise Exception('stop')
+    target_node_list = []
+    attack_edge_list = []
+    num_max_eq = 0
+    mean_rank_of_total_max = 0
+
+    for pp in tqdm(index):
+        target = sorted_rank['chemical'][pp][0]
+        target_node_list.append(target)
+
+        print('Target:', entity_raw_name[idtomeshid[str(target)]])
+
+        candidate_list = []
+        score_list = []
+        loss_list = []
+        main_dict = {}
+        for iid, score in sorted_rank['merged']:
+            a = G.number_of_edges(iid, target) + 1
+            if a != 1:
+                continue
+            b = G.out_degree(iid) + 1
+            tp = idtomeshid[str(iid)].split('_')[0]
+            edge_losses = []
+            r_list = []
+            for r in range(len(edgeid_to_edgetype)):
+                r_tp = edgeid_to_edgetype[str(r)]
+                if (edgeid_to_reversemask[str(r)] == 0 and r_tp.split('-')[0] == tp and r_tp.split('-')[1] == 'chemical'):
+                    train_trip = np.array([[iid, r, target]])
+                    train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+                    edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+                    edge_losses.append(edge_loss.unsqueeze(0).detach())
+                    r_list.append(r)
+                elif(edgeid_to_reversemask[str(r)] == 1 and r_tp.split('-')[0] == 'chemical' and r_tp.split('-')[1] == tp):
+                    train_trip = np.array([[iid, r, target]]) # add batch dim
+                    train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
+                    edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
+                    edge_losses.append(edge_loss.unsqueeze(0).detach())
+                    r_list.append(r)
+            if len(edge_losses)==0:
+                continue
+            min_index = torch.argmin(torch.cat(edge_losses, dim = 0))
+            r = r_list[min_index]
+            r_tp = edgeid_to_edgetype[str(r)]
+            
+
+            old_len = len(candidate_list)
+            if (edgeid_to_reversemask[str(r)] == 0):
+                bo, prob = check_reasonable(iid, r, target)
+                if bo:
+                    candidate_list.append((iid, r, target))
+                    score_list.append(score * a / b)
+                    loss_list.append(edge_losses[min_index].item())
+            if (edgeid_to_reversemask[str(r)] == 1):
+                bo, prob = check_reasonable(target, r, iid)
+                if bo:
+                    candidate_list.append((target, r, iid))
+                    score_list.append(score * a / b)
+                    loss_list.append(edge_losses[min_index].item())
+
+            if len(candidate_list) != old_len:
+                if int(iid) in main_iid:
+                    main_dict[iid] = len(candidate_list) - 1
+        
+        if len(candidate_list) == 0:
+            if args.added_edge_num == '' or int(args.added_edge_num) == 1:
+                attack_edge_list.append((-1,-1,-1))
+            else:
+                attack_edge_list.append([])
+            continue
+        norm_score = np.array(score_list) / np.sum(score_list)
+        norm_loss = np.exp(-np.array(loss_list)) / np.sum(np.exp(-np.array(loss_list)))
+
+        total_score = norm_score * norm_loss
+        total_score_index = list(zip(range(len(total_score)), total_score))
+        total_score_index.sort(key = lambda x: x[1], reverse = True)
+
+        norm_score_index = np.argsort(norm_score)[::-1]
+        norm_loss_index = np.argsort(norm_loss)[::-1]
+        total_index = np.argsort(total_score)[::-1]
+        assert total_index[0] == total_score_index[0][0]
+        # find rank of main index 
+        for k, v in main_dict.items():
+            k = int(k)
+            index = v
+            print(f'score rank of {entity_raw_name[idtomeshid[str(k)]]}: ', norm_score_index.tolist().index(index))
+            print(f'loss rank of {entity_raw_name[idtomeshid[str(k)]]}: ', norm_loss_index.tolist().index(index))
+            print(f'total rank of {entity_raw_name[idtomeshid[str(k)]]}: ', total_index.tolist().index(index))
+        
+        max_index = np.argmax(total_score)
+        assert max_index == total_score_index[0][0]
+
+        tmp_add = []
+        add_num = 1
+        if args.added_edge_num == '' or int(args.added_edge_num) == 1:
+            attack_edge_list.append(candidate_list[max_index])
+        else:
+            add_num = int(args.added_edge_num)
+            for i in range(add_num):
+                tmp_add.append(candidate_list[total_score_index[i][0]])
+            attack_edge_list.append(tmp_add)
+
+        score_max_index = np.argmax(norm_score)
+        if score_max_index == max_index:
+            num_max_eq += 1
+        score_index_list = list(zip(list(range(len(norm_score))), norm_score))
+        score_index_list.sort(key = lambda x: x[1], reverse = True)
+        max_index_in_score = score_index_list.index((max_index, norm_score[max_index]))
+        mean_rank_of_total_max += max_index_in_score / len(norm_score)
+    print('num_max_eq:', num_max_eq)
+    print('mean_rank_of_total_max:', mean_rank_of_total_max / 400)
+    Target_node_list = target_node_list
+    Attack_edge_list = attack_edge_list
+print(np.array(Target_node_list).shape)
+print(np.array(Attack_edge_list).shape)
+# with open(f'processed_data/target_{args.reasonable_rate}{args.init_mode}.pkl', 'wb') as fl:
+#     pkl.dump(Target_node_list, fl)
+# with open(f'processed_data/attack_edge_{args.model}_{args.reasonable_rate}{args.init_mode}{args.added_edge_num}.pkl', 'wb') as fl:
+#     pkl.dump(Attack_edge_list, fl)

DiseaseAgnostic/model.py

@@ -0,0 +1,520 @@
+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)
+
+        #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)
+
+        # complex space bilinear product (equivalent to HolE)
+#         realrealreal = torch.mm(s_real*rel_real, emb_e_real.transpose(1,0))
+#         realimgimg = torch.mm(s_real*rel_img, emb_e_img.transpose(1,0))
+#         imgrealimg = torch.mm(s_img*rel_real, emb_e_img.transpose(1,0))
+#         imgimgreal = torch.mm(s_img*rel_img, emb_e_real.transpose(1,0))
+#         pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+        
+        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)
+            
+            # 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
+        
+        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)
+            
+            # complex space bilinear product (equivalent to HolE)
+#             realrealreal = torch.mm(s_real*rel_real, emb_e_real.transpose(1,0))
+#             realimgimg = torch.mm(s_real*rel_img, emb_e_img.transpose(1,0))
+#             imgrealimg = torch.mm(s_img*rel_real, emb_e_img.transpose(1,0))
+#             imgimgreal = torch.mm(s_img*rel_img, emb_e_real.transpose(1,0))
+#             pred = realrealreal + realimgimg + imgrealimg - imgimgreal
+            
+            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 = torch.diagonal(pred)
+        
+        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 = x*obj_emb
+        
+        if sigmoid:
+            pred = torch.sigmoid(pred)
+
+        return pred

DiseaseAgnostic/utils.py

@@ -0,0 +1,187 @@
+'''
+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):
+    df = pd.read_csv(file_name, sep='\t', header=None, names=None, dtype=str)
+    df = df.drop_duplicates()
+    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')
+    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('--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)')
+
+    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, complex, transe, conve}')
+    
+    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