def __init__(self, datasetName, mode, pos_dataset, whole_dataset, batch_size, entity_set, relation_set, neg_rate, groundtruth=None, possible_entities=None, rdrop=False, pos_neg_dataset=None):

self.datasetName = datasetName

self.batch_size = batch_size

self.entity_set = entity_set

self.relation_set = relation_set

self.mode = mode

self.whole_dataset = whole_dataset

self.pos_neg_dataset = pos_neg_dataset # dataset with originally provided negative samples

self.neg_rate = neg_rate

self.groundtruth = groundtruth

self.possible_entities = possible_entities

self.rdrop = rdrop

if not os.path.exists('./sampler'):

os.makedirs('./sampler')

if mode == 'train':

global count_sampler

count_sampler += 1

dataset_path = 'sampler/{}-{}-{}-{}.pkl'.format(datasetName, mode, neg_rate, count_sampler)

else:

dataset_path = 'sampler/{}-{}-{}.pkl'.format(datasetName, mode, neg_rate)

if self.datasetName in ['fb13'] and mode != 'train':

self.dataset = [ ((i[0], i[1], i[2]), i[3]) for i in self.pos_neg_dataset]

else:

if os.path.exists(dataset_path):

with open(dataset_path, 'rb') as fil:

self.dataset = pickle.load(fil)

else:

self.dataset = self.create_dataset(pos_dataset)

with open(dataset_path, 'wb') as fil:

pickle.dump(self.dataset, fil)

self.n_batch = math.ceil(len(self.dataset) / batch_size)

self.i_batch = 0

assert(batch_size % (1+neg_rate) == 0)

def create_dataset(self, pos_dataset):

dataset = []

random.shuffle(pos_dataset)

pos_dataset_set = set(pos_dataset)

whole_dataset_set = set(self.whole_dataset)

if self.mode == 'train':

random_ratio = 1#1/3

constrain_ratio = 0#1/3

reverse_ratio = 0#1/3

viewable = 'train'

viewable_set = pos_dataset_set

else:

random_ratio = 1

constrain_ratio = 0

reverse_ratio = 0

viewable = 'all'

viewable_set = whole_dataset_set

for triple in tqdm(pos_dataset):

dataset.append((triple, 1))

choice = random_choose(random_ratio, constrain_ratio, reverse_ratio)

h, r, t = triple

for _ in range(self.neg_rate):

count = 0

while (True):

if (random.sample(range(2), 1)[0] == 1):

# replace head

if choice == 'random':

candidate_ents = self.entity_set - set(self.groundtruth[viewable]['head'][(r, t)])

replace_ent = random.sample(candidate_ents, 1)[0]

neg_triple = (replace_ent, r, t)

elif choice == 'constrain':

candidate_ents = self.possible_entities['train']['head'][r] - set(self.groundtruth[viewable]['head'][(r, t)])

# head that are head of rel, but not head of (rel, tail)

if len(candidate_ents) == 0:

candidate_ents = self.entity_set - set(self.groundtruth[viewable]['head'][(r, t)])

replace_ent = random.sample(candidate_ents, 1)[0]

neg_triple = (replace_ent, r, t)

else: # choice == 'reverse'

neg_triple = (t, r, h)

else:

# replace tail

if choice == 'random':

candidate_ents = self.entity_set - set(self.groundtruth[viewable]['tail'][(r, h)])

replace_ent = random.sample(candidate_ents, 1)[0]

neg_triple = (h, r, replace_ent)

elif choice == 'constrain':

candidate_ents = self.possible_entities['train']['tail'][r] - set(self.groundtruth[viewable]['tail'][(r, h)])

# tail that are tail of rel, but not tail of (rel, head)

if len(candidate_ents) == 0:

candidate_ents = self.entity_set - set(self.groundtruth[viewable]['tail'][(r, h)])

replace_ent = random.sample(candidate_ents, 1)[0]

neg_triple = (h, r, replace_ent)

else: # choice == 'reverse':

neg_triple = (t, r, h)

if neg_triple not in viewable_set:

dataset.append((neg_triple, 0))

break

elif choice == 'reverse':

dataset.append((neg_triple, 1))

return dataset

def __iter__(self):

return self

def __next__(self):

if self.i_batch == self.n_batch:

raise StopIteration()

