midification

Parameters.py

@@ -1,9 +1,9 @@
 from audioop import reverse
 
-GNBRfile = '../GNBRdata/'
-PubTatorfile = '../pubtator/'
-UMLSfile = '../umls/META/'
-Pubmedfile = '../pubmed/'
+GNBRfile = 'GNBRdata/'
+PubTatorfile = 'pubtator/'
+UMLSfile = 'umls/META/'
+Pubmedfile = 'pubmed/'
 
 edge_type_dict = {
     'chemical-gene':(['A+', 'A-', 'B', 'E+', 'E-', 'E', 'N', 'O', 'K', 'Z'],

model.py

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

server/server.py → server.py

@@ -9,7 +9,7 @@ import numpy as np
 import json 
 import networkx as nx
 import spacy
-os.system("python -m spacy download en-core-web-sm")
+# os.system("python -m spacy download en-core-web-sm")
 import pickle as pkl
 #%%
 
@@ -17,14 +17,12 @@ from torch.nn.modules.loss import CrossEntropyLoss
 from transformers import AutoTokenizer
 from transformers import BioGptForCausalLM, BartForConditionalGeneration
 
-import server_utils
+from server import server_utils
 
-sys.path.append("..")
 import Parameters
 from Openai.chat import generate_abstract
-sys.path.append("../DiseaseSpecific")
-import utils, attack
-from attack import calculate_edge_bound, get_model_loss_without_softmax
+from DiseaseSpecific import utils, attack
+from DiseaseSpecific.attack import calculate_edge_bound, get_model_loss_without_softmax
 
 
 specific_model =  None
@@ -51,8 +49,8 @@ np.set_printoptions(precision=5)
 cudnn.benchmark = False
 
 model_name = '{0}_{1}_{2}_{3}_{4}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop)
-model_path = '../DiseaseSpecific/saved_models/{0}_{1}.model'.format(args.data, model_name)
-data_path = os.path.join('../DiseaseSpecific/processed_data', args.data)
+model_path = 'DiseaseSpecific/saved_models/{0}_{1}.model'.format(args.data, model_name)
+data_path = os.path.join('DiseaseSpecific/processed_data', args.data)
 data  = utils.load_data(os.path.join(data_path, 'all.txt'))
 
 n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
@@ -596,11 +594,11 @@ def specific_func(start_entity, end_entity):
     o_name = entity_raw_name[id_to_entity[str(o)]]
     attack_data = np.array([[s, r, o]])
     path_list = []
-    with open(f'../DiseaseSpecific/generate_abstract/path/random_{args.reasonable_rate}_path.json', 'r') as fl:
+    with open(f'DiseaseSpecific/generate_abstract/path/random_{args.reasonable_rate}_path.json', 'r') as fl:
         for line in fl.readlines():
             line.replace('\n', '')
             path_list.append(line)
-    with open(f'../DiseaseSpecific/generate_abstract/random_{args.reasonable_rate}_sentence.json', 'r') as fl:
+    with open(f'DiseaseSpecific/generate_abstract/random_{args.reasonable_rate}_sentence.json', 'r') as fl:
         sentence_dict = json.load(fl)
     dpath = []
     for k, v in sentence_dict.items():

utils.py

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