947 lines
35 KiB
Python
947 lines
35 KiB
Python
from re import DEBUG
|
||
|
||
import contextlib
|
||
import sys
|
||
|
||
from collections import Counter
|
||
from multiprocessing import Pool
|
||
from torch._C import HOIST_CONV_PACKED_PARAMS
|
||
from torch.utils.data import Dataset, Sampler, IterableDataset
|
||
from collections import defaultdict
|
||
from functools import partial
|
||
from multiprocessing import Pool
|
||
import os
|
||
import random
|
||
import json
|
||
import torch
|
||
import copy
|
||
import numpy as np
|
||
import pickle
|
||
from tqdm import tqdm
|
||
from dataclasses import dataclass, asdict, replace
|
||
import inspect
|
||
|
||
from transformers.models.auto.tokenization_auto import AutoTokenizer
|
||
|
||
from models.utils import get_entity_spans_pre_processing
|
||
|
||
def lmap(a, b):
|
||
return list(map(a,b))
|
||
|
||
def cache_results(_cache_fp, _refresh=False, _verbose=1):
|
||
r"""
|
||
cache_results是fastNLP中用于cache数据的装饰器。通过下面的例子看一下如何使用::
|
||
|
||
import time
|
||
import numpy as np
|
||
from fastNLP import cache_results
|
||
|
||
@cache_results('cache.pkl')
|
||
def process_data():
|
||
# 一些比较耗时的工作，比如读取数据，预处理数据等，这里用time.sleep()代替耗时
|
||
time.sleep(1)
|
||
return np.random.randint(10, size=(5,))
|
||
|
||
start_time = time.time()
|
||
print("res =",process_data())
|
||
print(time.time() - start_time)
|
||
|
||
start_time = time.time()
|
||
print("res =",process_data())
|
||
print(time.time() - start_time)
|
||
|
||
# 输出内容如下，可以看到两次结果相同，且第二次几乎没有花费时间
|
||
# Save cache to cache.pkl.
|
||
# res = [5 4 9 1 8]
|
||
# 1.0042750835418701
|
||
# Read cache from cache.pkl.
|
||
# res = [5 4 9 1 8]
|
||
# 0.0040721893310546875
|
||
|
||
可以看到第二次运行的时候，只用了0.0001s左右，是由于第二次运行将直接从cache.pkl这个文件读取数据，而不会经过再次预处理::
|
||
|
||
# 还是以上面的例子为例，如果需要重新生成另一个cache，比如另一个数据集的内容，通过如下的方式调用即可
|
||
process_data(_cache_fp='cache2.pkl') # 完全不影响之前的‘cache.pkl'
|
||
|
||
上面的_cache_fp是cache_results会识别的参数，它将从'cache2.pkl'这里缓存/读取数据，即这里的'cache2.pkl'覆盖默认的
|
||
'cache.pkl'。如果在你的函数前面加上了@cache_results()则你的函数会增加三个参数[_cache_fp, _refresh, _verbose]。
|
||
上面的例子即为使用_cache_fp的情况，这三个参数不会传入到你的函数中，当然你写的函数参数名也不可能包含这三个名称::
|
||
|
||
process_data(_cache_fp='cache2.pkl', _refresh=True) # 这里强制重新生成一份对预处理的cache。
|
||
# _verbose是用于控制输出信息的，如果为0,则不输出任何内容;如果为1,则会提醒当前步骤是读取的cache还是生成了新的cache
|
||
|
||
:param str _cache_fp: 将返回结果缓存到什么位置;或从什么位置读取缓存。如果为None，cache_results没有任何效用，除非在
|
||
函数调用的时候传入_cache_fp这个参数。
|
||
:param bool _refresh: 是否重新生成cache。
|
||
:param int _verbose: 是否打印cache的信息。
|
||
:return:
|
||
"""
|
||
|
||
def wrapper_(func):
|
||
signature = inspect.signature(func)
|
||
for key, _ in signature.parameters.items():
|
||
if key in ('_cache_fp', '_refresh', '_verbose'):
|
||
raise RuntimeError("The function decorated by cache_results cannot have keyword `{}`.".format(key))
|
||
|
||
def wrapper(*args, **kwargs):
|
||
my_args = args[0]
|
||
mode = args[-1]
|
||
if '_cache_fp' in kwargs:
|
||
cache_filepath = kwargs.pop('_cache_fp')
|
||
assert isinstance(cache_filepath, str), "_cache_fp can only be str."
|
||
else:
|
||
cache_filepath = _cache_fp
|
||
if '_refresh' in kwargs:
|
||
refresh = kwargs.pop('_refresh')
|
||
assert isinstance(refresh, bool), "_refresh can only be bool."
|
||
else:
|
||
refresh = _refresh
|
||
if '_verbose' in kwargs:
|
||
verbose = kwargs.pop('_verbose')
|
||
assert isinstance(verbose, int), "_verbose can only be integer."
|
||
else:
|
||
verbose = _verbose
|
||
refresh_flag = True
|
||
|
||
model_name = my_args.model_name_or_path.split("/")[-1]
|
||
is_pretrain = my_args.pretrain
|
||
cache_filepath = os.path.join(my_args.data_dir, f"cached_{mode}_features{model_name}_pretrain{is_pretrain}_faiss{my_args.faiss_init}_seqlength{my_args.max_seq_length}_{my_args.litmodel_class}.pkl")
|
||
refresh = my_args.overwrite_cache
|
||
|
||
if cache_filepath is not None and refresh is False:
|
||
# load data
|
||
if os.path.exists(cache_filepath):
|
||
with open(cache_filepath, 'rb') as f:
|
||
results = pickle.load(f)
|
||
if verbose == 1:
|
||
logger.info("Read cache from {}.".format(cache_filepath))
|
||
refresh_flag = False
|
||
|
||
if refresh_flag:
|
||
results = func(*args, **kwargs)
|
||
if cache_filepath is not None:
|
||
if results is None:
|
||
raise RuntimeError("The return value is None. Delete the decorator.")
|
||
with open(cache_filepath, 'wb') as f:
|
||
pickle.dump(results, f)
|
||
logger.info("Save cache to {}.".format(cache_filepath))
|
||
|
||
return results
|
||
|
||
return wrapper
|
||
|
||
return wrapper_
|
||
|
||
|
||
import argparse
|
||
import csv
|
||
import logging
|
||
import os
|
||
import random
|
||
import sys
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
|
||
TensorDataset)
|
||
from torch.utils.data.distributed import DistributedSampler
|
||
from tqdm import tqdm, trange
|
||
|
||
# from torch.nn import CrossEntropyLoss, MSELoss
|
||
# from scipy.stats import pearsonr, spearmanr
|
||
# from sklearn.metrics import matthews_corrcoef, f1_scoreclass
|
||
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class InputExample(object):
|
||
"""A single training/test example for simple sequence classification."""
|
||
|
||
def __init__(self, guid, text_a, text_b=None, text_c=None, label=None, real_label=None, en=None, rel=None):
|
||
"""Constructs a InputExample.
|
||
|
||
Args:
|
||
guid: Unique id for the example.
|
||
text_a: string. The untokenized text of the first sequence. For single
|
||
sequence tasks, only this sequence must be specified.
|
||
text_b: (Optional) string. The untokenized text of the second sequence.
|
||
Only must be specified for sequence pair tasks.
|
||
text_c: (Optional) string. The untokenized text of the third sequence.
|
||
Only must be specified for sequence triple tasks.
|
||
label: (Optional) string. list of entities
|
||
"""
|
||
self.guid = guid
|
||
self.text_a = text_a
|
||
self.text_b = text_b
|
||
self.text_c = text_c
|
||
self.label = label
|
||
self.real_label = real_label
|
||
self.en = en
|
||
self.rel = rel # rel id
|
||
|
||
|
||
@dataclass
|
||
class InputFeatures:
|
||
"""A single set of features of data."""
|
||
|
||
input_ids: torch.Tensor
|
||
attention_mask: torch.Tensor
|
||
labels: torch.Tensor = None
|
||
label: torch.Tensor = None
|
||
en: torch.Tensor = 0
|
||
rel: torch.Tensor = 0
|
||
pos: torch.Tensor = 0
|
||
|
||
|
||
class DataProcessor(object):
|
||
"""Base class for data converters for sequence classification data sets."""
|
||
|
||
def get_train_examples(self, data_dir):
|
||
"""Gets a collection of `InputExample`s for the train set."""
|
||
raise NotImplementedError()
|
||
|
||
def get_dev_examples(self, data_dir):
|
||
"""Gets a collection of `InputExample`s for the dev set."""
|
||
raise NotImplementedError()
|
||
|
||
def get_labels(self, data_dir):
|
||
"""Gets the list of labels for this data set."""
|
||
raise NotImplementedError()
|
||
|
||
@classmethod
|
||
def _read_tsv(cls, input_file, quotechar=None):
|
||
"""Reads a tab separated value file."""
|
||
with open(input_file, "r", encoding="utf-8") as f:
|
||
reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
|
||
lines = []
|
||
for line in reader:
|
||
if sys.version_info[0] == 2:
|
||
line = list(unicode(cell, 'utf-8') for cell in line)
|
||
lines.append(line)
|
||
return lines
|
||
|
||
import copy
|
||
|
||
|
||
def solve_get_knowledge_store(line, set_type="train", pretrain=1):
|
||
"""
|
||
use the LM to get the entity embedding.
|
||
Transductive: triples + text description
|
||
Inductive: text description
|
||
|
||
"""
|
||
examples = []
|
||
|
||
head_ent_text = ent2text[line[0]]
|
||
tail_ent_text = ent2text[line[2]]
|
||
relation_text = rel2text[line[1]]
|
||
|
||
i=0
|
||
|
||
a = tail_filter_entities["\t".join([line[0],line[1]])]
|
||
b = head_filter_entities["\t".join([line[2],line[1]])]
|
||
|
||
guid = "%s-%s" % (set_type, i)
|
||
text_a = head_ent_text
|
||
text_b = relation_text
|
||
text_c = tail_ent_text
|
||
|
||
# use the description of c to predict A
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[PAD]", text_b=text_b + "[PAD]", text_c = "[PAD]" + " " + text_c, label=lmap(lambda x: ent2id[x], b), real_label=ent2id[line[0]], en=[ent2id[line[0]], rel2id[line[1]], ent2id[line[2]]], rel=0)
|
||
)
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[PAD]", text_b=text_b + "[PAD]", text_c = "[PAD]" + " " + text_a, label=lmap(lambda x: ent2id[x], b), real_label=ent2id[line[2]], en=[ent2id[line[0]], rel2id[line[1]], ent2id[line[2]]], rel=0)
|
||
)
|
||
return examples
|
||
|
||
|
||
def solve(line, set_type="train", pretrain=1):
|
||
examples = []
|
||
|
||
head_ent_text = ent2text[line[0]]
|
||
tail_ent_text = ent2text[line[2]]
|
||
relation_text = rel2text[line[1]]
|
||
|
||
i=0
|
||
|
||
a = tail_filter_entities["\t".join([line[0],line[1]])]
|
||
b = head_filter_entities["\t".join([line[2],line[1]])]
|
||
|
||
guid = "%s-%s" % (set_type, i)
|
||
text_a = head_ent_text
|
||
text_b = relation_text
|
||
text_c = tail_ent_text
|
||
|
||
|
||
if pretrain:
|
||
text_a_tokens = text_a.split()
|
||
for i in range(10):
|
||
st = random.randint(0, len(text_a_tokens))
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[MASK]", text_b=" ".join(text_a_tokens[st:min(st+64, len(text_a_tokens))]), text_c = "", label=ent2id[line[0]], real_label=ent2id[line[0]], en=0, rel=0)
|
||
)
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[MASK]", text_b=text_a, text_c = "", label=ent2id[line[0]], real_label=ent2id[line[0]], en=0, rel=0)
|
||
)
|
||
# examples.append(
|
||
# InputExample(guid=guid, text_a="[MASK]", text_b=text_c, text_c = "", label=ent2id[line[2]], real_label=ent2id[line[2]], en=0, rel=0)
|
||
# )
|
||
else:
|
||
|
||
# examples.append(
|
||
# InputExample(guid=guid, text_a="[MASK]", text_b=text_b + "[PAD]", text_c = "[UNK]" , label=lmap(lambda x: ent2id[x], b), real_label=ent2id[line[0]], en=ent2id[line[2]], rel=rel2id[line[1]]))
|
||
# examples.append(
|
||
# InputExample(guid=guid, text_a="[UNK] ", text_b=text_b + "[PAD]", text_c = "[MASK]", label=lmap(lambda x: ent2id[x], a), real_label=ent2id[line[2]], en=ent2id[line[0]], rel=rel2id[line[1]]))
|
||
|
||
# examples.append(
|
||
# InputExample(guid=guid, text_a="[UNK]" + " " + text_c, text_b=text_b + "[PAD]", text_c = "[MASK]", label=lmap(lambda x: ent2id[x], b), real_label=ent2id[line[0]], en=ent2id[line[2]], rel=rel2id[line[1]]))
|
||
# examples.append(
|
||
# InputExample(guid=guid, text_a="[MASK]", text_b=text_b + "[PAD]", text_c = "[UNK]" + text_a, label=lmap(lambda x: ent2id[x], a), real_label=ent2id[line[2]], en=ent2id[line[0]], rel=rel2id[line[1]]))
|
||
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[MASK]", text_b=text_b + "[PAD]", text_c = "[PAD]" + " " + text_c, label=lmap(lambda x: ent2id[x], b), real_label=ent2id[line[0]], en=[rel2id[line[1]], ent2id[line[2]]], rel=rel2id[line[1]]))
|
||
examples.append(
|
||
InputExample(guid=guid, text_a="[PAD] ", text_b=text_b + "[PAD]", text_c = "[MASK]" +" " + text_a, label=lmap(lambda x: ent2id[x], a), real_label=ent2id[line[2]], en=[ent2id[line[0]], rel2id[line[1]]], rel=rel2id[line[1]]))
|
||
return examples
|
||
|
||
def filter_init(head, tail, t1,t2, ent2id_, ent2token_, rel2id_):
|
||
global head_filter_entities
|
||
global tail_filter_entities
|
||
global ent2text
|
||
global rel2text
|
||
global ent2id
|
||
global ent2token
|
||
global rel2id
|
||
global negativeEntity
|
||
|
||
head_filter_entities = head
|
||
tail_filter_entities = tail
|
||
ent2text =t1
|
||
rel2text =t2
|
||
ent2id = ent2id_
|
||
ent2token = ent2token_
|
||
rel2id = rel2id_
|
||
negativeEntity = ent2id['[NEG]']
|
||
|
||
def delete_init(ent2text_):
|
||
global ent2text
|
||
ent2text = ent2text_
|
||
|
||
def getEntityIdByName(name):
|
||
global ent2id
|
||
return ent2id[name]
|
||
|
||
def getNegativeEntityId():
|
||
global negativeEntity
|
||
return negativeEntity
|
||
|
||
class KGProcessor(DataProcessor):
|
||
"""Processor for knowledge graph data set."""
|
||
def __init__(self, tokenizer, args):
|
||
self.labels = set()
|
||
self.tokenizer = tokenizer
|
||
self.args = args
|
||
self.entity_path = os.path.join(args.data_dir, "entity2textlong.txt") if os.path.exists(os.path.join(args.data_dir, 'entity2textlong.txt')) \
|
||
else os.path.join(args.data_dir, "entity2text.txt")
|
||
|
||
def get_train_examples(self, data_dir):
|
||
"""See base class."""
|
||
return self._create_examples(
|
||
self._read_tsv(os.path.join(data_dir, "train.tsv")), "train", data_dir, self.args)
|
||
|
||
def get_dev_examples(self, data_dir):
|
||
"""See base class."""
|
||
return self._create_examples(
|
||
self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev", data_dir, self.args)
|
||
|
||
def get_test_examples(self, data_dir, chunk=""):
|
||
"""See base class."""
|
||
return self._create_examples(
|
||
self._read_tsv(os.path.join(data_dir, f"test{chunk}.tsv")), "test", data_dir, self.args)
|
||
|
||
def get_relations(self, data_dir):
|
||
"""Gets all labels (relations) in the knowledge graph."""
|
||
# return list(self.labels)
|
||
with open(os.path.join(data_dir, "relations.txt"), 'r') as f:
|
||
lines = f.readlines()
|
||
relations = []
|
||
for line in lines:
|
||
relations.append(line.strip().split('\t')[0])
|
||
rel2token = {ent : f"[RELATION_{i}]" for i, ent in enumerate(relations)}
|
||
return list(rel2token.values())
|
||
|
||
def get_labels(self, data_dir):
|
||
"""Gets all labels (0, 1) for triples in the knowledge graph."""
|
||
relation = []
|
||
with open(os.path.join(data_dir, "relation2text.txt"), 'r') as f:
|
||
lines = f.readlines()
|
||
entities = []
|
||
for line in lines:
|
||
relation.append(line.strip().split("\t")[-1])
|
||
return relation
|
||
|
||
def get_entities(self, data_dir):
|
||
"""Gets all entities in the knowledge graph."""
|
||
with open(self.entity_path, 'r') as f:
|
||
lines = f.readlines()
|
||
lines.append('[NEG]\t')
|
||
entities = []
|
||
for line in lines:
|
||
entities.append(line.strip().split("\t")[0])
|
||
|
||
ent2token = {ent : f"[ENTITY_{i}]" for i, ent in enumerate(entities)}
|
||
return list(ent2token.values())
|
||
|
||
def get_train_triples(self, data_dir):
|
||
"""Gets training triples."""
|
||
return self._read_tsv(os.path.join(data_dir, "train.tsv"))
|
||
|
||
def get_dev_triples(self, data_dir):
|
||
"""Gets validation triples."""
|
||
return self._read_tsv(os.path.join(data_dir, "dev.tsv"))
|
||
|
||
def get_test_triples(self, data_dir, chunk=""):
|
||
"""Gets test triples."""
|
||
return self._read_tsv(os.path.join(data_dir, f"test{chunk}.tsv"))
|
||
|
||
def _create_examples(self, lines, set_type, data_dir, args):
|
||
"""Creates examples for the training and dev sets."""
|
||
# entity to text
|
||
ent2text = {}
|
||
ent2text_with_type = {}
|
||
with open(self.entity_path, 'r') as f:
|
||
ent_lines = f.readlines()
|
||
ent_lines.append('[NEG]\t')
|
||
for line in ent_lines:
|
||
temp = line.strip().split('\t')
|
||
try:
|
||
end = temp[1]#.find(',')
|
||
if "wiki" in data_dir:
|
||
assert "Q" in temp[0]
|
||
ent2text[temp[0]] = temp[1].replace("\\n", " ").replace("\\", "") #[:end]
|
||
except IndexError:
|
||
# continue
|
||
end = " "#.find(',')
|
||
if "wiki" in data_dir:
|
||
assert "Q" in temp[0]
|
||
ent2text[temp[0]] = end #[:end]
|
||
|
||
entities = list(ent2text.keys())
|
||
ent2token = {ent : f"[ENTITY_{i}]" for i, ent in enumerate(entities)}
|
||
ent2id = {ent : i for i, ent in enumerate(entities)}
|
||
|
||
rel2text = {}
|
||
with open(os.path.join(data_dir, "relation2text.txt"), 'r') as f:
|
||
rel_lines = f.readlines()
|
||
for line in rel_lines:
|
||
temp = line.strip().split('\t')
|
||
rel2text[temp[0]] = temp[1]
|
||
relation_names = {}
|
||
with open(os.path.join(data_dir, "relations.txt"), "r") as file:
|
||
for line in file.readlines():
|
||
t = line.strip()
|
||
relation_names[t] = rel2text[t]
|
||
|
||
tmp_lines = []
|
||
not_in_text = 0
|
||
for line in tqdm(lines, desc="delete entities without text name."):
|
||
if (line[0] not in ent2text) or (line[2] not in ent2text) or (line[1] not in rel2text):
|
||
not_in_text += 1
|
||
continue
|
||
tmp_lines.append(line)
|
||
lines = tmp_lines
|
||
print(f"total entity not in text : {not_in_text} ")
|
||
|
||
# rel id -> relation token id
|
||
num_entities = len(self.get_entities(args.data_dir))
|
||
rel2id = {w:i+num_entities for i,w in enumerate(relation_names.keys())}
|
||
|
||
# add reverse relation
|
||
# tmp_rel2id = {}
|
||
# num_relations = len(rel2id)
|
||
# cnt = 0
|
||
# for k, v in rel2id.items():
|
||
# tmp_rel2id[k + " (reverse)"] = num_relations + cnt
|
||
# cnt += 1
|
||
# rel2id.update(tmp_rel2id)
|
||
|
||
examples = []
|
||
# head filter head entity
|
||
head_filter_entities = defaultdict(list)
|
||
tail_filter_entities = defaultdict(list)
|
||
|
||
dataset_list = ["train.tsv", "dev.tsv", "test.tsv"]
|
||
# in training, only use the train triples
|
||
if set_type == "train" and not args.pretrain: dataset_list = dataset_list[0:1]
|
||
for m in dataset_list:
|
||
with open(os.path.join(data_dir, m), 'r') as file:
|
||
train_lines = file.readlines()
|
||
for idx in range(len(train_lines)):
|
||
train_lines[idx] = train_lines[idx].strip().split("\t")
|
||
|
||
for line in train_lines:
|
||
tail_filter_entities["\t".join([line[0], line[1]])].append(line[2])
|
||
head_filter_entities["\t".join([line[2], line[1]])].append(line[0])
|
||
|
||
|
||
|
||
max_head_entities = max(len(_) for _ in head_filter_entities.values())
|
||
max_tail_entities = max(len(_) for _ in tail_filter_entities.values())
|
||
|
||
|
||
# use bce loss, ignore the mlm
|
||
if set_type == "train" and args.bce:
|
||
lines = []
|
||
for k, v in tail_filter_entities.items():
|
||
h, r = k.split('\t')
|
||
t = v[0]
|
||
lines.append([h, r, t])
|
||
for k, v in head_filter_entities.items():
|
||
t, r = k.split('\t')
|
||
h = v[0]
|
||
lines.append([h, r, t])
|
||
|
||
|
||
# for training , select each entity as for get mask embedding.
|
||
if args.pretrain:
|
||
rel = list(rel2text.keys())[0]
|
||
lines = []
|
||
for k in ent2text.keys():
|
||
lines.append([k, rel, k])
|
||
|
||
print(f"max number of filter entities : {max_head_entities} {max_tail_entities}")
|
||
|
||
from os import cpu_count
|
||
threads = min(1, cpu_count())
|
||
filter_init(head_filter_entities, tail_filter_entities,ent2text, rel2text, ent2id, ent2token, rel2id
|
||
)
|
||
|
||
if hasattr(args, "faiss_init") and args.faiss_init:
|
||
annotate_ = partial(
|
||
solve_get_knowledge_store,
|
||
pretrain=self.args.pretrain
|
||
)
|
||
else:
|
||
annotate_ = partial(
|
||
solve,
|
||
pretrain=self.args.pretrain
|
||
)
|
||
examples = list(
|
||
tqdm(
|
||
map(annotate_, lines),
|
||
total=len(lines),
|
||
desc="convert text to examples"
|
||
)
|
||
)
|
||
|
||
|
||
# with Pool(threads, initializer=filter_init, initargs=(head_filter_entities, tail_filter_entities,ent2text, rel2text,
|
||
# ent2text_with_type, rel2id,)) as pool:
|
||
# annotate_ = partial(
|
||
# solve,
|
||
# )
|
||
# examples = list(
|
||
# tqdm(
|
||
# map(annotate_, lines, chunksize= 128),
|
||
# total=len(lines),
|
||
# desc="convert text to examples"
|
||
# )
|
||
# )
|
||
tmp_examples = []
|
||
for e in examples:
|
||
for ee in e:
|
||
tmp_examples.append(ee)
|
||
examples = tmp_examples
|
||
# delete vars
|
||
del head_filter_entities, tail_filter_entities, ent2text, rel2text, ent2id, ent2token, rel2id
|
||
return examples
|
||
|
||
class Verbalizer(object):
|
||
def __init__(self, args):
|
||
if "WN18RR" in args.data_dir:
|
||
self.mode = "WN18RR"
|
||
elif "FB15k" in args.data_dir:
|
||
self.mode = "FB15k"
|
||
elif "umls" in args.data_dir:
|
||
self.mode = "umls"
|
||
elif "codexs" in args.data_dir:
|
||
self.mode = "codexs"
|
||
elif "codexl" in args.data_dir:
|
||
self.mode = "codexl"
|
||
elif "FB13" in args.data_dir:
|
||
self.mode = "FB13"
|
||
elif "WN11" in args.data_dir:
|
||
self.mode = "WN11"
|
||
|
||
|
||
def _convert(self, head, relation, tail):
|
||
if self.mode == "umls":
|
||
return f"The {relation} {head} is "
|
||
|
||
return f"{head} {relation}"
|
||
|
||
|
||
class KGCDataset(Dataset):
|
||
def __init__(self, features):
|
||
self.features = features
|
||
|
||
def __getitem__(self, index):
|
||
return self.features[index]
|
||
|
||
def __len__(self):
|
||
return len(self.features)
|
||
|
||
def convert_examples_to_features_init(tokenizer_for_convert):
|
||
global tokenizer
|
||
tokenizer = tokenizer_for_convert
|
||
|
||
def convert_examples_to_features(example, max_seq_length, mode, pretrain=1):
|
||
"""Loads a data file into a list of `InputBatch`s."""
|
||
# tokens_a = tokenizer.tokenize(example.text_a)
|
||
# tokens_b = tokenizer.tokenize(example.text_b)
|
||
# tokens_c = tokenizer.tokenize(example.text_c)
|
||
|
||
# _truncate_seq_triple(tokens_a, tokens_b, tokens_c, max_length= max_seq_length)
|
||
text_a = " ".join(example.text_a.split()[:128])
|
||
text_b = " ".join(example.text_b.split()[:128])
|
||
text_c = " ".join(example.text_c.split()[:128])
|
||
|
||
if pretrain:
|
||
input_text_a = text_a
|
||
input_text_b = text_b
|
||
else:
|
||
input_text_a = tokenizer.sep_token.join([text_a, text_b])
|
||
input_text_b = text_c
|
||
|
||
|
||
inputs = tokenizer(
|
||
input_text_a,
|
||
input_text_b,
|
||
truncation="longest_first",
|
||
max_length=max_seq_length,
|
||
padding="longest",
|
||
add_special_tokens=True,
|
||
)
|
||
# assert tokenizer.mask_token_id in inputs.input_ids, "mask token must in input"
|
||
|
||
features = asdict(InputFeatures(input_ids=inputs["input_ids"],
|
||
attention_mask=inputs['attention_mask'],
|
||
labels=torch.tensor(example.label),
|
||
label=torch.tensor(example.real_label)
|
||
)
|
||
)
|
||
return features
|
||
|
||
|
||
def _truncate_seq_pair(tokens_a, tokens_b, max_length):
|
||
"""Truncates a sequence pair in place to the maximum length."""
|
||
|
||
# This is a simple heuristic which will always truncate the longer sequence
|
||
# one token at a time. This makes more sense than truncating an equal percent
|
||
# of tokens from each, since if one sequence is very short then each token
|
||
# that's truncated likely contains more information than a longer sequence.
|
||
while True:
|
||
total_length = len(tokens_a) + len(tokens_b)
|
||
if total_length <= max_length:
|
||
break
|
||
if len(tokens_a) > len(tokens_b):
|
||
tokens_a.pop()
|
||
else:
|
||
tokens_b.pop()
|
||
|
||
def _truncate_seq_triple(tokens_a, tokens_b, tokens_c, max_length):
|
||
"""Truncates a sequence triple in place to the maximum length."""
|
||
|
||
# This is a simple heuristic which will always truncate the longer sequence
|
||
# one token at a time. This makes more sense than truncating an equal percent
|
||
# of tokens from each, since if one sequence is very short then each token
|
||
# that's truncated likely contains more information than a longer sequence.
|
||
while True:
|
||
total_length = len(tokens_a) + len(tokens_b) + len(tokens_c)
|
||
if total_length <= max_length:
|
||
break
|
||
if len(tokens_a) > len(tokens_b) and len(tokens_a) > len(tokens_c):
|
||
tokens_a.pop()
|
||
elif len(tokens_b) > len(tokens_a) and len(tokens_b) > len(tokens_c):
|
||
tokens_b.pop()
|
||
elif len(tokens_c) > len(tokens_a) and len(tokens_c) > len(tokens_b):
|
||
tokens_c.pop()
|
||
else:
|
||
tokens_c.pop()
|
||
|
||
|
||
@cache_results(_cache_fp="./dataset")
|
||
def get_dataset(args, processor, label_list, tokenizer, mode):
|
||
|
||
assert mode in ["train", "dev", "test"], "mode must be in train dev test!"
|
||
|
||
# use training data to construct the entity embedding
|
||
combine_train_and_test = False
|
||
if args.faiss_init and mode == "test" and not args.pretrain:
|
||
mode = "train"
|
||
if "ind" in args.data_dir: combine_train_and_test = True
|
||
else:
|
||
pass
|
||
|
||
if mode == "train":
|
||
train_examples = processor.get_train_examples(args.data_dir)
|
||
elif mode == "dev":
|
||
train_examples = processor.get_dev_examples(args.data_dir)
|
||
else:
|
||
train_examples = processor.get_test_examples(args.data_dir)
|
||
|
||
if combine_train_and_test:
|
||
logger.info("use all the dataset for getting the entity mask embedding in pretraining pretraining")
|
||
logger.info("use all the dataset for getting the entity mask embedding in pretraining pretraining")
|
||
train_examples = processor.get_test_examples(args.data_dir) + processor.get_train_examples(args.data_dir) + processor.get_dev_examples(args.data_dir)
|
||
|
||
from os import cpu_count
|
||
with open(os.path.join(args.data_dir, f"examples_{mode}.txt"), 'w') as file:
|
||
for line in train_examples:
|
||
d = {}
|
||
d.update(line.__dict__)
|
||
file.write(json.dumps(d) + '\n')
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=False)
|
||
|
||
features = []
|
||
# # with open(os.path.join(args.data_dir, "cached_relation_pattern.pkl"), "rb") as file:
|
||
# # pattern = pickle.load(file)
|
||
# pattern = None
|
||
# convert_examples_to_features_init(tokenizer)
|
||
# annotate_ = partial(
|
||
# convert_examples_to_features,
|
||
# max_seq_length=args.max_seq_length,
|
||
# mode = mode,
|
||
# pretrain = args.pretrain
|
||
# )
|
||
# features = list(
|
||
# tqdm(
|
||
# map(annotate_, train_examples),
|
||
# total=len(train_examples)
|
||
# )
|
||
# )
|
||
# encoder = MultiprocessingEncoder(tokenizer, args)
|
||
# encoder.initializer()
|
||
# for t in tqdm(train_examples):
|
||
# features.append(encoder.encode_lines([json.dumps(t.__dict__)]))
|
||
|
||
# for example in tqdm(train_examples):
|
||
# text_a = example.text_a
|
||
# text_b = example.text_b
|
||
# text_c = example.text_c
|
||
|
||
# bpe = tokenizer
|
||
# if 0:
|
||
# input_text_a = text_a
|
||
# input_text_b = text_b
|
||
# else:
|
||
# if text_a == "[MASK]":
|
||
# input_text_a = bpe.sep_token.join([text_a, text_b])
|
||
# input_text_b = text_c
|
||
# else:
|
||
# input_text_a = text_a
|
||
# input_text_b = bpe.sep_token.join([text_b, text_c])
|
||
|
||
|
||
# inputs = tokenizer(
|
||
# input_text_a,
|
||
# # input_text_b,
|
||
# truncation="longest_first",
|
||
# max_length=128,
|
||
# padding="longest",
|
||
# add_special_tokens=True,
|
||
# )
|
||
# # assert tokenizer.mask_token_id in inputs.input_ids, "mask token must in input"
|
||
|
||
# # features.append(asdict(InputFeatures(input_ids=inputs["input_ids"],
|
||
# # attention_mask=inputs['attention_mask'],
|
||
# # labels=example.label,
|
||
# # label=example.real_label
|
||
# # )
|
||
# # ))
|
||
|
||
|
||
file_inputs = [os.path.join(args.data_dir, f"examples_{mode}.txt")]
|
||
file_outputs = [os.path.join(args.data_dir, f"features_{mode}.txt")]
|
||
|
||
with contextlib.ExitStack() as stack:
|
||
inputs = [
|
||
stack.enter_context(open(input, "r", encoding="utf-8"))
|
||
if input != "-" else sys.stdin
|
||
for input in file_inputs
|
||
]
|
||
outputs = [
|
||
stack.enter_context(open(output, "w", encoding="utf-8"))
|
||
if output != "-" else sys.stdout
|
||
for output in file_outputs
|
||
]
|
||
|
||
encoder = MultiprocessingEncoder(tokenizer, args)
|
||
pool = Pool(16, initializer=encoder.initializer)
|
||
encoder.initializer()
|
||
encoded_lines = pool.imap(encoder.encode_lines, zip(*inputs), 1000)
|
||
# encoded_lines = map(encoder.encode_lines, zip(*inputs))
|
||
|
||
stats = Counter()
|
||
for i, (filt, enc_lines) in tqdm(enumerate(encoded_lines, start=1), total=len(train_examples)):
|
||
if filt == "PASS":
|
||
for enc_line, output_h in zip(enc_lines, outputs):
|
||
features.append(eval(enc_line))
|
||
# features.append(enc_line)
|
||
# print(enc_line, file=output_h)
|
||
else:
|
||
stats["num_filtered_" + filt] += 1
|
||
# if i % 10000 == 0:
|
||
# print("processed {} lines".format(i), file=sys.stderr)
|
||
|
||
for k, v in stats.most_common():
|
||
print("[{}] filtered {} lines".format(k, v), file=sys.stderr)
|
||
# threads = min(16, cpu_count())
|
||
# with Pool(threads, initializer=convert_examples_to_features_init, initargs=(tokenizer,)) as pool:
|
||
# annotate_ = partial(
|
||
# convert_examples_to_features,
|
||
# max_seq_length=args.max_seq_length,
|
||
# mode = mode,
|
||
# pretrain = args.pretrain
|
||
# )
|
||
# features = list(
|
||
# tqdm(
|
||
# pool.imap_unordered(annotate_, train_examples),
|
||
# total=len(train_examples),
|
||
# desc="convert examples to features",
|
||
# )
|
||
# )
|
||
|
||
# num_entities = len(processor.get_entities(args.data_dir))
|
||
for f_id, f in enumerate(features):
|
||
en = features[f_id].pop("en")
|
||
rel = features[f_id].pop("rel")
|
||
real_label = f['label']
|
||
cnt = 0
|
||
if not isinstance(en, list): break
|
||
|
||
pos = 0
|
||
for i,t in enumerate(f['input_ids']):
|
||
if t == tokenizer.pad_token_id:
|
||
features[f_id]['input_ids'][i] = en[cnt] + len(tokenizer)
|
||
cnt += 1
|
||
if features[f_id]['input_ids'][i] == real_label + len(tokenizer):
|
||
pos = i
|
||
if cnt == len(en): break
|
||
assert not (args.faiss_init and pos == 0)
|
||
features[f_id]['pos'] = pos
|
||
|
||
|
||
# for i,t in enumerate(f['input_ids']):
|
||
# if t == tokenizer.pad_token_id:
|
||
# features[f_id]['input_ids'][i] = rel + len(tokenizer) + num_entities
|
||
# break
|
||
|
||
|
||
|
||
features = KGCDataset(features)
|
||
return features
|
||
|
||
|
||
class MultiprocessingEncoder(object):
|
||
def __init__(self, tokenizer, args):
|
||
self.tokenizer = tokenizer
|
||
self.pretrain = args.pretrain
|
||
self.max_seq_length = args.max_seq_length
|
||
|
||
def initializer(self):
|
||
global bpe
|
||
bpe = self.tokenizer
|
||
|
||
def encode(self, line):
|
||
global bpe
|
||
ids = bpe.encode(line)
|
||
return list(map(str, ids))
|
||
|
||
def decode(self, tokens):
|
||
global bpe
|
||
return bpe.decode(tokens)
|
||
|
||
def encode_lines(self, lines):
|
||
"""
|
||
Encode a set of lines. All lines will be encoded together.
|
||
"""
|
||
enc_lines = []
|
||
for line in lines:
|
||
line = line.strip()
|
||
if len(line) == 0:
|
||
return ["EMPTY", None]
|
||
# enc_lines.append(" ".join(tokens))
|
||
enc_lines.append(json.dumps(self.convert_examples_to_features(example=eval(line))))
|
||
# enc_lines.append(" ")
|
||
# enc_lines.append("123")
|
||
return ["PASS", enc_lines]
|
||
|
||
def decode_lines(self, lines):
|
||
dec_lines = []
|
||
for line in lines:
|
||
tokens = map(int, line.strip().split())
|
||
dec_lines.append(self.decode(tokens))
|
||
return ["PASS", dec_lines]
|
||
|
||
def convert_examples_to_features(self, example):
|
||
pretrain = self.pretrain
|
||
max_seq_length = self.max_seq_length
|
||
global bpe
|
||
"""Loads a data file into a list of `InputBatch`s."""
|
||
# tokens_a = tokenizer.tokenize(example.text_a)
|
||
# tokens_b = tokenizer.tokenize(example.text_b)
|
||
# tokens_c = tokenizer.tokenize(example.text_c)
|
||
|
||
# _truncate_seq_triple(tokens_a, tokens_b, tokens_c, max_length= max_seq_length)
|
||
# text_a = " ".join(example['text_a'].split()[:128])
|
||
# text_b = " ".join(example['text_b'].split()[:128])
|
||
# text_c = " ".join(example['text_c'].split()[:128])
|
||
|
||
text_a = example['text_a']
|
||
text_b = example['text_b']
|
||
text_c = example['text_c']
|
||
|
||
if pretrain:
|
||
# the des of xxx is [MASK] .
|
||
input_text = f"The description of {text_a} is that {text_b} ."
|
||
inputs = bpe(
|
||
input_text,
|
||
truncation="longest_first",
|
||
max_length=max_seq_length,
|
||
padding="longest",
|
||
add_special_tokens=True,
|
||
)
|
||
else:
|
||
if text_a == "[MASK]":
|
||
input_text_a = bpe.sep_token.join([text_a, text_b])
|
||
input_text_b = text_c
|
||
else:
|
||
input_text_a = text_a
|
||
input_text_b = bpe.sep_token.join([text_b, text_c])
|
||
|
||
|
||
inputs = bpe(
|
||
input_text_a,
|
||
input_text_b,
|
||
truncation="longest_first",
|
||
max_length=max_seq_length,
|
||
padding="longest",
|
||
add_special_tokens=True,
|
||
)
|
||
# assert bpe.mask_token_id in inputs.input_ids, "mask token must in input"
|
||
|
||
features = asdict(InputFeatures(input_ids=inputs["input_ids"],
|
||
attention_mask=inputs['attention_mask'],
|
||
labels=example['label'],
|
||
label=example['real_label'],
|
||
en=example['en'],
|
||
rel=example['rel']
|
||
)
|
||
)
|
||
return features
|
||
if __name__ == "__main__":
|
||
dataset = KGCDataset('./dataset') |