def __init__(self,
data,
label_map,
tokenizer,
max_length=512,
pad_to_max_length=False):
super(DuIEDataset, self).__init__()
self.data = data
self.chn_punc_extractor = ChineseAndPunctuationExtractor()
self.tokenizer = tokenizer
self.max_seq_length = max_length
self.pad_to_max_length = pad_to_max_length
self.label_map = label_map
def from_file(cls,
file_path: Union[str, os.PathLike],
tokenizer: AutoTokenizer,
max_length: Optional[int] = 512,
pad_to_max_length: Optional[bool] = None):
assert os.path.exists(file_path) and os.path.isfile(
file_path), f"{file_path} dose not exists or is not a file."
label_map_path = os.path.join(os.path.dirname(file_path),
"predicate2id.json")
assert os.path.exists(label_map_path) and os.path.isfile(
label_map_path
), f"{label_map_path} dose not exists or is not a file."
with open(label_map_path, 'r', encoding='utf8') as fp:
label_map = json.load(fp)
chineseandpunctuationextractor = ChineseAndPunctuationExtractor()
input_ids, seq_lens, tok_to_orig_start_index, tok_to_orig_end_index, labels = (
[] for _ in range(5))
dataset_scale = sum(1
for line in open(file_path, 'r', encoding='utf8'))
print("Preprocessing data, loaded from %s" % file_path)
with open(file_path, "r", encoding="utf-8") as fp:
lines = fp.readlines()
for line in tqdm(lines):
example = json.loads(line)
input_feature = convert_example_to_feature(
example, tokenizer, chineseandpunctuationextractor,
label_map, max_length, pad_to_max_length)
input_ids.append(input_feature.input_ids)
seq_lens.append(input_feature.seq_len)
tok_to_orig_start_index.append(
input_feature.tok_to_orig_start_index)
tok_to_orig_end_index.append(
input_feature.tok_to_orig_end_index)
labels.append(input_feature.labels)
return cls(input_ids, seq_lens, tok_to_orig_start_index,
tok_to_orig_end_index, labels)
def convert_example_to_feature(
example,
tokenizer: AutoTokenizer,
chineseandpunctuationextractor: ChineseAndPunctuationExtractor,
label_map,
max_length: Optional[int] = 512,
pad_to_max_length: Optional[bool] = None):
spo_list = example['spo_list'] if "spo_list" in example.keys() else None
text_raw = example['text']
sub_text = [] # 放置中文字符
buff = "" # 存放非中文字符
for char in text_raw:
if chineseandpunctuationextractor.is_chinese_or_punct(char):
if buff != "":
sub_text.append(buff)
buff = ""
sub_text.append(char)
else:
buff += char
if buff != "":
sub_text.append(buff)
tok_to_orig_start_index = []
tok_to_orig_end_index = []
orig_to_tok_index = []
tokens = []
text_tmp = ''
for (i, token) in enumerate(sub_text):
orig_to_tok_index.append(len(tokens))
sub_tokens = tokenizer.tokenize(token)
text_tmp += token
for sub_token in sub_tokens:
tok_to_orig_start_index.append(len(text_tmp) - len(token))
tok_to_orig_end_index.append(len(text_tmp) - 1)
tokens.append(sub_token)
if len(tokens) >= max_length - 2:
break
else:
continue
break
# print("tok_to_orig_start_index: ", tok_to_orig_start_index)
# print("tok_to_orig_end_index: ", tok_to_orig_end_index)
# print("orig_to_tok_index: ", orig_to_tok_index)
# print("tokens: ", tokens)
seq_len = len(tokens)
# 2 tags for each predicate + I tag + O tag
num_labels = 2 * (len(label_map.keys()) - 2) + 2
# initialize tag
labels = [[0] * num_labels for i in range(seq_len)] # 每个字都要生成标签表示，用于预测
if spo_list is not None:
labels = parse_label(spo_list, label_map, tokens, tokenizer)
# add [CLS] and [SEP] token, they are tagged into "O" for outside
if seq_len > max_length - 2:
tokens = tokens[0:(max_length - 2)]
labels = labels[0:(max_length - 2)]
tok_to_orig_start_index = tok_to_orig_start_index[0:(max_length - 2)]
tok_to_orig_end_index = tok_to_orig_end_index[0:(max_length - 2)]
tokens = ["[CLS]"] + tokens + ["[SEP]"]
# "O" tag for [PAD], [CLS], [SEP] token
outside_label = [[1] + [0] * (num_labels - 1)]
labels = outside_label + labels + outside_label
tok_to_orig_start_index = [-1] + tok_to_orig_start_index + [-1]
tok_to_orig_end_index = [-1] + tok_to_orig_end_index + [-1]
if seq_len < max_length:
tokens = tokens + ["[PAD]"] * (max_length - seq_len - 2)
labels = labels + outside_label * (max_length - len(labels))
tok_to_orig_start_index = tok_to_orig_start_index + [-1] * (
max_length - len(tok_to_orig_start_index))
tok_to_orig_end_index = tok_to_orig_end_index + [-1] * (
max_length - len(tok_to_orig_end_index))
token_ids = tokenizer.convert_tokens_to_ids(tokens)
return InputFeature(
input_ids=np.array(token_ids),
seq_len=np.array(seq_len),
tok_to_orig_start_index=np.array(tok_to_orig_start_index),
tok_to_orig_end_index=np.array(tok_to_orig_end_index),
labels=np.array(labels),
)