batch = self.dataset[self.i_batch*self.batch_size: (self.i_batch+1)*self.batch_size]

if self.rdrop:

batch = batch + batch

self.i_batch += 1

return batch

def __len__(self):

return self.n_batch

def get_dataset_size(self):

def __init__(self, in_paths, tokenizer, batch_size = 16, neg_rate = 1, max_desc_length = 512, add_tokens=False, p_tuning=False, rdrop=False, model='bert'):

self.datasetName = in_paths['dataset']

self.train_set = self.load_dataset(in_paths['train'])

if self.datasetName not in ['fb13']:

self.valid_set = self.load_dataset(in_paths['valid'])

self.test_set = self.load_dataset(in_paths['test'])

self.valid_set_with_neg = None

self.test_set_with_neg = None

else:

self.valid_set, self.valid_set_with_neg = self.load_dataset_with_neg(in_paths['valid'])

self.test_set, self.test_set_with_neg = self.load_dataset_with_neg(in_paths['test'])

self.whole_set = self.train_set + self.valid_set + self.test_set

self.uid2text = {}

self.uid2tokens = {}

self.entity_set = set([t[0] for t in (self.train_set + self.valid_set + self.test_set)] + [t[-1] for t in (self.train_set + self.valid_set + self.test_set)])

self.relation_set = set([t[1] for t in (self.train_set + self.valid_set + self.test_set)])

self.tokenizer = tokenizer

for p in in_paths['text']:

self.load_text(p)

self.batch_size = batch_size

self.step_per_epc = math.ceil(len(self.train_set) * (1+neg_rate) / batch_size)

self.train_entity_set = set([t[0] for t in self.train_set] + [t[-1] for t in self.train_set])

self.train_relation_set = set([t[1] for t in self.train_set])

self.entity_list = sorted(self.entity_set)

self.relation_list = sorted(self.relation_set)

self.ent2id = {e:i for i, e in enumerate(sorted(self.entity_set))}

self.rel2id = {r:i for i, r in enumerate(sorted(self.relation_set))}

self.id2ent = {i:e for i, e in enumerate(sorted(self.entity_set))}

self.id2rel = {i:r for i, r in enumerate(sorted(self.relation_set))}

self.neg_rate = neg_rate

self.max_desc_length = max_desc_length

self.groundtruth, self.possible_entities= self.count_groundtruth()

self.add_tokens = add_tokens

self.p_tuning = p_tuning

self.rdrop = rdrop

self.model = model

self.orig_vocab_size = len(tokenizer)

self.count_degrees()

def load_dataset(self, in_path):

dataset = []

with open(in_path, 'r', encoding='utf8') as fil:

for line in fil.readlines():

if in_path[-3:] == 'txt':

h, t, r = line.strip('\n').split('\t')

else:

h, r, t = line.strip('\n').split('\t')

dataset.append((h, r, t))

return dataset

def load_dataset_with_neg(self, in_path):

dataset = []

dataset_with_neg = []

with open(in_path, 'r', encoding='utf8') as fil:

for line in fil.readlines():

h, r, t, l = line.strip('\n').split('\t')

if l == '-1':

l = 0

else:

l = 1

dataset.append((h, r, t))

dataset_with_neg.append((h, r, t, l))

return dataset, dataset_with_neg

def load_name_wiki(self, in_path):

uid2name = {}

with open(in_path, 'r', encoding='utf8') as fil:

for line in fil.readlines():

_ = line.strip('\n').split('\t')

uid = _[0]

name = _[1]

uid2name[uid] = name

return uid2name

def load_text(self, in_path):

uid2text = self.uid2text

uid2tokens = self.uid2tokens

tokenizer = self.tokenizer

with open(in_path, 'r', encoding='utf8') as fil:

for line in fil.readlines():

uid, text = line.strip('\n').split('\t', 1)

text = text.replace('@en', '').strip('"')

if uid not in uid2text.keys():

uid2text[uid] = text

tokens = tokenizer.tokenize(text)

if uid not in uid2tokens.keys():

uid2tokens[uid] = tokens

self.uid2text = uid2text

self.uid2tokens = uid2tokens

def triple_to_text(self, triple, with_text):

tokenizer = self.tokenizer

ent2id = self.ent2id

rel2id = self.rel2id

if True:

# 512 tokens, 1 CLS, 1 SEP, 1 head, 1 rel, 1 tail, so 507 remaining.

