| import json |
| from collections import defaultdict |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def split_span(span): |
| sub_spans = [[span[0]]] |
| for i in range(1, len(span)): |
| if span[i - 1] == span[i] - 1: |
| sub_spans[-1].append(span[i]) |
| else: |
| sub_spans.append([span[i]]) |
| return sub_spans |
|
|
|
|
| class OIE4ReaderForJointDecoding(): |
| """Define text data reader and preprocess data for entity relation |
| joint decoding on ACE dataset. |
| """ |
|
|
| def __init__(self, file_path, is_test=False, max_len=dict()): |
| """This function defines file path and some settings |
| |
| Arguments: |
| file_path {str} -- file path |
| |
| Keyword Arguments: |
| is_test {bool} -- indicate training or testing (default: {False}) |
| max_len {dict} -- max length for some namespace (default: {dict()}) |
| """ |
|
|
| self.file_path = file_path |
| self.is_test = is_test |
| self.max_len = dict(max_len) |
| self.seq_lens = defaultdict(list) |
|
|
| def __iter__(self): |
| """Generator function |
| """ |
|
|
| with open(self.file_path, 'r') as fin: |
| for line in fin: |
| line = json.loads(line) |
| sentence = {} |
|
|
| state, results = self.get_tokens(line) |
| self.seq_lens['tokens'].append(len(results['tokens'])) |
| if not state or ('tokens' in self.max_len and len(results['tokens']) > self.max_len['tokens'] |
| and not self.is_test): |
| if not self.is_test: |
| continue |
| sentence.update(results) |
|
|
| state, results = self.get_wordpiece_tokens(line) |
| self.seq_lens['wordpiece_tokens'].append(len(results['wordpiece_tokens'])) |
| if not state or ('wordpiece_tokens' in self.max_len |
| and len(results['wordpiece_tokens']) > self.max_len['wordpiece_tokens']): |
| if not self.is_test: |
| continue |
| sentence.update(results) |
|
|
| if len(sentence['tokens']) != len(sentence['wordpiece_tokens_index']): |
| logger.error( |
| "sentence id: {} wordpiece_tokens_index length is not equal to tokens.".format(line['sentId'])) |
| continue |
|
|
| if len(sentence['wordpiece_tokens']) != len(sentence['wordpiece_segment_ids']): |
| logger.error( |
| "sentence id: {} wordpiece_tokens length is not equal to wordpiece_segment_ids.". |
| format(line['sentId'])) |
| continue |
|
|
| state, results = self.get_entity_relation_label(line, len(sentence['tokens'])) |
| for key, result in results.items(): |
| self.seq_lens[key].append(len(result)) |
| if key in self.max_len and len(result) > self.max_len[key]: |
| state = False |
| if not state: |
| continue |
| sentence.update(results) |
|
|
| yield sentence |
|
|
| def get_tokens(self, line): |
| """This function splits text into tokens |
| |
| Arguments: |
| line {dict} -- text |
| |
| Returns: |
| bool -- execute state |
| dict -- results: tokens |
| """ |
|
|
| results = {} |
|
|
| if 'sentText' not in line: |
| logger.error("sentence id: {} doesn't contain 'sentText'.".format( |
| line['sentId'])) |
| return False, results |
|
|
| results['text'] = line['sentText'] |
|
|
| if 'tokens' in line: |
| results['tokens'] = line['tokens'] |
| else: |
| results['tokens'] = line['sentText'].strip().split(' ') |
|
|
| return True, results |
|
|
| def get_wordpiece_tokens(self, line): |
| """This function splits wordpiece text into wordpiece tokens |
| |
| Arguments: |
| line {dict} -- text |
| |
| Returns: |
| bool -- execute state |
| dict -- results: tokens |
| """ |
|
|
| results = {} |
|
|
| if 'wordpieceSentText' not in line or 'wordpieceTokensIndex' not in line or 'wordpieceSegmentIds' not in line: |
| logger.error( |
| "sentence id: {} doesn't contain 'wordpieceSentText' or 'wordpieceTokensIndex' or 'wordpieceSegmentIds'." |
| .format(line['sentId'])) |
| return False, results |
|
|
| wordpiece_tokens = line['wordpieceSentText'].strip().split(' ') |
| results['wordpiece_tokens'] = wordpiece_tokens |
| results['wordpiece_tokens_index'] = [span[0] for span in line['wordpieceTokensIndex']] |
| results['wordpiece_segment_ids'] = list(line['wordpieceSegmentIds']) |
|
|
| return True, results |
|
|
| def get_entity_relation_label(self, line, sentence_length): |
| """This function constructs mapping relation from span to entity label |
| and span pair to relation label, and joint entity relation label matrix. |
| |
| Arguments: |
| line {dict} -- text |
| sentence_length {int} -- sentence length |
| |
| Returns: |
| bool -- execute state |
| dict -- ent2rel: entity span mapping to entity label, |
| span2rel: two entity span mapping to relation label, |
| joint_label_matrix: joint entity relation label matrix |
| """ |
|
|
| results = {} |
|
|
| if 'entityMentions' not in line: |
| logger.error("sentence id: {} doesn't contain 'entityMentions'.".format(line['sentId'])) |
| return False, results |
|
|
| entity_pos = [0] * sentence_length |
| idx2ent = {} |
| span2ent = {} |
|
|
| separate_positions = [] |
| for entity in line['entityMentions']: |
| entity_sub_spans = [] |
| |
| st, ed = entity['span_ids'][0], entity['span_ids'][-1] |
| sub_spans = split_span(entity['span_ids']) |
| if len(sub_spans) == 1: |
| if st > 0: |
| separate_positions.append(st - 1) |
| if ed < sentence_length - 1: |
| separate_positions.append(ed) |
| entity_sub_spans.append((st, ed + 1)) |
| else: |
| |
| for sub in sub_spans: |
| if sub[0] > 0: |
| separate_positions.append(sub[0] - 1) |
| if sub[-1] < sentence_length - 1: |
| separate_positions.append(sub[-1]) |
| entity_sub_spans.append((sub[0], sub[-1] + 1)) |
|
|
| idx2ent[entity['emId']] = (tuple(entity_sub_spans), entity['text']) |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| span2ent[tuple(entity_sub_spans)] = entity['label'] |
|
|
| j = 0 |
| for s_i in entity['span_ids']: |
| if entity_pos[s_i] != 0: |
| logger.error("sentence id: {} entity span overlap. {}".format(line['sentId'], entity['span_ids'])) |
| return False, results |
| entity_pos[s_i] = 1 |
| j += 1 |
|
|
| separate_positions = list(set(separate_positions)) |
| results['separate_positions'] = sorted(separate_positions) |
| results['span2ent'] = span2ent |
|
|
| if 'relationMentions' not in line: |
| logger.error("sentence id: {} doesn't contain 'relationMentions'.".format(line['sentId'])) |
| return False, results |
|
|
| span2rel = {} |
| for relation in line['relationMentions']: |
| if relation['arg1']['emId'] not in idx2ent or relation['arg2']['emId'] not in idx2ent: |
| logger.error("sentence id: {} entity not exists .".format(line['sentId'])) |
| continue |
|
|
| entity1_span, entity1_text = idx2ent[relation['arg1']['emId']] |
| entity2_span, entity2_text = idx2ent[relation['arg2']['emId']] |
|
|
| if entity1_text != relation['arg1']['text'] or entity2_text != relation['arg2']['text']: |
| logger.error("sentence id: {} entity text doesn't match realtiaon text.".format(line['sentId'])) |
| return False, None |
|
|
| span2rel[(entity1_span, entity2_span)] = relation['label'] |
|
|
| results['span2rel'] = span2rel |
|
|
| if 'jointLabelMatrix' not in line: |
| logger.error("sentence id: {} doesn't contain 'jointLabelMatrix'.".format(line['sentId'])) |
| return False, results |
|
|
| results['joint_label_matrix'] = line['jointLabelMatrix'] |
| return True, results |
|
|
| def get_seq_lens(self): |
| return self.seq_lens |
|
|
