Spaces:
Runtime error
Runtime error
''' | |
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 | |