h_n_tokens = min(241, self.max_desc_length)

t_n_tokens = min(241, self.max_desc_length)

r_n_tokens = min(16, self.max_desc_length)

h, r, t = triple

h_text_tokens = self.uid2tokens.get(h, [])[:h_n_tokens] if with_text['h'] else []

r_text_tokens = self.uid2tokens.get(r, [])[:r_n_tokens] if with_text['r'] else []

t_text_tokens = self.uid2tokens.get(t, [])[:t_n_tokens] if with_text['t'] else []

if self.add_tokens:

if self.p_tuning:

h_token = ["[head_b1]", "[head_b2]"] + (['[ent_{}]'.format(ent2id[h])] if with_text['h'] else [tokenizer.mask_token]) + ["[head_a1]", "[head_a2]"]

r_token = ["[rel_b1]", "[rel_b2]"] + (['[rel_{}]'.format(rel2id[r])] if with_text['r'] else [tokenizer.mask_token]) + ["[rel_a1]", "[rel_a2]"]

t_token = ["[tail_b1]", "[tail_b2]"] + (['[ent_{}]'.format(ent2id[t])] if with_text['t'] else [tokenizer.mask_token]) + ["[tail_a1]", "[tail_a2]"]

else:

h_token = ['[ent_{}]'.format(ent2id[h])] if with_text['h'] else [tokenizer.mask_token]

r_token = ['[rel_{}]'.format(rel2id[r])] if with_text['r'] else [tokenizer.mask_token]

t_token = ['[ent_{}]'.format(ent2id[t])] if with_text['t'] else [tokenizer.mask_token]

else:

h_token = [self.tokenizer.cls_token] if with_text['h'] else [tokenizer.mask_token]

r_token = [self.tokenizer.cls_token] if with_text['r'] else [tokenizer.mask_token]

t_token = [self.tokenizer.cls_token] if with_text['t'] else [tokenizer.mask_token]

tokens = h_token + h_text_tokens + r_token + r_text_tokens + t_token + t_text_tokens

text = tokenizer.convert_tokens_to_string(tokens)

return text, tokens

def element_to_text(self, target):

tokenizer = self.tokenizer

ent2id = self.ent2id

rel2id = self.rel2id

n_tokens = min(508, self.max_desc_length)

text_tokens = self.uid2tokens.get(target, [])[:n_tokens]

if self.add_tokens:

if target in ent2id.keys():

token = ['[ent_{}]'.format(ent2id[target])]

else:

token = ['[rel_{}]'.format(rel2id[target])]

else:

token = [self.tokenizer.cls_token]#'[CLS]']

tokens = token + text_tokens

text = tokenizer.convert_tokens_to_string(tokens)

return text, tokens

def batch_tokenize(self, batch_triples, mode):

batch_texts = []

batch_tokens = []

batch_positions = []

ent2id = self.ent2id

rel2id = self.rel2id

if mode in ['triple_classification']:

with_text = {'h': True, 'r': True, 't': True}

elif mode == "link_prediction_h":

with_text = {'h': False, 'r': True, 't': True}

elif mode == "link_prediction_r":

with_text = {'h': True, 'r': False, 't': True}

elif mode == "link_prediction_t":

with_text = {'h': True, 'r': True, 't': False}

for triple in batch_triples:

text, tokens = self.triple_to_text(triple, with_text)

batch_texts.append(text)

batch_tokens.append(tokens)

h, r, t = triple

#batch_tokens_ = self.tokenizer(batch_texts, truncation = True, max_length = 512, return_tensors='pt', padding=True )

batch_tokens = self.my_tokenize(batch_tokens, max_length=512, padding=True, model=self.model)

orig_vocab_size = self.orig_vocab_size

num_ent_rel_tokens = len(ent2id) + len(rel2id)

mask_idx = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)

cls_idx = self.tokenizer.convert_tokens_to_ids(self.tokenizer.cls_token)

for i, _ in enumerate(batch_tokens['input_ids']):

triple = batch_triples[i]

h, r, t = triple

#import pdb

#pdb.set_trace()

if not self.add_tokens:

cls_pos, h_pos, r_pos, t_pos = torch.where((_==mask_idx) + (_==cls_idx))[0]

else:

