process_ere.py (65605B)
1 import os 2 import re 3 import json 4 import glob 5 import random 6 from lxml import etree 7 from bs4 import BeautifulSoup 8 from dataclasses import dataclass 9 from collections import defaultdict 10 from typing import List, Dict, Any, Tuple 11 from nltk import sent_tokenize, wordpunct_tokenize 12 from transformers import (BertTokenizer, RobertaTokenizer, 13 XLMRobertaTokenizer, PreTrainedTokenizer, 14 AutoTokenizer) 15 from argparse import ArgumentParser, Namespace 16 17 18 ERE_V1 = 'LDC2015E29_DEFT_Rich_ERE_English_Training_Annotation_V1' 19 ERE_V2 = 'LDC2015E29_DEFT_Rich_ERE_English_Training_Annotation_V2' 20 ERE_R2_V2 = 'LDC2015E68_DEFT_Rich_ERE_English_Training_Annotation_R2_V2' 21 ERE_PARL_V2 = 'LDC2015E78_DEFT_Rich_ERE_Chinese_and_English_Parallel_Annotation_V2' 22 SPANISH = 'LDC2015E107_DEFT_Rich_ERE_Spanish_Annotation_V2' 23 24 WRAPPED_DOCS = {'XIN_ENG_20101125.0137'} 25 DOCS_TO_REVISE_SENT = {'NYT_ENG_20131225.0200', 26 'NYT_ENG_20130716.0217', 27 '4fbb1eec7dfd5c2fefb94a2d873ddfa5', 28 '1f60eb9697e240af089b134b69c2042d', 29 'XIN_SPA_20050402.0105', 30 'APW_SPA_19980914.0097'} 31 32 EVENT_TYPE_WEIGHT = { 33 'conflict': 5, 34 'justice': 4, 35 'life': 3, 36 } 37 38 relation_type_mapping = { 39 'orgaffiliation': 'ORG-AFF', 40 'personalsocial': 'PER-SOC', 41 'physical': 'PHYS', 42 'generalaffiliation': 'GEN-AFF', 43 'partwhole': 'PART-WHOLE', 44 } 45 relation_subtype_mapping = { 46 'opra': 'OPRA', 47 'subsidiary': 'Subsidiary', 48 'business': 'Business', 49 'membership': 'Membership', 50 'founder': 'Founder', 51 'studentalum': 'StudentAlum', 52 'employmentmembership': 'EmploymentMembership', 53 'family': 'Family', 54 'unspecified': 'Unspecified', 55 'ownership': 'Ownership', 56 'investorshareholder': 'InvestorShareholder', 57 'orgheadquarter': 'OrgHeadquarter', 58 'locatednear': 'LocatedNear', 59 'more': 'MORE', 60 'resident': 'Resident', 61 'orglocationorigin': 'OrgLocationOrigin', 62 'leadership': 'Leadership' 63 } 64 65 event_type_mapping = { 66 'business:declarebankruptcy': 'Business:Declare-Bankruptcy', 67 'business:endorg': 'Business:End-Org', 68 'business:mergeorg': 'Business:Merge-Org', 69 'business:startorg': 'Business:Start-Org', 70 'conflict:attack': 'Conflict:Attack', 71 'conflict:demonstrate': 'Conflict:Demonstrate', 72 'contact:broadcast': 'Contact:Broadcast', 73 'contact:contact': 'Contact:Contact', 74 'contact:correspondence': 'Contact:Correspondence', 75 'contact:meet': 'Contact:Meet', 76 'justice:acquit': 'Justice:Acquit', 77 'justice:appeal': 'Justice:Appeal', 78 'justice:arrestjail': 'Justice:Arrest-Jail', 79 'justice:chargeindict': 'Justice:Charge-Indict', 80 'justice:convict': 'Justice:Convict', 81 'justice:execute': 'Justice:Execute', 82 'justice:extradite': 'Justice:Extradite', 83 'justice:fine': 'Justice:Fine', 84 'justice:pardon': 'Justice:Pardon', 85 'justice:releaseparole': 'Justice:Release-Parole', 86 'justice:sentence': 'Justice:Sentence', 87 'justice:sue': 'Justice:Sue', 88 'justice:trialhearing': 'Justice:Trial-Hearing', 89 'life:beborn': 'Life:Be-Born', 90 'life:die': 'Life:Die', 91 'life:divorce': 'Life:Divorce', 92 'life:injure': 'Life:Injure', 93 'life:marry': 'Life:Marry', 94 'manufacture:artifact': 'Manufacture:Artifact', 95 'movement:transportartifact': 'Movement:Transport-Artifact', 96 'movement:transportperson': 'Movement:Transport-Person', 97 'personnel:elect': 'Personnel:Elect', 98 'personnel:endposition': 'Personnel:End-Position', 99 'personnel:nominate': 'Personnel:Nominate', 100 'personnel:startposition': 'Personnel:Start-Position', 101 'transaction:transaction': 'Transaction:Transaction', 102 'transaction:transfermoney': 'Transaction:Transfer-Money', 103 'transaction:transferownership': 'Transaction:Transfer-Ownership', 104 } 105 106 role_type_mapping = { 107 'victim': 'Victim', 108 'attacker': 'Attacker', 109 'person': 'Person', 110 'plaintiff': 'Plaintiff', 111 'audience': 'Audience', 112 'destination': 'Destination', 113 'prosecutor': 'Prosecutor', 114 'target': 'Target', 115 'origin': 'Origin', 116 'recipient': 'Recipient', 117 'beneficiary': 'Beneficiary', 118 'adjudicator': 'Adjudicator', 119 'thing': 'Thing', 120 'giver': 'Giver', 121 'defendant': 'Defendant', 122 'entity': 'Entity', 123 'org': 'Org', 124 'agent': 'Agent', 125 'place': 'Place', 126 'artifact': 'Artifact', 127 'instrument': 'Instrument' 128 } 129 130 131 def mask_escape(text: str) -> str: 132 """"Replaces escaped characters with rare sequences. 133 134 Args: 135 text (str): text to mask. 136 137 Returns: 138 str: masked string. 139 """ 140 return text.replace('&', 'ҪҪҪҪҪ').replace('<', 'ҚҚҚҚ').replace('>', 'ҺҺҺҺ') 141 142 143 def unmask_escape(text: str) -> str: 144 """Replaces masking sequences with the original escaped characters. 145 146 Args: 147 text (str): masked string. 148 149 Returns: 150 str: unmasked string. 151 """ 152 return text.replace('ҪҪҪҪҪ', '&').replace('ҚҚҚҚ', '<').replace('ҺҺҺҺ', '>') 153 154 155 def recover_escape(text: str) -> str: 156 """Converts named character references in the given string to the corresponding 157 Unicode characters. I didn't notice any numeric character references in this 158 dataset. 159 160 Args: 161 text (str): text to unescape. 162 163 Returns: 164 str: unescaped string. 165 """ 166 return text.replace('&', '&').replace('<', '<').replace('>', '>') 167 168 169 def sentence_tokenize(sentence: Tuple[int, int, str], 170 language: str = 'english') -> List[Tuple[int, int, str]]: 171 """Split a sentence and adds offsets. 172 173 Args: 174 sentence (Tuple[int, int, str]): a sentence tuple consisting of start_offset, 175 end_offset, and sentence text. 176 language (str, optional): sentence language. 177 178 Returns: 179 List[Tuple[int, int, str]]: a list of sentence tuples. 180 """ 181 start, end, text = sentence 182 sents = sent_tokenize(text, language='english') 183 184 last = 0 185 sents_ = [] 186 for sent in sents: 187 index = text[last:].find(sent) 188 if index == -1: 189 print(text, sent) 190 else: 191 sents_.append((last + index + start, last + index + len(sent) + start, sent)) 192 last += index + len(sent) 193 return(sents_) 194 195 196 @dataclass 197 class Span: 198 start: int 199 end: int 200 text: str 201 202 def __post_init__(self): 203 self.start = int(self.start) 204 self.end = int(self.end) 205 self.text = self.text.replace('\n', ' ') 206 207 def char_offsets_to_token_offsets(self, tokens: List[Tuple[int, int, str]]): 208 """Converts self.start and self.end from character offsets to token 209 offsets. 210 211 Args: 212 tokens (List[Tuple[int, int, str]]): a list of token tuples. Each 213 item in the list is a triple (start_offset, end_offset, text). 214 """ 215 start_ = end_ = -1 216 for i, (s, e, _) in enumerate(tokens): 217 if s == self.start: 218 start_ = i 219 if e == self.end: 220 end_ = i + 1 221 if start_ == -1 or end_ == -1 or start_ > end_: 222 raise ValueError('Failed to update offsets for {}-{}:{} in {}'.format( 223 self.start, self.end, self.text, tokens)) 224 self.start, self.end = start_, end_ 225 226 def to_dict(self) -> Dict[str, Any]: 227 """Converts instance variables to a dict. 228 229 Returns: 230 Dict[str, Any]: a dict of instance variables. 231 """ 232 return { 233 'text': recover_escape(self.text), 234 'start': self.start, 235 'end': self.end 236 } 237 238 def remove_space(self): 239 """Removes heading and trailing spaces in the span text.""" 240 # heading spaces 241 text = self.text.lstrip(' ') 242 self.start += len(self.text) - len(text) 243 # trailing spaces 244 after_text = text.rstrip(' ') 245 #self.end = self.start + len(text) 246 #self.text = text 247 self.end += len(after_text) - len(text) 248 self.text = after_text 249 250 def copy(self): 251 """Makes a copy of itself. 252 253 Returns: 254 Span: a copy of itself.""" 255 return Span(self.start, self.end, self.text) 256 257 @dataclass 258 class Entity(Span): 259 entity_id: str 260 mention_id: str 261 entity_type: str 262 mention_type: str 263 264 def copy(self): 265 return Entity(self.start, self.end, self.text, 266 self.entity_id, self.mention_id, self.entity_type, 267 self.mention_type) 268 269 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 270 """Converts instance variables to a dict. 271 272 Args: 273 sent_id (str): Sentence ID. 274 275 Returns: 276 Dict[str, Any]: a dict of instance variables. 277 """ 278 if sent_id: 279 entity_id = '{}-{}-{}'.format(sent_id, 280 self.entity_id.split('-')[-1], 281 self.mention_id.split('-')[-1]) 282 else: 283 entity_id = '{}-{}'.format(self.entity_id.split('-')[-1], 284 self.mention_id.split('-')[-1]) 285 return { 286 'entity_id': entity_id, 287 'entity_type': self.entity_type, 288 'mention_type': self.mention_type, 289 'start': self.start, 290 'end': self.end, 291 'text': recover_escape(self.text) 292 } 293 294 @dataclass 295 class Entity_whole: 296 entity_id: str 297 entity_type: str 298 entity_mentions: List[Entity] 299 300 def __str__(self, sent_id: str = None): 301 return self.to_dict(sent_id) 302 303 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 304 """Converts instance variables to a dict. 305 306 Returns: 307 Dict: a dict of instance variables. 308 """ 309 if sent_id: 310 entity_id = '{}-{}'.format(sent_id, 311 self.entity_id.split('-')[-1]) 312 else: 313 entity_id = '{}'.format(self.entity_id.split('-')[-1]) 314 entity_dict = { 315 'entity_id': entity_id, 316 'entity_type': self.entity_type, 317 'entity_mentions': [e.to_dict(sent_id) for e in self.entity_mentions] 318 } 319 return entity_dict 320 321 def copy(self): 322 """Makes a copy of itself. 323 """ 324 return Entity_whole(self.entity_id, self.entity_type, self.entity_mentions) 325 326 @dataclass 327 class RelationArgument: 328 entity_id: str 329 mention_id: str 330 role: str 331 text: str 332 333 def to_dict(self, sent_id:str = None) -> Dict[str, Any]: 334 """Converts instance variables to a dict. 335 336 Returns: 337 Dict[str, Any]: a dict of instance variables. 338 """ 339 if sent_id: 340 entity_id = '{}-{}-{}'.format(sent_id, 341 self.entity_id.split('-')[-1], 342 self.mention_id.split('-')[-1]) 343 else: 344 entity_id = '{}-{}'.format(self.entity_id.split('-')[-1], 345 self.mention_id.split('-')[-1]) 346 347 return { 348 'entity_id': entity_id, 349 'role': self.role, 350 'text': recover_escape(self.text) 351 } 352 353 354 @dataclass 355 class Relation: 356 relation_id: str 357 mention_id: str 358 relation_type: str 359 relation_subtype: str 360 arg1: RelationArgument 361 arg2: RelationArgument 362 363 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 364 """Converts instance variables to a dict. 365 366 Returns: 367 Dict[str, Any]: a dict of instance variables. 368 """ 369 if sent_id: 370 relation_id = '{}-{}-{}'.format(sent_id, 371 self.relation_id.split('-')[-1], 372 self.mention_id.split('-')[-1]) 373 else: 374 relation_id = '{}-{}'.format(self.relation_id.split('-')[-1], 375 self.mention_id.split('-')[-1]) 376 377 return { 378 'relation_id': relation_id, 379 'relation_type': self.relation_type, 380 'relation_subtype': self.relation_subtype, 381 'arg1': self.arg1.to_dict(sent_id), 382 'arg2': self.arg2.to_dict(sent_id) 383 } 384 385 386 @dataclass 387 class EventArgument: 388 entity_id: str 389 mention_id: str 390 role: str 391 text: str 392 393 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 394 """Converts instance variables to a dict. 395 396 Returns: 397 Dict[str, Any]: a dict of instance variables. 398 """ 399 if sent_id: 400 entity_id = '{}-{}-{}'.format(sent_id, 401 self.entity_id.split('-')[-1], 402 self.mention_id.split('-')[-1]) 403 else: 404 entity_id = '{}-{}'.format(self.entity_id.split('-')[-1], 405 self.mention_id.split('-')[-1]) 406 407 return { 408 'entity_id': entity_id, 409 'role': self.role, 410 'text': recover_escape(self.text), 411 } 412 413 414 @dataclass 415 class Event: 416 event_id: str 417 mention_id: str 418 event_type: str 419 event_subtype: str 420 trigger: Span 421 arguments: List[EventArgument] 422 423 def copy(self): 424 return Event(self.event_id, self.mention_id, self.event_type, 425 self.event_subtype, self.trigger.copy(), self.arguments) 426 427 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 428 """Converts instance variables to a dict. 429 430 Returns: 431 Dict[str, Any]: a dict of instance variables. 432 """ 433 if sent_id: 434 event_id = '{}-{}-{}'.format(sent_id, 435 self.event_id.split('-')[-1], 436 self.mention_id.split('-')[-1]) 437 else: 438 event_id = '{}-{}'.format(self.event_id.split('-')[-1], 439 self.mention_id.split('-')[-1]) 440 return { 441 'event_id': event_id, 442 'event_type': self.event_type, 443 'event_subtype': self.event_subtype, 444 'trigger': self.trigger.to_dict(), 445 'arguments': [arg.to_dict(sent_id) for arg in self.arguments] 446 } 447 448 @dataclass 449 class EventArgument_whole: 450 entity_id: str 451 role: str 452 453 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 454 """Converts instance variables to a dict. 455 456 Returns: 457 Dict[str, Any]: a dict of instance variables. 458 """ 459 if sent_id: 460 entity_id = '{}-{}'.format(sent_id, self.entity_id.split('-')[-1]) 461 else: 462 entity_id = '{}'.format(self.entity_id.split('-')[-1]) 463 return { 464 'entity_id': entity_id, 465 'role': self.role, 466 } 467 468 @dataclass 469 class Event_whole: 470 event_id: str 471 event_type: str 472 event_subtype: str 473 arguments: List[EventArgument_whole] 474 event_mentions: List[Event] 475 476 def to_dict(self, sent_id: str = None) -> Dict[str, Any]: 477 """Converts instance variables to a dict. 478 479 Returns: 480 Dict[str, Any]: a dict of instance variables. 481 """ 482 if sent_id: 483 event_id = '{}-{}'.format(sent_id, 484 self.event_id.split('-')[-1]) 485 else: 486 event_id = '{}'.format(self.event_id.split('-')[-1]) 487 return { 488 'event_id': event_id, 489 'event_type': self.event_type, 490 'event_subtype': self.event_subtype, 491 'arguments': [arg.to_dict(sent_id) for arg in self.arguments], 492 'event_mentions': [e.to_dict(sent_id) for e in self.event_mentions] 493 } 494 495 496 497 @dataclass 498 class Sentence(Span): 499 doc_id: str 500 sent_id: str 501 tokens: List[str] 502 entities: List[Entity] 503 relations: List[Relation] 504 events: List[Event] 505 sent_starts: List[int] 506 entity_cluster: List[Entity_whole] 507 event_cluster: List[Event_whole] 508 509 def to_dict(self) -> Dict[str, Any]: 510 """Converts instance variables to a dict. 511 512 Returns: 513 Dict[str, Any]: a dict of instance variables. 514 """ 515 return { 516 'doc_id': self.doc_id, 517 'wnd_id': self.sent_id, 518 'tokens': [recover_escape(t) for t in self.tokens], 519 'entities': [entity.to_dict(self.sent_id) for entity in self.entities], 520 'relations': [relation.to_dict(self.sent_id) for relation in self.relations], 521 'events': [event.to_dict(self.sent_id) for event in self.events], 522 'sentence_starts': self.sent_starts, 523 'entity_cluster': [entity.to_dict(self.sent_id) for entity in self.entity_cluster], 524 'event_cluster': [event.to_dict(self.sent_id) for event in self.event_cluster], 525 'start': self.start, 526 'end': self.end, 527 'text': recover_escape(self.text).replace('\t', ' ') 528 } 529 530 531 @dataclass 532 class Document: 533 doc_id: str 534 sentences: List[Sentence] 535 536 537 def process_wrapped_text(text: str) -> str: 538 """Handles wrapped text in some documents by replacing linebreaks between a 539 a pair of <p> and </p> tags with spaces. 540 541 Args: 542 text (str): text string to process. 543 544 Returns: 545 str: processed text string. 546 """ 547 segments = text.split('\n') 548 segments_new = [] 549 in_p = False 550 for segment in segments: 551 if in_p: 552 if segment == '</P>': 553 segments_new.append(segment) 554 in_p = False 555 else: 556 if segments_new[-1]: 557 segments_new[-1] += ' ' 558 segments_new[-1] += segment 559 else: 560 segments_new.append(segment) 561 if segment == '<P>': 562 segments_new.append('') 563 in_p = True 564 return '\n'.join(segments_new) 565 566 567 def revise_sentences(sentences: List[Tuple[int, int, str]], 568 doc_id: str) -> List[Tuple[int, int, str]]: 569 """Automatic sentence tokenization may have errors for a few documents. This 570 function joins sentences that are mistakenly split into multiple segments. 571 572 Args: 573 sentences (List[Tuple[int, int, str]]): a list of sentence tuples. 574 doc_id (str): document ID. 575 576 Returns: 577 List[Tuple[int, int, str]]: a list of revised sentence tuples. 578 """ 579 sentences_ = [] 580 if doc_id == 'NYT_ENG_20131225.0200': 581 for sentence_idx, (start, end, text) in enumerate(sentences): 582 if start == 2089: 583 next_start, next_end, next_text = sentences[sentence_idx + 1] 584 space = ' ' * (next_start - end) 585 sentences_.append((start, next_end, text + space + next_text)) 586 elif start == 2163: 587 continue 588 else: 589 sentences_.append((start, end, text)) 590 elif doc_id == 'NYT_ENG_20130716.0217': 591 for sentence_idx, (start, end, text) in enumerate(sentences): 592 if start == 1781: 593 next_start, next_end, next_text = sentences[sentence_idx + 1] 594 space = ' ' * (next_start - end) 595 sentences_.append((start, next_end, text + space + next_text)) 596 elif start == 1932: 597 continue 598 else: 599 sentences_.append((start, end, text)) 600 elif doc_id == '4fbb1eec7dfd5c2fefb94a2d873ddfa5': 601 for sentence_idx, (start, end, text) in enumerate(sentences): 602 if start == 4540: 603 next_start, next_end, next_text = sentences[sentence_idx + 1] 604 space = ' ' * (next_start - end) 605 sentences_.append((start, next_end, text + space + next_text)) 606 elif start == 4543: 607 continue 608 else: 609 sentences_.append((start, end, text)) 610 elif doc_id == '1f60eb9697e240af089b134b69c2042d': 611 for sentence_idx, (start, end, text) in enumerate(sentences): 612 if start == 5373: 613 next_start, next_end, next_text = sentences[sentence_idx + 1] 614 space = ' ' * (next_start - end) 615 sentences_.append((start, next_end, text + space + next_text)) 616 elif start == 5391: 617 continue 618 else: 619 sentences_.append((start, end, text)) 620 elif doc_id == 'XIN_SPA_20050402.0105': 621 for sentence_idx, (start, end, text) in enumerate(sentences): 622 if start == 58: 623 next_start, next_end, next_text = sentences[sentence_idx + 1] 624 space = ' ' * (next_start - end) 625 sentences_.append((start, next_end, text + space + next_text)) 626 elif start == 113: 627 continue 628 else: 629 sentences_.append((start, end, text)) 630 elif doc_id == 'APW_SPA_19980914.0097': 631 cur_start, cur_end, cur_text = 0, 0, '' 632 for start, end, text in sentences: 633 if '--' in text: 634 sentences_.append((cur_start, cur_end, cur_text)) 635 sentences_.append((start, end, text)) 636 cur_start, cur_end, cur_text = 0, 0, '' 637 else: 638 if cur_end == 0: 639 cur_start, cur_end, cur_text = start, end, text 640 else: 641 cur_text += ' ' * (start - cur_end) + text 642 cur_end = end 643 if cur_text: 644 sentences_.append((cur_start, cur_end, cur_text)) 645 646 return sentences_ 647 648 649 def read_source_file(path: str, 650 language: str = 'english', 651 wrapped: bool = False, 652 ) -> List[Tuple[int, int, str]]: 653 """Reads text file. 654 655 Args: 656 path (str): path to the text file. 657 language (str, optional): document language. Defaults to 'english'. 658 wrapped (bool, optional): if the document is wrapped. Defaults to False. 659 660 Returns: 661 List[Tuple[int, int, str]]: a list of sentence tuples. 662 """ 663 data = open(path, 'r', encoding='utf-8').read() 664 665 if wrapped: 666 data = process_wrapped_text(data) 667 # data = data.replace('\n<a', ' <a').replace('</a>\n', '</a> ') 668 669 min_offset = max(0, data.find('<HEADLINE>')) 670 671 intag = False 672 in_a = False 673 linebreak = True 674 sentences = [] 675 start = end = 0 676 sentence = '' 677 for i, c in enumerate(data): 678 if c == '<': 679 intag = True 680 if (len(data) > i + 2 681 and (data[i + 1: i + 3] == 'a ' 682 or data[i + 1: i + 3] == 'a>' 683 or data[i + 1: i + 3] == '/a') 684 ): 685 in_a = True 686 if in_a: 687 sentence += ' ' 688 end = i + 1 689 else: 690 # linebreak = False 691 if sentence and start >= min_offset: 692 sentences.append((start, end, sentence)) 693 start = end = i + 1 694 sentence = '' 695 elif c == '>': 696 intag = False 697 if in_a: 698 sentence += ' ' 699 end = i + 1 700 else: 701 start = end = i + 1 702 in_a = False 703 elif not intag and linebreak: 704 if c == '\n': 705 if sentence: 706 if start >= min_offset: 707 sentences.append((start, end, sentence)) 708 sentence = '' 709 start = end = i + 1 710 else: 711 sentence += c 712 end = i + 1 713 elif in_a: 714 sentence += ' ' 715 end = i + 1 716 if c == '\n': 717 linebreak = True 718 start = end = i + 1 719 if sentence: 720 if start >= min_offset: 721 sentences.append((start, end, sentence)) 722 # for s, e, t in sentences: 723 # if t != data[s:e]: 724 # print(t, data[s:e]) 725 726 # Re-tokenize sentences 727 sentences_ = [] 728 for sent in sentences: 729 if not sent[-1].startswith('http') and '</a>' not in sent[-1]: 730 sentences_.extend(sentence_tokenize(sent, language=language)) 731 return sentences_ 732 733 734 def read_annotation(path: str 735 ) -> (str, str, List[Entity], List[Relation], List[Event]): 736 """Reads annotation file. 737 738 Args: 739 path (str): path to the annotation file. 740 741 Returns: 742 doc_id (str): document ID. 743 source_type (str): type of the source of the document. 744 entity_list (List[Entity]): a list of Entity objects. 745 relation_list (List[Relation]) a list of Relation objects. 746 event_list (List[Event]): a list of Event objectss. 747 """ 748 data = open(path, 'r', encoding='utf-8').read() 749 750 soup = BeautifulSoup(data, 'lxml') 751 752 # metadata 753 root = soup.find('deft_ere') 754 doc_id = root['doc_id'] 755 source_type = root['source_type'] 756 757 # entities 758 entity_list = [] 759 entity_clusters = [] 760 entities_node = root.find('entities') 761 if entities_node: 762 for entity_node in entities_node.find_all('entity'): 763 entity_id = entity_node['id'] 764 entity_type = entity_node['type'] 765 entity_cluster = [] 766 for entity_mention_node in entity_node.find_all('entity_mention'): 767 mention_id = entity_mention_node['id'] 768 mention_type = entity_mention_node['noun_type'] 769 if mention_type == 'NOM': 770 mention_offset = int(entity_mention_node.find('nom_head')['offset']) 771 mention_length = int(entity_mention_node.find('nom_head')['length']) 772 mention_text = entity_mention_node.find('nom_head').text 773 else: 774 mention_offset = int(entity_mention_node['offset']) 775 mention_length = int(entity_mention_node['length']) 776 mention_text = entity_mention_node.find('mention_text').text 777 entity_list.append(Entity( 778 entity_id=entity_id, mention_id=mention_id, 779 entity_type=entity_type, mention_type=mention_type, 780 start=mention_offset, end=mention_offset + mention_length, 781 text=mention_text 782 )) 783 entity_cluster.append(Entity( 784 entity_id=entity_id, mention_id=mention_id, 785 entity_type=entity_type, mention_type=mention_type, 786 start=mention_offset, end=mention_offset + mention_length, 787 text=mention_text 788 )) 789 entity_clusters.append(Entity_whole(entity_id, entity_type, entity_cluster)) 790 fillers_node = root.find('fillers') 791 if fillers_node: 792 for filler_node in fillers_node.find_all('filler'): 793 entity_id = filler_node['id'] 794 entity_type = filler_node['type'] 795 if entity_type == 'weapon': 796 entity_type = 'WEA' 797 elif entity_type == 'vehicle': 798 entity_type = 'VEH' 799 else: 800 continue 801 mention_offset = int(filler_node['offset']) 802 mention_length = int(filler_node['length']) 803 mention_text = filler_node.text 804 entity_list.append( 805 Entity( 806 entity_id=entity_id, mention_id=entity_id, 807 entity_type=entity_type, mention_type='NOM', 808 start=mention_offset, end=mention_offset + mention_length, 809 text=mention_text 810 ) 811 ) 812 813 # relations 814 relation_list = [] 815 relations_node = root.find('relations') 816 if relations_node: 817 for relation_node in relations_node.find_all('relation'): 818 relation_id = relation_node['id'] 819 relation_type = relation_node['type'] 820 relation_subtype = relation_node['subtype'] 821 for relation_mention_node in relation_node.find_all('relation_mention'): 822 mention_id = relation_mention_node['id'] 823 arg1 = relation_mention_node.find('rel_arg1') 824 arg2 = relation_mention_node.find('rel_arg2') 825 if arg1 and arg2: 826 if arg1.has_attr('entity_id'): 827 arg1_entity_id = arg1['entity_id'] 828 arg1_mention_id = arg1['entity_mention_id'] 829 else: 830 arg1_entity_id = arg1['filler_id'] 831 arg1_mention_id = arg1['filler_id'] 832 arg1_role = arg1['role'] 833 arg1_text = arg1.text 834 if arg2.has_attr('entity_id'): 835 arg2_entity_id = arg2['entity_id'] 836 arg2_mention_id = arg2['entity_mention_id'] 837 else: 838 arg2_entity_id = arg2['filler_id'] 839 arg2_mention_id = arg2['filler_id'] 840 arg2_role = arg2['role'] 841 arg2_text = arg2.text 842 relation_list.append(Relation( 843 relation_id=relation_id, mention_id=mention_id, 844 relation_type=relation_type, 845 relation_subtype=relation_subtype, 846 arg1=RelationArgument(entity_id=arg1_entity_id, 847 mention_id=arg1_mention_id, 848 role=arg1_role, 849 text=arg1_text), 850 arg2=RelationArgument(entity_id=arg2_entity_id, 851 mention_id=arg2_mention_id, 852 role=arg2_role, 853 text=arg2_text))) 854 855 # events 856 event_list = [] 857 event_clusters = [] 858 events_node = root.find('hoppers') 859 if events_node: 860 for event_node in events_node.find_all('hopper'): 861 event_cluster = [] 862 event_arguments = [] 863 event_id = event_node['id'] 864 for event_mention_node in event_node.find_all('event_mention'): 865 trigger = event_mention_node.find('trigger') 866 trigger_offset = int(trigger['offset']) 867 trigger_length = int(trigger['length']) 868 arguments = [] 869 for arg in event_mention_node.find_all('em_arg'): 870 if arg['realis'] == 'false': 871 continue 872 if arg.has_attr('entity_id'): 873 arguments.append(EventArgument( 874 entity_id=arg['entity_id'], 875 mention_id=arg['entity_mention_id'], 876 role=arg['role'], 877 text=arg.text)) 878 if (arg['entity_id'], arg['role']) not in event_arguments: 879 event_arguments.append((arg['entity_id'], arg['role'])) 880 elif arg.has_attr('filler_id'): 881 arguments.append(EventArgument( 882 entity_id=arg['filler_id'], 883 mention_id=arg['filler_id'], 884 role=arg['role'], 885 text=arg.text 886 )) 887 if (arg['filler_id'], arg['role']) not in event_arguments: 888 event_arguments.append((arg['filler_id'], arg['role'])) 889 event_list.append(Event( 890 event_id=event_id, 891 mention_id=event_mention_node['id'], 892 event_type=event_mention_node['type'], 893 event_subtype=event_mention_node['subtype'], 894 trigger=Span(start=trigger_offset, 895 end=trigger_offset + trigger_length, 896 text=trigger.text), 897 arguments=arguments)) 898 event_cluster.append(Event( 899 event_id=event_id, 900 mention_id=event_mention_node['id'], 901 event_type=event_mention_node['type'], 902 event_subtype=event_mention_node['subtype'], 903 trigger=Span(start=trigger_offset, 904 end=trigger_offset + trigger_length, 905 text=trigger.text), 906 arguments=arguments)) 907 event_arguments = [EventArgument_whole(ent_id, ent_role) for ent_id, ent_role in event_arguments] 908 event_clusters.append(Event_whole(event_id, event_cluster[0].event_type, event_cluster[0].event_subtype, event_arguments, event_cluster)) 909 910 return doc_id, source_type, entity_list, relation_list, event_list, entity_clusters, event_clusters 911 912 913 def clean_entities(entities: List[Entity], 914 sentences: List[List[Tuple[int, int, str]]], 915 ) -> List[List[Entity]]: 916 """Cleans entities and assigns them to the corresponding sentences. 917 918 Args: 919 entities (List[Entity]): a list of Entity objects. 920 sentences (List[List[Tuple[int, int, str]]]): a list of window of sentences. 921 922 Returns: 923 List[List[Entity]]: a list of lists of Entity objects. The i-th list 924 represents Entity objects in the i-th window. 925 """ 926 sentence_entities = [[] for _ in range(len(sentences))] 927 for entity in entities: 928 start, end = entity.start, entity.end 929 for i, sentence in enumerate(sentences): # each window 930 flag = False 931 for s, e, text in sentence: 932 if start >= s and end <= e: 933 entity.remove_space() 934 # put this entity into this window 935 flag = True 936 break 937 if flag: 938 sentence_entities[i].append(entity) 939 940 941 # Remove overlapping entities 942 # TODO: we keep it becauce of fair comparison to OneIE, but this should be prevented for golds. 943 sentence_entities_ = [[] for _ in range(len(sentences))] 944 for i, entities in enumerate(sentence_entities): 945 if not entities: 946 continue 947 # Prefer longer entities 948 entities.sort(key=lambda x: (x.end - x.start), reverse=True) 949 chars = [0] * max([x.end for x in entities]) 950 for entity in entities: 951 overlap = False 952 for j in range(entity.start, entity.end): 953 if chars[j] == 1: 954 overlap = True 955 break 956 if overlap: 957 continue 958 else: 959 chars[entity.start:entity.end] = [1] * (entity.end - entity.start) 960 sentence_entities_[i].append(entity) 961 sentence_entities_[i].sort(key=lambda x: x.start) 962 963 return sentence_entities_ 964 965 966 def detect_multievent_triggers(sentence_events: List[List[Event]]) -> List[int]: 967 """Find sentences that contain multi-event triggers. 968 969 Args: 970 sentence_events (List[List[Event]]): a list of lists of Event objects. 971 972 Returns: 973 List[int]: a list of sentence indices. 974 """ 975 multievent_trigger_sents = [] 976 for sent_idx, events in enumerate(sentence_events): 977 trigger_event_types = defaultdict(set) 978 for event in events: 979 event_type = event.event_type 980 key = (event.trigger.start, event.trigger.end) 981 if event_type in trigger_event_types[key]: 982 multievent_trigger_sents.append(sent_idx) 983 break 984 else: 985 trigger_event_types[key].add(event_type) 986 return multievent_trigger_sents 987 988 989 def clean_events(events: List[Event], 990 sentence_entities: List[List[Entity]], 991 sentences: List[List[Tuple[int, int, str]]]) -> List[List[Event]]: 992 """Cleans events and assigns them to the corresponding sentences. 993 994 Args: 995 events (List[Event]): a list of Event objects. 996 sentence_entities (List[List[Entity]]): a list of lists of Entity 997 objects. The output of the `clean_entities`. 998 sentences (List[List[Tuple[int, int, str]]]): a list of window of sentences. 999 1000 Returns: 1001 List[List[Event]]: a list of lists of Event objects. The i-th list 1002 represents Event objects in the i-th window. 1003 """ 1004 sentence_events = [[] for _ in range(len(sentences))] 1005 for event in events: 1006 start, end = event.trigger.start, event.trigger.end 1007 for i, sentence in enumerate(sentences): # each window 1008 for s, e, _ in sentence: 1009 flag = False 1010 if start >= s and end <= e: 1011 event.trigger.remove_space() 1012 flag = True 1013 break 1014 if flag: 1015 # Clean the argument list 1016 arguments = [] 1017 entities = sentence_entities[i] 1018 for argument in event.arguments: 1019 entity_id = argument.entity_id 1020 mention_id = argument.mention_id 1021 for entity in entities: 1022 if (entity.entity_id == entity_id and 1023 entity.mention_id == mention_id): 1024 arguments.append(argument) 1025 event_ = Event(event_id=event.event_id, 1026 mention_id=event.mention_id, 1027 event_type=event.event_type, 1028 event_subtype=event.event_subtype, 1029 trigger=event.trigger.copy(), 1030 arguments=arguments 1031 ) 1032 sentence_events[i].append(event_) 1033 1034 1035 multievent_trigger_sents = detect_multievent_triggers(sentence_events) 1036 # Remove overlapping events 1037 # TODO: we keep it becauce of fair comparison to OneIE, but this should be prevented for golds. 1038 sentence_events_ = [[] for _ in range(len(sentences))] 1039 for i, events in enumerate(sentence_events): 1040 if not events: 1041 continue 1042 events.sort(key=lambda x: (EVENT_TYPE_WEIGHT.get(x.event_type, 0), 1043 x.trigger.end - x.trigger.start), 1044 reverse=True) 1045 chars = [0] * max([x.trigger.end for x in events]) 1046 for event in events: 1047 overlap = False 1048 for j in range(event.trigger.start, event.trigger.end): 1049 if chars[j] == 1: 1050 overlap = True 1051 break 1052 if overlap: 1053 continue 1054 else: 1055 chars[event.trigger.start:event.trigger.end] = \ 1056 [1] * (event.trigger.end - event.trigger.start) 1057 sentence_events_[i].append(event) 1058 sentence_events_[i].sort(key=lambda x: x.trigger.start) 1059 1060 return sentence_events_, multievent_trigger_sents 1061 1062 1063 def clean_relations(relations: List[Relation], 1064 sentence_entities: List[List[Entity]], 1065 sentences: List[List[Tuple[int, int, str]]] 1066 ) -> List[List[Relation]]: 1067 """Cleans relations and assigns them to the corresponding sentence. 1068 1069 Args: 1070 relations (List[Relation]): a list of Relation objects. 1071 sentence_entities (List[List[Entity]]): a list of lists of Entity 1072 objects. The output of the `clean_entities` function. 1073 sentences (List[List[Tuple[int, int, str]]]): a list of sentence tuples. 1074 1075 Returns: 1076 List[List[Relation]]: a list of lists of Relation objects. The i-th list 1077 represents Relation objects in the i-th window. 1078 """ 1079 sentence_relations = [[] for _ in range(len(sentences))] 1080 for relation in relations: 1081 keep = False 1082 entity_id_1, mention_id_1 = relation.arg1.entity_id, relation.arg1.mention_id 1083 entity_id_2, mention_id_2 = relation.arg2.entity_id, relation.arg2.mention_id 1084 for i, entities in enumerate(sentence_entities): 1085 arg1_in_sent = arg2_in_sent = False 1086 for entity in entities: 1087 if entity.entity_id == entity_id_1 and entity.mention_id == mention_id_1: 1088 arg1_in_sent = True 1089 if entity.entity_id == entity_id_2 and entity.mention_id == mention_id_2: 1090 arg2_in_sent = True 1091 if arg1_in_sent and arg2_in_sent: 1092 sentence_relations[i].append(relation) 1093 keep = True 1094 break 1095 elif arg1_in_sent != arg2_in_sent: 1096 # Stop searching because we find only one entity in the current 1097 # sentence. 1098 break 1099 return sentence_relations 1100 1101 def process_entity_cluster(entity_clusters: List[Entity_whole], 1102 sentence_entities: List[List[Entity]], 1103 sentences: List[List[Tuple[int, int, str]]] 1104 ) -> List[List[Entity_whole]]: 1105 """Cleans entity_cluster and splits them into lists 1106 1107 Args: 1108 entity_clusters (List[Entity_whole]): a list of Entity_whole instances. 1109 entities (List[List[Entity]]): a list of sentence entity list. 1110 sentences (List[List[Tuple[int, int, str]]]): a list of window of sentences. 1111 1112 Returns: 1113 List[List[Entity_whole]]: a list of sentence Entity_whole lists. 1114 """ 1115 sentence_entity_cluster = [[] for _ in range(len(sentences))] 1116 1117 for entity_cluster in entity_clusters: 1118 for i, sentence_entity in enumerate(sentence_entities): 1119 within_entities = [ent.copy() for ent in sentence_entity if ent.entity_id == entity_cluster.entity_id] 1120 if any(within_entities): 1121 sentence_entity_cluster[i].append(Entity_whole(entity_cluster.entity_id, entity_cluster.entity_type, within_entities)) 1122 return sentence_entity_cluster 1123 1124 def process_event_cluster(event_clusters: List[Event_whole], 1125 sentence_events: List[List[Event]], 1126 sentences: List[List[Tuple[int, int, str]]] 1127 ) -> List[List[Event_whole]]: 1128 """Cleans and assigns event cluster. 1129 1130 Args: 1131 event_clusters (List[Event_whole]): A list of event cluster objects 1132 sentence_events (List[List[Event]]): A list of sentence event lists. 1133 sentences (List[List[Tuple[int, int, str]]]): a list of window of sentences. 1134 1135 Returns: 1136 List[List[Event_whole]]: a list of sentence event cluster lists. 1137 """ 1138 sentence_events_cluster = [[] for _ in range(len(sentences))] 1139 # Reconstruct event_whole based on the event mentions that within the sentence(window) span 1140 for event_cluster in event_clusters: 1141 for i, sentence_event in enumerate(sentence_events): 1142 within_events = [eve.copy() for eve in sentence_event if (eve.event_id == event_cluster.event_id)] 1143 if any(within_events): 1144 arguments = [] 1145 for argu in event_cluster.arguments: 1146 flag = False 1147 for within_event in within_events: 1148 for arg in within_event.arguments: 1149 if arg.entity_id == argu.entity_id: 1150 flag = True 1151 break 1152 if flag: 1153 break 1154 if flag: 1155 arguments.append(argu) 1156 sentence_events_cluster[i].append(Event_whole(event_cluster.event_id, event_cluster.event_type, 1157 event_cluster.event_subtype, arguments, within_events)) 1158 1159 return sentence_events_cluster 1160 1161 def tokenize(sentences: List[Tuple[int, int, str]], 1162 entities: List[Entity], 1163 events: List[Event], 1164 language: str = 'english' 1165 ) -> List[Tuple[int, int, str]]: 1166 """Tokenizes a sentence and makes sure entity and event spans are compatible 1167 with the tokenization result. 1168 1169 Args: 1170 sentence (List[Tuple[int, int, str]]): a window of sentences. 1171 entities (List[Entity]): a list of Entity objects. 1172 events (List[Entity]): a list of Event objects. 1173 language (str, optional): sentence langauge. 1174 1175 Returns: 1176 tokens (List[Tuple[int, int, str]]): a list of token tuples, where each 1177 item consists of start_offset, end_offset, and text. 1178 """ 1179 all_tokens = [] 1180 token_starts = [0] 1181 for sentence in sentences: 1182 start, end, text = sentence 1183 text = mask_escape(text) 1184 # Split the sentence into chunks 1185 splits = {0, len(text)} 1186 for entity in entities: 1187 if entity.start >= start and entity.end <= end: 1188 splits.add(entity.start - start) 1189 splits.add(entity.end - start) 1190 for event in events: 1191 if event.trigger.start >= start and event.trigger.end <= end: 1192 splits.add(event.trigger.start - start) 1193 splits.add(event.trigger.end - start) 1194 splits = sorted(list(splits)) 1195 chunks = [(splits[i], splits[i + 1], text[splits[i]:splits[i + 1]]) 1196 for i in range(len(splits) - 1)] 1197 1198 # Tokenize each chunk 1199 chunks = [(s, e, t, wordpunct_tokenize(t)) for s, e, t in chunks] 1200 1201 # Merge chunks and add word offsets 1202 tokens = [] 1203 for chunk_start, chunk_end, chunk_text, chunk_tokens in chunks: 1204 last = 0 1205 chunk_tokens_ = [] 1206 for token in chunk_tokens: 1207 token_start = chunk_text[last:].find(token) 1208 if token_start == -1: 1209 raise ValueError('Cannot find token {} in {}'.format(token, text)) 1210 token_end = token_start + len(token) 1211 chunk_tokens_.append((token_start + start + last + chunk_start, 1212 token_end + start + last + chunk_start, 1213 unmask_escape(token))) 1214 last += token_end 1215 tokens.extend(chunk_tokens_) 1216 1217 all_tokens.extend(tokens) 1218 token_starts.append(len(all_tokens)) 1219 return all_tokens, token_starts 1220 1221 def sentence2window(sentences, window_size_): 1222 """Convert sentences list to windows 1223 1224 Args: 1225 sentences (List[Tuple[int, int, str]]) 1226 window_size : int 1227 1228 Returns: 1229 List[List[Tuple[int, int, str]]] 1230 1231 """ 1232 if window_size_ > len(sentences): 1233 window_size = len(sentences) 1234 else: 1235 window_size = window_size_ 1236 1237 output = [] 1238 for i in range(len(sentences)-window_size+1): 1239 output.append(sentences[i:i+window_size]) 1240 return output 1241 1242 def extract(source_path: str, 1243 ere_path: str, 1244 doc_id: str, 1245 window_size: int = 1, 1246 language: str = 'english', 1247 discard_sentences_with_multievent_triggers: bool = True) -> Document: 1248 """Generates a Document object from a text file and an annotation file. 1249 1250 Args: 1251 source_path (str): path to the source file. 1252 ere_path (str): path to the annotation file. 1253 doc_id (str): document ID. 1254 language (str, optional): document language. Defaults to 'english'. 1255 1256 Returns: 1257 Document: a Document object that contains tokens, entities, relations, 1258 and events of a document. 1259 """ 1260 wrapped = (doc_id in WRAPPED_DOCS 1261 or (language == 'spanish' and 'newswire' in source_path )) 1262 sentences = read_source_file(source_path, 1263 language=language, 1264 wrapped=wrapped) 1265 # Revise sentences 1266 if doc_id in DOCS_TO_REVISE_SENT: 1267 sentences = revise_sentences(sentences, doc_id) 1268 doc_id_, source_type, entities, relations, events, entity_cluster, event_cluster = read_annotation(ere_path) 1269 1270 # concate windows 1271 sentences = sentence2window(sentences, window_size) 1272 1273 # Remove entities and events out of extracted sentences 1274 sentence_entities = clean_entities(entities, sentences) 1275 sentence_events, multievent_trigger_sents = clean_events( 1276 events, sentence_entities, sentences) 1277 sentence_relations = clean_relations(relations, sentence_entities, sentences) 1278 1279 # Process entity_cluster, event_cluster, relation(entity_base) 1280 sentence_entity_cluster = process_entity_cluster(entity_cluster, sentence_entities, sentences) 1281 sentence_event_cluster = process_event_cluster(event_cluster, sentence_events, sentences) 1282 1283 # Tokenization 1284 sentence_tokens = [] 1285 sentence_starts = [] 1286 for s, ent, evt in zip(sentences, sentence_entities, sentence_events): 1287 sent_tokens, sent_start = tokenize(s, ent, evt, language=language) 1288 sentence_tokens.append(sent_tokens) 1289 sentence_starts.append(sent_start) 1290 1291 1292 # Convert span character offsets to token index offsets 1293 sentence_objs = [] 1294 for i, (tokens, entities, events, relations, sentence, entity_cluster, event_cluster, starts) in enumerate(zip( 1295 sentence_tokens, sentence_entities, sentence_events, 1296 sentence_relations, sentences, sentence_entity_cluster, sentence_event_cluster, sentence_starts)): 1297 wnd_id = '{}-{}'.format(doc_id, i) 1298 if (discard_sentences_with_multievent_triggers 1299 and i in multievent_trigger_sents): 1300 continue 1301 for entity in entities: 1302 entity.char_offsets_to_token_offsets(tokens) 1303 for entity in entity_cluster: 1304 for mention in entity.entity_mentions: 1305 mention.char_offsets_to_token_offsets(tokens) 1306 for event in events: 1307 event.trigger.char_offsets_to_token_offsets(tokens) 1308 for event in event_cluster: 1309 for mention in event.event_mentions: 1310 mention.trigger.char_offsets_to_token_offsets(tokens) 1311 sentence_objs.append(Sentence( 1312 doc_id=doc_id, sent_id=wnd_id, 1313 tokens=[t for _, _, t in tokens], 1314 entities=entities, 1315 relations=relations, 1316 events=events, 1317 sent_starts=starts, 1318 entity_cluster=entity_cluster, 1319 event_cluster=event_cluster, 1320 start=sentence[0][0], 1321 end=sentence[-1][1], 1322 text=' '.join([s[2] for s in sentence]) 1323 )) 1324 return Document(doc_id=doc_id, sentences=sentence_objs) 1325 1326 1327 def process_batch(input_dir, output_file, window_size=1, dataset='normal', language='english'): 1328 """Processes a batch of documents. 1329 1330 Args: 1331 input_dir (str): path to the input directory. 1332 output_file (str): path to the output file. 1333 dataset (str, optional): dataset type. Defaults to 'normal'. 1334 language (str, optional): dataset language. Defaults to 'english'. 1335 """ 1336 if dataset == 'normal': 1337 source_files = glob.glob( 1338 os.path.join(input_dir, 'source', 'cmptxt', '*.txt')) 1339 elif dataset == 'r2v2': 1340 source_files = glob.glob(os.path.join(input_dir, 'source', '*.txt')) 1341 elif dataset == 'parallel': 1342 source_files = glob.glob( 1343 os.path.join(input_dir, 'eng', 'translation', '*.txt')) 1344 elif dataset == 'spanish': 1345 source_files = glob.glob( 1346 os.path.join(input_dir, 'source', '**', '*.txt')) 1347 else: 1348 raise ValueError('Unknown dataset type: {}'.format(dataset)) 1349 1350 with open(output_file, 'w', encoding='utf-8') as w: 1351 for source_file in source_files: 1352 doc_id = os.path.basename(source_file).replace('.txt', '') \ 1353 .replace('.cmp', '').replace('.mp', '') 1354 if dataset == 'normal': 1355 annotation_file = os.path.join(input_dir, 'ere', 'cmptxt', 1356 '{}.rich_ere.xml'.format(doc_id)) 1357 elif dataset == 'r2v2': 1358 annotation_file = os.path.join(input_dir, 'ere', 1359 '{}.rich_ere.xml'.format(doc_id)) 1360 elif dataset == 'parallel': 1361 annotation_file = os.path.join(input_dir, 'eng', 'ere', 1362 '{}.rich_ere.xml'.format(doc_id)) 1363 elif dataset == 'spanish': 1364 annotation_file = (source_file.replace('/source/', '/ere/') 1365 .replace('.txt', '.rich_ere.xml')) 1366 doc = extract(source_file, annotation_file, doc_id, window_size, language) 1367 for sent in doc.sentences: 1368 w.write(json.dumps(sent.to_dict()) + '\n') 1369 1370 1371 def ere_to_oneie(input_file: str, 1372 output_file: str, 1373 tokenizer: PreTrainedTokenizer): 1374 """Converts to OneIE format. 1375 1376 Args: 1377 input_file (str): path to the input file. 1378 output_file (str): path to the output file. 1379 tokenizer (PreTrainedTokenizer): a tokenizer that converts tokens to 1380 word pieces. 1381 """ 1382 skip_num = 0 1383 with open(input_file, 'r', encoding='utf-8') as r, \ 1384 open(output_file, 'w', encoding='utf-8') as w: 1385 for line in r: 1386 inst = json.loads(line) 1387 # tokens 1388 tokens = inst['tokens'] 1389 pieces = [tokenizer.tokenize(t) for t in tokens] 1390 token_lens = [len(x) for x in pieces] 1391 if 0 in token_lens: 1392 skip_num += 1 1393 continue 1394 pieces = [p for ps in pieces for p in ps] 1395 sentence = inst['text'] 1396 # entities 1397 entity_mentions = [] 1398 entity_text = {} 1399 for entity in inst['entities']: 1400 entity_mentions.append({ 1401 'id': entity['entity_id'], 1402 'start': entity['start'], 1403 'end': entity['end'], 1404 'entity_type': entity['entity_type'], 1405 'mention_type': entity['mention_type'], 1406 'text': entity['text'] 1407 }) 1408 entity_text[entity['entity_id']] = entity['text'] 1409 # relations 1410 relation_mentions = [] 1411 for relation in inst['relations']: 1412 relation_type = relation_type_mapping[relation['relation_type']] 1413 relation_subtype = relation_subtype_mapping[ 1414 relation['relation_subtype']] 1415 relation_mentions.append({ 1416 'id': relation['relation_id'], 1417 'relation_type': relation_type, 1418 'relation_subtype': '{}:{}'.format(relation_type, 1419 relation_subtype), 1420 'arguments': [ 1421 { 1422 'entity_id': relation['arg1']['entity_id'], 1423 'role': 'Arg-1', 1424 'text': entity_text[relation['arg1']['entity_id']] 1425 }, 1426 { 1427 'entity_id': relation['arg2']['entity_id'], 1428 'role': 'Arg-2', 1429 'text': entity_text[relation['arg2']['entity_id']] 1430 }, 1431 ] 1432 }) 1433 # events 1434 event_mentions = [] 1435 for event in inst['events']: 1436 event_mentions.append({ 1437 'id': event['event_id'], 1438 'event_type': event_type_mapping['{}:{}'.format( 1439 event['event_type'], event['event_subtype'])], 1440 'trigger': { 1441 'start': event['trigger']['start'], 1442 'end': event['trigger']['end'], 1443 'text': event['trigger']['text']}, 1444 'arguments': [{ 1445 'entity_id': arg['entity_id'], 1446 'text': entity_text[arg['entity_id']], 1447 'role': role_type_mapping[arg['role']] 1448 } for arg in event['arguments']] 1449 }) 1450 1451 # coreference 1452 corefs = [] 1453 for entity in inst['entity_cluster']: 1454 if len(entity['entity_mentions']) > 1: 1455 corefs.append({ 1456 'id': entity['entity_id'], 1457 'entities': [{ 1458 'id': mention['entity_id'], 1459 'text': mention['text'], 1460 'entity_type': mention['entity_type'], 1461 'mention_type': mention['mention_type'], 1462 'start': mention['start'], 1463 'end': mention['end'] 1464 } for mention in entity['entity_mentions']], 1465 'entity_type': entity['entity_type'] 1466 }) 1467 1468 # event coreference 1469 event_corefs = [] 1470 for event in inst['event_cluster']: 1471 if len(event['event_mentions']) > 1 : 1472 event_corefs.append({ 1473 'id': event['event_id'], 1474 'events': [{ 1475 'id': mention['event_id'], 1476 'event_type': event_type_mapping['{}:{}'.format( 1477 mention['event_type'], mention['event_subtype'])], 1478 'trigger': { 1479 'start': mention['trigger']['start'], 1480 'end': mention['trigger']['end'], 1481 'text': mention['trigger']['text']}, 1482 'arguments':[{ 1483 'entity_id': arg['entity_id'], 1484 'text': entity_text[arg['entity_id']], 1485 'role': role_type_mapping[arg['role']] 1486 } for arg in mention['arguments']], 1487 } for mention in event['event_mentions']], 1488 'event_type': event_type_mapping['{}:{}'.format( 1489 event['event_type'], event['event_subtype'])] 1490 }) 1491 1492 w.write(json.dumps({ 1493 'doc_id': inst['doc_id'], 1494 'wnd_id': inst['wnd_id'], 1495 'tokens': tokens, 1496 'pieces': pieces, 1497 'token_lens': token_lens, 1498 'sentence': inst['text'], 1499 'entity_mentions': entity_mentions, 1500 'relation_mentions': relation_mentions, 1501 'event_mentions': event_mentions, 1502 'entity_coreference': corefs, 1503 'event_coreference': event_corefs, 1504 'sentence_starts': inst['sentence_starts'][:-1] 1505 }) + '\n') 1506 print('#Skip: {}'.format(skip_num)) 1507 1508 1509 def split_data(input_file, output_dir, split_path): 1510 """Splits the input file into train/dev/test sets. 1511 1512 Args: 1513 input_file (str): path to the input file. 1514 output_dir (str): path to the output directory. 1515 split_path (str): path to the split directory that contains three files, 1516 train.doc.txt, dev.doc.txt, and test.doc.txt . Each line in these 1517 files is a document ID. 1518 """ 1519 print('Splitting the dataset into train/dev/test sets') 1520 train_docs, dev_docs, test_docs = set(), set(), set() 1521 # Load doc ids 1522 with open(os.path.join(split_path, 'train.doc.txt')) as r: 1523 train_docs.update(r.read().strip('\n').split('\n')) 1524 with open(os.path.join(split_path, 'dev.doc.txt')) as r: 1525 dev_docs.update(r.read().strip('\n').split('\n')) 1526 with open(os.path.join(split_path, 'test.doc.txt')) as r: 1527 test_docs.update(r.read().strip('\n').split('\n')) 1528 1529 real_train, real_dev, real_test = set(), set(), set() 1530 # Split the dataset 1531 with open(input_file, 'r', encoding='utf-8') as r, \ 1532 open(os.path.join(output_dir, 'train.oneie.json'), 'w') as w_train, \ 1533 open(os.path.join(output_dir, 'dev.oneie.json'), 'w') as w_dev, \ 1534 open(os.path.join(output_dir, 'test.oneie.json'), 'w') as w_test: 1535 for line in r: 1536 inst = json.loads(line) 1537 doc_id = inst['doc_id'] 1538 if doc_id in train_docs: 1539 w_train.write(line) 1540 real_train.add(doc_id) 1541 elif doc_id in dev_docs: 1542 w_dev.write(line) 1543 real_dev.add(doc_id) 1544 elif doc_id in test_docs: 1545 w_test.write(line) 1546 real_test.add(doc_id) 1547 else: 1548 print('this doc is not been covered!') 1549 print(train_docs.difference(real_train)) 1550 print(dev_docs.difference(real_dev)) 1551 print(test_docs.difference(real_test)) 1552 1553 1554 def parse_arguments() -> Namespace: 1555 """Parses command line arguments. 1556 1557 Returns: 1558 args (Namespace): a Namespace of parsed arguments. 1559 """ 1560 parser = ArgumentParser() 1561 parser.add_argument('-i', 1562 '--input', 1563 required=True, 1564 help='Path to the input folder') 1565 parser.add_argument('-o', 1566 '--output', 1567 required=True, 1568 help='Path to the output folder') 1569 parser.add_argument('-b', 1570 '--bert', 1571 help='BERT model name', 1572 default='bert-large-cased') 1573 parser.add_argument('-c', 1574 '--bert_cache_dir', 1575 help='Path to the BERT cahce directory') 1576 1577 parser.add_argument('-s', '--split', default=None, 1578 help='Path to the split folder') 1579 1580 parser.add_argument('-w', '--window', default=1, help='Integer for window size', type=int) 1581 args = parser.parse_args() 1582 1583 return args 1584 1585 1586 def main(): 1587 args = parse_arguments() 1588 1589 model_name = args.bert 1590 cache_dir = args.bert_cache_dir 1591 1592 # Create a tokenizer based on the model name 1593 if model_name.startswith('bert-'): 1594 tokenizer = BertTokenizer.from_pretrained(model_name, 1595 cache_dir=cache_dir) 1596 elif model_name.startswith('roberta-'): 1597 tokenizer = RobertaTokenizer.from_pretrained(model_name, 1598 cache_dir=cache_dir) 1599 elif model_name.startswith('xlm-roberta-'): 1600 tokenizer = XLMRobertaTokenizer.from_pretrained(model_name, 1601 cache_dir=cache_dir) 1602 elif model_name.startswith('lanwuwei'): 1603 tokenizer = BertTokenizer.from_pretrained(model_name, 1604 cache_dir=cache_dir, 1605 do_lower_case=True) 1606 else: 1607 tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir, do_lower_case=False, use_fast=False) 1608 1609 if args.window > 10000: 1610 f_size = 'doc' 1611 else: 1612 f_size = 'w{}'.format(args.window) 1613 1614 sub_set = [ 1615 ("LDC2015E29_DEFT_Rich_ERE_English_Training_Annotation_V2/data", "normal"), 1616 ("LDC2015E68_DEFT_Rich_ERE_English_Training_Annotation_R2_V2/data", "r2v2"), 1617 ("LDC2015E78_DEFT_Rich_ERE_Chinese_and_English_Parallel_Annotation_V2/data", "parallel") 1618 ] 1619 1620 for d_path, dataset in sub_set: 1621 # Convert to JSON format 1622 json_path = os.path.join(args.output, '{}.{}.{}.json'.format("english", 1623 dataset, 1624 f_size)) 1625 process_batch(os.path.join(args.input, d_path), 1626 json_path, 1627 window_size=args.window, 1628 dataset=dataset, 1629 language="english" 1630 ) 1631 1632 # Convert to OneIE foramt 1633 oneie_path = os.path.join(args.output, '{}.{}.{}.oneie.json'.format("english", 1634 dataset, 1635 f_size)) 1636 ere_to_oneie(json_path, oneie_path, tokenizer=tokenizer) 1637 1638 if args.split: 1639 all_data = [] 1640 for ds in ["normal", "r2v2", "parallel"]: 1641 with open(os.path.join(args.output,'{}.{}.{}.oneie.json'.format("english",ds,f_size)), 'r', encoding='utf-8') as r: 1642 for line in r: 1643 all_data.append(json.loads(line)) 1644 train_docs, dev_docs, test_docs = set(), set(), set() 1645 # load doc ids 1646 with open(os.path.join(args.split, 'train.doc.txt')) as r: 1647 train_docs.update(r.read().strip('\n').split('\n')) 1648 with open(os.path.join(args.split, 'dev.doc.txt')) as r: 1649 dev_docs.update(r.read().strip('\n').split('\n')) 1650 with open(os.path.join(args.split, 'test.doc.txt')) as r: 1651 test_docs.update(r.read().strip('\n').split('\n')) 1652 1653 with open(os.path.join(args.output, 'train.{}.oneie.json'.format(f_size)), 'w') as w_train, \ 1654 open(os.path.join(args.output, 'dev.{}.oneie.json'.format(f_size)), 'w') as w_dev, \ 1655 open(os.path.join(args.output, 'test.{}.oneie.json'.format(f_size)), 'w') as w_test: 1656 for inst in all_data: 1657 doc_id = inst['doc_id'] 1658 if doc_id in train_docs: 1659 w_train.write(json.dumps(inst) + '\n') 1660 elif doc_id in dev_docs: 1661 w_dev.write(json.dumps(inst) + '\n') 1662 elif doc_id in test_docs: 1663 w_test.write(json.dumps(inst) + '\n') 1664 else: 1665 print('missing!! {}'.format(doc_id)) 1666 1667 if __name__ == '__main__': 1668 main()