| | import json |
| | from collections import defaultdict |
| | import logging |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | class ReaderForRelationDecoding(): |
| | """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['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_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_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 = {} |
| | for entity in line['entityMentions']: |
| | st, ed = entity['span_ids'][0], entity['span_ids'][-1] |
| | idx2ent[entity['emId']] = ((st, ed), entity['text']) |
| | if st > ed or st < 0 or st > sentence_length or ed < 0 or ed > sentence_length: |
| | logger.error("sentence id: {} offset error. st: {}, ed: {}".format(line['sentId'], st, ed)) |
| | return False, results |
| |
|
| | j = 0 |
| | for i in range(st, ed): |
| | if entity_pos[i] != 0: |
| | logger.error("sentence id: {} entity span overlap.".format(line['sentId'])) |
| | return False, results |
| | entity_pos[i] = 1 |
| | j += 1 |
| |
|
| | if 'relationMentions' not in line: |
| | logger.error("sentence id: {} doesn't contain 'relationMentions'.".format(line['sentId'])) |
| | return False, results |
| |
|
| | 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 |
| |
|
| | if 'labelIds' not in line: |
| | logger.error("sentence id: {} doesn't contain 'labelIds'.".format(line['sentId'])) |
| | return False, results |
| |
|
| | if 'relationIds' not in line: |
| | logger.error("sentence id: {} doesn't contain 'relationIds'.".format(line['sentId'])) |
| | return False, results |
| |
|
| | if 'argumentIds' not in line: |
| | logger.error("sentence id: {} doesn't contain 'argumentIds'.".format(line['sentId'])) |
| | return False, results |
| |
|
| | results["label_ids"] = line["labelIds"] |
| | results["relation_ids"] = line["relationIds"] |
| | results["argument_ids"] = line["argumentIds"] |
| |
|
| | return True, results |
| |
|
| | def get_seq_lens(self): |
| | return self.seq_lens |
| |
|