h_pos, r_pos, t_pos = torch.where( (_ >= orig_vocab_size) * (_ < orig_vocab_size + num_ent_rel_tokens) + (_ == mask_idx) )[0]

batch_positions.append({'head': (ent2id[h], h_pos.item()), 'rel': (rel2id[r], r_pos.item()), 'tail': (ent2id[t], t_pos.item())})

return batch_tokens, batch_positions

def batch_tokenize_target(self, batch_triples=None, mode=None, targets = None):

batch_texts = []

batch_tokens = []

batch_positions = []

ent2id = self.ent2id

rel2id = self.rel2id

if targets == None:

if mode == "link_prediction_h":

targets = [ triple[0] for triple in batch_triples]

elif mode == "link_prediction_r":

targets = [ triple[1] for triple in batch_triples]

elif mode == "link_prediction_t":

targets = [ triple[2] for triple in batch_triples]

for target in targets:

text, tokens = self.element_to_text(target)

batch_texts.append(text)

batch_tokens.append(tokens)

batch_tokens = self.my_tokenize(batch_tokens, max_length=512, padding=True, model=self.model)

for i, _ in enumerate(batch_tokens['input_ids']):

target = targets[i]

target_pos = 1

if target in ent2id.keys():

target_idx = ent2id[target]

else:

target_idx = rel2id[target]

batch_positions.append( (target_idx, target_pos) )

return batch_tokens, batch_positions

def train_data_sampler(self):

return DataSampler(datasetName = self.datasetName, mode='train', pos_dataset=self.train_set, whole_dataset=self.whole_set, batch_size=self.batch_size,

entity_set=self.train_entity_set, relation_set=self.train_relation_set, neg_rate=self.neg_rate, groundtruth=self.groundtruth,

possible_entities=self.possible_entities, rdrop=self.rdrop)

def valid_data_sampler(self, batch_size, neg_rate):

return DataSampler(datasetName = self.datasetName, mode='valid', pos_dataset=self.valid_set, whole_dataset=self.whole_set, batch_size=batch_size,

entity_set=self.entity_set, relation_set=self.relation_set, neg_rate=neg_rate, groundtruth=self.groundtruth, pos_neg_dataset=self.valid_set_with_neg)

def test_data_sampler(self, batch_size, neg_rate):

return DataSampler(datasetName = self.datasetName, mode='test', pos_dataset=self.test_set, whole_dataset=self.whole_set, batch_size=batch_size,

entity_set=self.entity_set, relation_set=self.relation_set, neg_rate=neg_rate, groundtruth=self.groundtruth, pos_neg_dataset=self.test_set_with_neg)

def get_dataset_size(self, split='train'):

if split == 'train':

return len(self.train_set) * (1+self.neg_rate)

def count_groundtruth(self):

groundtruth = { split: {'head': {}, 'rel': {}, 'tail': {}} for split in ['all', 'train', 'valid', 'test']}

possible_entities = { split: {'head': {}, 'tail': {}} for split in ['train']}

for triple in self.train_set:

h, r, t = triple

groundtruth['all']['head'].setdefault((r, t), [])

groundtruth['all']['head'][(r, t)].append(h)

groundtruth['all']['tail'].setdefault((r, h), [])

groundtruth['all']['tail'][(r, h)].append(t)

groundtruth['all']['rel'].setdefault((h, t), [])

groundtruth['all']['rel'][(h, t)].append(r)

groundtruth['train']['head'].setdefault((r, t), [])

groundtruth['train']['head'][(r, t)].append(h)

groundtruth['train']['tail'].setdefault((r, h), [])

groundtruth['train']['tail'][(r, h)].append(t)

groundtruth['train']['rel'].setdefault((h, t), [])

groundtruth['train']['rel'][(h, t)].append(r)

possible_entities['train']['head'].setdefault(r, set())

possible_entities['train']['head'][r].add(h)

possible_entities['train']['tail'].setdefault(r, set())

possible_entities['train']['tail'][r].add(t)

for triple in self.valid_set:

h, r, t = triple

groundtruth['all']['head'].setdefault((r, t), [])

groundtruth['all']['head'][(r, t)].append(h)

groundtruth['all']['tail'].setdefault((r, h), [])

groundtruth['all']['tail'][(r, h)].append(t)

groundtruth['all']['rel'].setdefault((h, t), [])

groundtruth['all']['rel'][(h, t)].append(r)

