event_extraction.py (6406B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 import os 4 5 from data_convert.task_format.task_format import TaskFormat 6 7 8 class DyIEPP(TaskFormat): 9 def __init__(self, doc_json): 10 self.doc_key = doc_json['doc_key'] 11 self.sentences = doc_json['sentences'] 12 self.ner = doc_json['ner'] 13 self.relations = doc_json['relations'] 14 self.events = doc_json['events'] 15 self.sentence_start = doc_json.get('sentence_start', doc_json['_sentence_start']) 16 17 def generate_sentence(self, type_format='subtype'): 18 for sentence, events_in_sentence, sentence_start in zip(self.sentences, self.events, self.sentence_start): 19 events = list() 20 for event in events_in_sentence: 21 trigger, event_type = event[0] 22 trigger -= sentence_start 23 24 suptype, subtype = event_type.split('.') 25 26 if type_format == 'subtype': 27 event_type = subtype 28 elif type_format == 'suptype': 29 event_type = suptype 30 else: 31 event_type = suptype + type_format + subtype 32 33 arguments = list() 34 for start, end, role in event[1:]: 35 start -= sentence_start 36 end -= sentence_start 37 arguments += [[role, list(range(start, end + 1))]] 38 39 event = {'type': event_type, 'tokens': [trigger], 'arguments': arguments} 40 41 events += [event] 42 yield {'tokens': sentence, 'events': events} 43 44 45 class Event(TaskFormat): 46 """ 47 { 48 "doc_id": "NYT_ENG_20130914.0094", 49 "sent_id": "NYT_ENG_20130914.0094-1", 50 "tokens": ["LARGO", "\u2014", "A", "judge", "on", "Friday", "refused", "to", "stop", "''", "Hiccup", "Girl", "\u2019'", "Jennifer", "Mee", "from", "giving", "interviews", "to", "the", "media", "in", "the", "final", "days", "before", "her", "murder", "trial", "."], 51 "entities": [ 52 {"entity_id": "NYT_ENG_20130914.0094-1-8-42", "entity_type": "PER", "mention_type": "NOM", "start": 3, "end": 4, "text": "judge"}, 53 {"entity_id": "NYT_ENG_20130914.0094-1-39-2096", "entity_type": "PER", "mention_type": "NAM", "start": 10, "end": 12, "text": "Hiccup Girl"}, 54 {"entity_id": "NYT_ENG_20130914.0094-1-39-48", "entity_type": "PER", "mention_type": "NAM", "start": 13, "end": 15, "text": "Jennifer Mee"}, 55 {"entity_id": "NYT_ENG_20130914.0094-1-39-54", "entity_type": "PER", "mention_type": "PRO", "start": 26, "end": 27, "text": "her"} 56 ], 57 "relations": [], 58 "events": [ 59 { 60 "event_id": "NYT_ENG_20130914.0094-1-25-998", "event_type": "contact", "event_subtype": "broadcast", 61 "trigger": {"text": "interviews", "start": 17, "end": 18}, 62 "arguments": [ 63 {"entity_id": "NYT_ENG_20130914.0094-1-39-48", "role": "entity", "text": "Jennifer Mee"} 64 ] 65 }, 66 { 67 "event_id": "NYT_ENG_20130914.0094-1-1040-1019", "event_type": "justice", "event_subtype": "trialhearing", 68 "trigger": {"text": "trial", "start": 28, "end": 29}, 69 "arguments": [{"entity_id": "NYT_ENG_20130914.0094-1-39-54", "role": "defendant", "text": "her"} 70 ] 71 } 72 ], 73 "start": 231, "end": 380, 74 "text": "LARGO \u2014 A judge on Friday refused to stop ''Hiccup Girl\u2019' Jennifer Mee from giving interviews to the media in the final days before her murder trial." 75 } 76 """ 77 78 def __init__(self, doc_json): 79 self.doc_key = doc_json['doc_id'] 80 self.sentence = doc_json['tokens'] 81 self.entities = {entity['id']: entity for entity in doc_json['entity_mentions']} 82 self.relations = doc_json['relation_mentions'] 83 self.events = doc_json['event_mentions'] 84 # self.sentence_start = doc_json['start'] 85 # self.sentence_end = doc_json['end'] 86 # self.text = doc_json['text'] 87 88 def generate_sentence(self, type_format='subtype'): 89 events = list() 90 91 for event in self.events: 92 arguments = list() 93 for argument in event['arguments']: 94 argument_entity = self.entities[argument['entity_id']] 95 arguments += [[argument['role'], list(range(argument_entity['start'], argument_entity['end']))]] 96 97 suptype, subtype = event['event_type'].split(':') 98 99 if type_format == 'subtype': 100 event_type = subtype 101 elif type_format == 'suptype': 102 event_type = suptype 103 else: 104 event_type = suptype + type_format + subtype 105 106 events += [{ 107 'type': event_type, 108 'tokens': list(range(event['trigger']['start'], event['trigger']['end'])), 109 'arguments': arguments 110 }] 111 112 yield {'tokens': self.sentence, 'events': events} 113 114 115 def DyIEPP_ace2005_file_tuple(output_folder): 116 if not os.path.exists(output_folder): 117 os.makedirs(output_folder, exist_ok=True) 118 119 conll_2012_folder = "data/raw_data/dyiepp_ace2005" 120 121 file_tuple = [ 122 (conll_2012_folder + "/train.json", output_folder + '/train'), 123 (conll_2012_folder + "/dev.json", output_folder + '/val'), 124 (conll_2012_folder + "/test.json", output_folder + '/test'), 125 ] 126 127 return file_tuple 128 129 130 def ace2005_en_file_tuple(output_folder): 131 if not os.path.exists(output_folder): 132 os.makedirs(output_folder, exist_ok=True) 133 134 conll_2012_folder = "data/raw_data/ace05-EN/" 135 136 file_tuple = [ 137 (conll_2012_folder + "/train.oneie.json", output_folder + '/train'), 138 (conll_2012_folder + "/dev.oneie.json", output_folder + '/val'), 139 (conll_2012_folder + "/test.oneie.json", output_folder + '/test'), 140 ] 141 142 return file_tuple 143 144 145 def ere_en_file_tuple(output_folder): 146 if not os.path.exists(output_folder): 147 os.makedirs(output_folder, exist_ok=True) 148 149 conll_2012_folder = "data/raw_data/ERE-EN" 150 151 file_tuple = [ 152 (conll_2012_folder + "/train.oneie.json", output_folder + '/train'), 153 (conll_2012_folder + "/dev.oneie.json", output_folder + '/val'), 154 (conll_2012_folder + "/test.oneie.json", output_folder + '/test'), 155 ] 156 157 return file_tuple 158 159 160 if __name__ == "__main__": 161 pass