groundtruth['valid']['head'].setdefault((r, t), [])

groundtruth['valid']['head'][(r, t)].append(h)

groundtruth['valid']['tail'].setdefault((r, h), [])

groundtruth['valid']['tail'][(r, h)].append(t)

for triple in self.test_set:

h, r, t = triple

groundtruth['all']['head'].setdefault((r, t), [])

groundtruth['all']['head'][(r, t)].append(h)

groundtruth['all']['tail'].setdefault((r, h), [])

groundtruth['all']['tail'][(r, h)].append(t)

groundtruth['all']['rel'].setdefault((h, t), [])

groundtruth['all']['rel'][(h, t)].append(r)

groundtruth['test']['head'].setdefault((r, t), [])

groundtruth['test']['head'][(r, t)].append(h)

groundtruth['test']['tail'].setdefault((r, h), [])

groundtruth['test']['tail'][(r, h)].append(t)

return groundtruth, possible_entities

def get_groundtruth(self):

return self.groundtruth

def get_dataset(self, split):

assert (split in ['train', 'valid', 'test'])

if split == 'train':

return self.train_set

elif split == 'valid':

return self.valid_set

elif split == 'test':

return self.test_set

def my_tokenize(self, batch_tokens, max_length=512, padding=True, model='roberta'):

'''

if model == 'roberta':

start_tokens = ['<s>']

end_tokens = ['</s>']

pad_token = '<pad>'

elif model == 'bert':

start_tokens = ['[CLS]']

end_tokens = ['[SEP]']

pad_token = '[PAD]'

'''

start_tokens = [self.tokenizer.cls_token]

end_tokens = [self.tokenizer.sep_token]

batch_tokens = [ start_tokens + i + end_tokens for i in batch_tokens]

batch_size = len(batch_tokens)

longest = min(max([len(i) for i in batch_tokens]), 512)

if model == 'bert':

input_ids = torch.zeros((batch_size, longest)).long()

elif model == 'roberta':

input_ids = torch.ones((batch_size, longest)).long()

token_type_ids = torch.zeros((batch_size, longest)).long()

attention_mask = torch.zeros((batch_size, longest)).long()

for i in range(batch_size):

tokens = self.tokenizer.convert_tokens_to_ids(batch_tokens[i])

input_ids[i, :len(tokens)] = torch.tensor(tokens).long()

attention_mask[i, :len(tokens)] = 1

if model == 'roberta':

return BatchEncoding(data = {'input_ids': input_ids, 'attention_mask': attention_mask})

elif model == 'bert':

return BatchEncoding(data = {'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids})

def adding_tokens(self):

n_ent = len(self.ent2id)

n_rel = len(self.rel2id)

new_tokens = ["[ent_{}]".format(i) for i in range(n_ent)] + ["[rel_{}]".format(i) for i in range(n_rel)]

if self.p_tuning:

new_tokens += ["[head_b1]", "[head_b2]", "[head_a1]", "[head_a2]",

"[rel_b1]", "[rel_b2]", "[rel_a1]", "[rel_a2]",

"[tail_b1]", "[tail_b2]", "[tail_a1]", "[tail_a2]"] # continuous prompt

self.tokenizer.add_tokens(new_tokens)

def count_degrees(self):

train_set = self.train_set #+ self.valid_set + self.test_set

degrees = {}

for triple in train_set:

h, r, t=triple

degrees[h] = degrees.get(h, 0) + 1

degrees[t] = degrees.get(t, 0) + 1

degrees[r] = degrees.get(r, 0) + 1

raw_degrees = copy.deepcopy(degrees)

max_degree = 0

for k, v in degrees.items():

max_degree = max(max_degree, v)

max_degree = math.floor(math.log(max_degree) / math.log(2))

count_degree_group = { i:0 for i in range(0, max_degree+1)}

for k, v in degrees.items():

degrees[k] = math.floor(math.log(v) / math.log(2)) + 1

count_degree_group[degrees[k]] = count_degree_group.get(degrees[k], 0) + 1

self.statistics = {

'degrees': raw_degrees,

'degree_group': degrees,

'count_degree_group': count_degree_group,

'max_degree': max_degree

x = random.random()

if x <= random_ratio:

return 'random'

elif x <= (random_ratio + constrain_ratio):

return 'constrain'

else: